| |
| """Numerical optimization for parameters and opted-in numeric literals.""" |
| from __future__ import annotations |
|
|
| from dataclasses import dataclass, replace as dataclass_replace |
| from typing import Any, Dict |
| import warnings |
|
|
| import numpy as np |
| import sr_harness_engine as engine |
| from scipy.optimize import differential_evolution, least_squares |
| from sr_harness_engine.evaluation import grouped_parameter_key |
| from sr_harness_engine.tree import children, with_children, transform |
|
|
|
|
| @dataclass(frozen=True) |
| class ConstantOptimizerConfig: |
| """Numerical budget and reproducibility controls for ``fit_constants``.""" |
|
|
| random_state: int = 0 |
| n_restarts: int = 3 |
| max_nfev: int = 700 |
| differential_evolution_maxiter: int = 15 |
| differential_evolution_popsize: int = 5 |
| use_global_search: bool = True |
| snap_exponents: bool = True |
| use_eps: float = 1e-8 |
|
|
|
|
| @dataclass(frozen=True) |
| class _ParameterSpec: |
| name: str |
| label: Any | None |
| initial: float |
| lower: float |
| upper: float |
| literal: bool = False |
|
|
|
|
| def _parameterize_literals( |
| expression: engine.Expression, |
| data: Dict[str, np.ndarray], |
| ) -> tuple[engine.Expression, list[_ParameterSpec]]: |
| specs: list[_ParameterSpec] = [] |
| named: dict[str, int] = {} |
| grouped: dict[str, engine.GroupedParameter] = {} |
| literal_index = 0 |
|
|
| def visit( |
| node: engine.Expression, |
| parent: engine.Expression | None = None, |
| child_index: int | None = None, |
| ) -> engine.Expression: |
| nonlocal literal_index |
| if isinstance(node, engine.Number): |
| name = f"__literal_{literal_index}" |
| literal_index += 1 |
| value = float(node.value) |
| limit = max(100.0, 10.0 * abs(value) + 1.0) |
| lower, upper = -limit, limit |
| if isinstance(parent, engine.Binary) and parent.operator == "**" and child_index == 1: |
| lower, upper = -10.0, 10.0 |
| if ( |
| isinstance(parent, engine.Binary) |
| and parent.operator == "-" |
| and child_index == 1 |
| and isinstance(parent.left, engine.Symbol) |
| and parent.left.name in data |
| ): |
| feature = np.asarray(data[parent.left.name], dtype=float).reshape(-1) |
| feature = feature[np.isfinite(feature)] |
| if feature.size: |
| width = float(np.max(feature) - np.min(feature)) |
| margin = max(0.2 * width, 10.0) |
| lower = float(np.min(feature) - margin) |
| upper = float(np.max(feature) + margin) |
| specs.append(_ParameterSpec(name, None, value, lower, upper, literal=True)) |
| return engine.Parameter(name, value) |
| if isinstance(node, engine.Parameter): |
| if node.name not in named: |
| initial = 1.0 if node.value is None else float(node.value) |
| limit = max(100.0, 10.0 * abs(initial) + 1.0) |
| named[node.name] = len(specs) |
| specs.append(_ParameterSpec(node.name, None, initial, -limit, limit)) |
| return node |
| if isinstance(node, engine.GroupedParameter): |
| key = grouped_parameter_key(node) |
| grouped.setdefault(key, node) |
| return dataclass_replace(node, by=visit(node.by, node, 0)) |
| transformed_children = tuple( |
| visit(child, node, index) for index, child in enumerate(children(node)) |
| ) |
| return with_children(node, transformed_children) |
|
|
| parameterized = visit(expression) |
| for key, node in grouped.items(): |
| labels = np.asarray(node.by.evaluate(data), dtype=object).reshape(-1) |
| unique = [] |
| for label in labels: |
| label = label.item() if isinstance(label, np.generic) else label |
| if label not in unique: |
| unique.append(label) |
| supplied = dict(node.value or {}) |
| for label in unique: |
| initial = float(supplied.get(label, node.default if node.default is not None else 1.0)) |
| limit = max(100.0, 10.0 * abs(initial) + 1.0) |
| specs.append(_ParameterSpec(key, label, initial, -limit, limit)) |
| return parameterized, specs |
|
|
|
|
| def _unpack(raw: np.ndarray, specs: list[_ParameterSpec]) -> dict[str, Any]: |
| result: dict[str, Any] = {} |
| for value, spec in zip(raw, specs): |
| if spec.label is None: |
| result[spec.name] = float(value) |
| else: |
| result.setdefault(spec.name, {})[spec.label] = float(value) |
| return result |
|
|
|
|
| def _bind( |
| expression: engine.Expression, |
| parameters: dict[str, Any], |
| literal_names: set[str], |
| ) -> engine.Expression: |
| def bind(node: engine.Expression) -> engine.Expression: |
| if isinstance(node, engine.Parameter) and node.name in parameters: |
| value = float(parameters[node.name]) |
| if node.name in literal_names: |
| return engine.Number(value) |
| return engine.Parameter(node.name, value) |
| if isinstance(node, engine.GroupedParameter): |
| key = grouped_parameter_key(node) |
| if key in parameters: |
| return dataclass_replace(node, value=dict(parameters[key])) |
| return node |
|
|
| return transform(expression, bind) |
|
|
|
|
| def fit_constants( |
| f: engine.Expression, |
| data: Dict[str, np.ndarray], |
| y: np.ndarray, |
| *, |
| config: ConstantOptimizerConfig | None = None, |
| num_nodes: int | None = None, |
| ) -> engine.Expression: |
| """Return a fitted expression without modifying ``f``. |
| |
| This function is an explicit opt-in optimization operation, so numeric |
| literals are temporarily promoted to independent parameters. Outside this |
| function, literals remain fixed and only ``param(...)`` nodes are fitable. |
| """ |
| if not isinstance(f, engine.Expression): |
| raise TypeError("f must be an sr_harness_engine.Expression") |
| config = config or ConstantOptimizerConfig() |
| target = np.asarray(y, dtype=float).reshape(-1) |
| if target.size == 0 or not np.all(np.isfinite(target)): |
| raise ValueError("y must contain at least one finite-only sample") |
| arrays = {name: np.asarray(value) for name, value in data.items()} |
| if num_nodes is None and any( |
| np.asarray(value).reshape(-1).size != target.size |
| for value in arrays.values() |
| ): |
| raise ValueError("all data arrays must have the same number of samples as y") |
|
|
| expression, specs = _parameterize_literals(f.copy(), arrays) |
| if not specs: |
| return expression.copy() |
| initial = np.asarray([spec.initial for spec in specs], dtype=float) |
| lower = np.asarray([spec.lower for spec in specs], dtype=float) |
| upper = np.asarray([spec.upper for spec in specs], dtype=float) |
| scale = max(float(np.std(target)), float(np.max(np.abs(target))) * 1e-12, 1e-12) |
| invalid = np.full(target.size, 1e6, dtype=float) |
|
|
| def prediction(raw: np.ndarray) -> np.ndarray | None: |
| try: |
| with np.errstate(all="ignore"): |
| value = np.asarray( |
| expression.evaluate( |
| arrays, |
| parameters=_unpack(raw, specs), |
| num_nodes=num_nodes, |
| ), |
| dtype=float, |
| ).reshape(-1) |
| if value.size == 1 and target.size != 1: |
| value = np.full(target.size, float(value[0])) |
| if value.size != target.size or not np.all(np.isfinite(value)): |
| return None |
| return value |
| except (ArithmeticError, FloatingPointError, KeyError, TypeError, ValueError): |
| return None |
|
|
| def residual(raw: np.ndarray) -> np.ndarray: |
| value = prediction(raw) |
| return invalid if value is None else (value - target) / scale |
|
|
| def loss(raw: np.ndarray) -> float: |
| value = prediction(raw) |
| return 1e100 if value is None else float(np.mean(np.square(value - target))) |
|
|
| starts = [np.clip(initial, lower, upper)] |
| rng = np.random.default_rng(config.random_state) |
| for _ in range(max(0, config.n_restarts - 1)): |
| starts.append(rng.uniform(lower, upper)) |
| if config.use_global_search and len(specs) <= 12: |
| with warnings.catch_warnings(): |
| warnings.simplefilter("ignore") |
| global_result = differential_evolution( |
| loss, |
| list(zip(lower, upper)), |
| seed=config.random_state, |
| maxiter=config.differential_evolution_maxiter, |
| popsize=config.differential_evolution_popsize, |
| polish=False, |
| updating="immediate", |
| ) |
| starts.append(global_result.x) |
|
|
| best = starts[0] |
| best_loss = loss(best) |
| for start in starts: |
| with warnings.catch_warnings(): |
| warnings.simplefilter("ignore") |
| result = least_squares( |
| residual, |
| start, |
| bounds=(lower, upper), |
| max_nfev=config.max_nfev, |
| xtol=1e-13, |
| ftol=1e-13, |
| gtol=1e-13, |
| ) |
| current_loss = loss(result.x) |
| if current_loss < best_loss: |
| best, best_loss = result.x, current_loss |
|
|
| parameters = _unpack(best, specs) |
| literal_names = {spec.name for spec in specs if spec.literal} |
| return _bind(expression, parameters, literal_names) |
|
|
