| """Parameter optimization for symbolic expressions.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass, replace |
| from typing import Any, Mapping |
|
|
| import numpy as np |
| from scipy.optimize import minimize |
|
|
| from .evaluation import grouped_parameter_key, walk |
| from .expression import Expression, GroupedParameter, Parameter |
|
|
|
|
| def bind_parameters(expression: Expression, parameters: Mapping[str, Any]) -> Expression: |
| """Return an expression whose parameter nodes contain fitted values.""" |
| from .tree import transform |
|
|
| def bind(node: Expression) -> Expression: |
| if isinstance(node, Parameter) and node.name in parameters: |
| return replace(node, value=float(parameters[node.name])) |
| if isinstance(node, GroupedParameter): |
| key = grouped_parameter_key(node) |
| if key in parameters: |
| return replace(node, value=dict(parameters[key])) |
| return node |
|
|
| return transform(expression, bind) |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class FitResult: |
| """Result of fitting an expression to target observations.""" |
| expression: Expression |
| parameters: dict[str, Any] |
| loss: float |
| success: bool |
| message: str |
| n_iter: int |
|
|
| def evaluate( |
| self, values: Mapping[str, Any], *, time: Any = None, |
| delay_resolver: Any = None, num_nodes: int | None = None, |
| ) -> Any: |
| """Evaluate the supplied model or expression. |
| |
| Args: |
| values: Values keyed by symbol name or an evaluation context. |
| time: Optional sample times. |
| delay_resolver: Optional callback that resolves delayed values. |
| num_nodes: Explicit node count for indexed expressions. |
| |
| Returns: |
| Predictions from the fitted expression. |
| """ |
| return self.expression.evaluate( |
| values, parameters=self.parameters, time=time, |
| delay_resolver=delay_resolver, num_nodes=num_nodes, |
| ) |
|
|
| predict = evaluate |
|
|
|
|
| def fit( |
| expression: Expression, |
| values: Mapping[str, Any], |
| target: Any, |
| *, |
| initial: Mapping[str, Any] | None = None, |
| method: str = "BFGS", |
| options: Mapping[str, Any] | None = None, |
| num_nodes: int | None = None, |
| ) -> FitResult: |
| """Minimize mean squared error and return fitted parameter values. |
| |
| Args: |
| expression: Symbolic expression to process. |
| values: Values keyed by symbol name. |
| target: Target name or target values. |
| initial: Optional initial parameter values. |
| method: Optimization method name. |
| options: Optional optimizer settings. |
| num_nodes: Explicit node count for indexed expressions. |
| |
| Returns: |
| Fitted parameters, expression, predictions, and loss. |
| """ |
| initial = dict(initial or {}) |
| target = np.asarray(target, dtype=float) |
| specs: list[tuple[str, Any | None]] = [] |
| named: dict[str, float | None] = {} |
| grouped: dict[str, tuple[GroupedParameter, list[Any]]] = {} |
|
|
| for node in walk(expression): |
| if isinstance(node, Parameter): |
| previous = named.get(node.name) |
| if previous is not None and node.value is not None and previous != node.value: |
| raise ValueError(f"Conflicting initial values for parameter {node.name!r}.") |
| if node.name not in named or node.value is not None: |
| named[node.name] = node.value |
| elif isinstance(node, GroupedParameter): |
| key = grouped_parameter_key(node) |
| labels = np.asarray(node.by.evaluate(values), 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) |
| grouped.setdefault(key, (node, unique)) |
|
|
| vector = [] |
| for name, default in named.items(): |
| value = initial.get(name, default if default is not None else 1.0) |
| vector.append(float(value)) |
| specs.append((name, None)) |
| for key, (node, labels) in grouped.items(): |
| supplied = initial.get(key, node.value or {}) |
| if not isinstance(supplied, Mapping): |
| raise TypeError(f"Initial value for grouped parameter {key!r} must be a mapping.") |
| for label in labels: |
| default = node.default if node.default is not None else 1.0 |
| vector.append(float(supplied.get(label, default))) |
| specs.append((key, label)) |
|
|
| if not specs: |
| prediction = np.asarray(expression.evaluate(values, num_nodes=num_nodes), dtype=float) |
| loss = float(np.mean((prediction - target) ** 2)) |
| return FitResult(expression, {}, loss, True, "No parameters to optimize.", 0) |
|
|
| def unpack(raw: np.ndarray) -> dict[str, Any]: |
| parameters: dict[str, Any] = {} |
| for value, (name, label) in zip(raw, specs): |
| if label is None: |
| parameters[name] = float(value) |
| else: |
| parameters.setdefault(name, {})[label] = float(value) |
| return parameters |
|
|
| def objective(raw: np.ndarray) -> float: |
| try: |
| prediction = np.asarray( |
| expression.evaluate(values, parameters=unpack(raw), num_nodes=num_nodes), |
| dtype=float, |
| ) |
| if prediction.shape != target.shape: |
| prediction = np.broadcast_to(prediction, target.shape) |
| loss = np.mean((prediction - target) ** 2) |
| return float(loss) if np.isfinite(loss) else 1e100 |
| except (FloatingPointError, OverflowError, ValueError): |
| return 1e100 |
|
|
| result = minimize( |
| objective, |
| np.asarray(vector, dtype=float), |
| method=method, |
| options=dict(options or {}), |
| ) |
| parameters = unpack(result.x) |
| return FitResult( |
| expression=bind_parameters(expression, parameters), |
| parameters=parameters, |
| loss=float(result.fun), |
| success=bool(result.success), |
| message=str(result.message), |
| n_iter=int(getattr(result, "nit", 0)), |
| ) |
|
|
