sindy-exp 0.2.0__py3-none-any.whl

sindy_exp/_odes.py

@@ -0,0 +1,287 @@
+from importlib import resources
+from logging import getLogger
+from typing import Callable, TypeVar
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pysindy as ps
+import sympy as sp
+from dysts.base import DynSys
+
+from ._plotting import (
+    compare_coefficient_plots_from_dicts,
+    plot_test_trajectory,
+    plot_training_data,
+)
+from ._typing import (
+    DynamicsTrialData,
+    FullDynamicsTrialData,
+    ProbData,
+    SINDyTrialUpdate,
+    _BaseSINDy,
+)
+from ._utils import (
+    _simulate_test_data,
+    coeff_metrics,
+    integration_metrics,
+    unionize_coeff_dicts,
+)
+
+metric_ordering = {
+    "coeff_precision": "max",
+    "coeff_f1": "max",
+    "coeff_recall": "max",
+    "coeff_mae": "min",
+    "coeff_mse": "min",
+    "mse_plot": "min",
+    "mae_plot": "min",
+}
+
+
+T = TypeVar("T", bound=int)
+DType = TypeVar("DType", bound=np.dtype)
+MOD_LOG = getLogger(__name__)
+LOCAL_DYNAMICS_PATH = resources.files("sindy_exp").joinpath("addl_attractors.json")
+
+
+def _add_forcing(
+    forcing_func: Callable[[float], np.ndarray[tuple[T], DType]],
+    auto_func: Callable[
+        [float, np.ndarray[tuple[T], DType]], np.ndarray[tuple[T], DType]
+    ],
+) -> Callable[[float, np.ndarray], np.ndarray]:
+    """Add a time-dependent forcing term to a rhs func
+
+    Args:
+        forcing_func: The forcing function to add
+        auto_func: An existing rhs func for solve_ivp
+
+    Returns:
+        A rhs function for integration
+    """
+
+    def sum_of_terms(
+        t: float, state: np.ndarray[tuple[T], DType]
+    ) -> np.ndarray[tuple[T], DType]:
+        return np.array(forcing_func(t)) + np.array(auto_func(t, state))
+
+    return sum_of_terms
+
+
+class LotkaVolterra(DynSys):
+    """Lotka-Volterra (predator-prey) dynamical system."""
+
+    nonnegative = True
+
+    def __init__(self):
+        super().__init__(metadata_path=LOCAL_DYNAMICS_PATH)
+
+    @staticmethod
+    def _rhs(x, y, t: float, alpha, beta, gamma, delta) -> np.ndarray:
+        """LV dynamics
+
+        Args:
+            x: prey population
+            y: predator population
+            t: time (ignored, since autonomous)
+            alpha: prey growth rate
+            beta: predation rate
+            delta: predator reproduction rate
+            gamma: predator death rate
+        """
+        dxdt = alpha * x - beta * x * y
+        dydt = delta * x * y - gamma * y
+        return np.array([dxdt, dydt])
+
+
+class Hopf(DynSys):
+    """Hopf normal form dynamical system."""
+
+    def __init__(self):
+        super().__init__(metadata_path=LOCAL_DYNAMICS_PATH)
+
+    @staticmethod
+    def _rhs(x, y, t: float, mu, omega, A) -> np.ndarray:
+        dxdt = mu * x - omega * y - A * (x**3 + x * y**2)
+        dydt = omega * x + mu * y - A * (x**2 * y + y**3)
+        return np.array([dxdt, dydt])
+
+
+class SHO(DynSys):
+    """Linear damped simple harmonic oscillator"""
+
+    def __init__(self):
+        super().__init__(metadata_path=LOCAL_DYNAMICS_PATH)
+
+    @staticmethod
+    def _rhs(x, y, t: float, a, b, c, d) -> np.ndarray:
+        dxdt = a * x + b * y
+        dydt = c * x + d * y
+        return np.array([dxdt, dydt])
+
+
+class CubicHO(DynSys):
+    """Cubic damped harmonic oscillator."""
+
+    def __init__(self):
+        super().__init__(metadata_path=LOCAL_DYNAMICS_PATH)
+
+    @staticmethod
+    def _rhs(x, y, t: float, a, b, c, d) -> np.ndarray:
+        dxdt = a * x**3 + b * y**3
+        dydt = c * x**3 + d * y**3
+        return np.array([dxdt, dydt])
+
+
+class VanDerPol(DynSys):
+    """Van der Pol oscillator.
+
+    dx/dt = y
+    dy/dt = mu * (1 - x^2) * y - x
+    """
+
+    def __init__(self):
+        super().__init__(metadata_path=LOCAL_DYNAMICS_PATH)
+
+    @staticmethod
+    def _rhs(x, x_dot, t: float, mu) -> np.ndarray:
+        dxdt = x_dot
+        dx2dt2 = mu * (1 - x**2) * x_dot - x
+        return np.array([dxdt, dx2dt2])
+
+
+class Kinematics(DynSys):
+    """One-dimensional kinematics with constant acceleration.
+
+    dx/dt = v
+    dv/dt = a
+    """
+
+    def __init__(self):
+        super().__init__(metadata_path=LOCAL_DYNAMICS_PATH)
+
+    @staticmethod
+    def _rhs(x, v, t: float, a) -> np.ndarray:
+        dxdt = v
+        dvdt = a
+        return np.array([dxdt, dvdt])
+
+
+def fit_eval(
+    data: tuple[list[ProbData], list[dict[sp.Expr, float]]],
+    model: _BaseSINDy,
+    simulations: bool = True,
+    display: bool = True,
+    return_all: bool = True,
+) -> dict | tuple[dict, DynamicsTrialData | FullDynamicsTrialData]:
+    """Fit and evaluate a SINDy model on a set of trajectories.
+
+    Args:
+        data: Tuple of (trajectories, true_equations), where ``trajectories`` is
+            a list of ProbData objects and ``true_equations`` is a list of
+            dictionaries mapping SymPy symbols to their true coefficients for
+            each state coordinate.
+        model: A SINDy-like model implementing the _BaseSINDy protocol.
+        simulations: Whether to run forward simulations for evaluation.
+        display: Whether to generate plots as part of evaluation.
+        return_all: If True, return a dictionary containing metrics and the
+            assembled DynamicsTrialData; otherwise return only the metrics
+            dictionary.
+    """
+
+    trajectories, true_equations = data
+    input_features = trajectories[0].input_features
+
+    x_train = [traj.x_train for traj in trajectories]
+    t_train = [traj.t_train for traj in trajectories]
+    model.fit(x_train, t=t_train, feature_names=input_features)
+
+    MOD_LOG.info(f"Fitting a model: {model}")
+    coeff_true_dicts, coeff_est_dicts = unionize_coeff_dicts(model, true_equations)
+    if isinstance(model.feature_library, ps.WeakPDELibrary):
+        # WeakPDE library fails to simulate, so insert nonweak library
+        # to Pipeline and SINDy model.
+        inner_lib = model.feature_library.function_library
+        model.feature_library = inner_lib  # type: ignore  # TODO: Fix in pysindy
+    if isinstance(model, ps.SINDy) and hasattr(
+        model.differentiation_method, "smoothed_x_"
+    ):
+        smooth_x = []
+        for traj in trajectories:
+            model.differentiation_method(traj.x_train, t=traj.t_train)
+            smooth_x.append(model.differentiation_method.smoothed_x_)
+    else:  # using WeakPDELibrary
+        smooth_x = x_train
+    trial_data = DynamicsTrialData(
+        trajectories=trajectories,
+        true_equations=coeff_true_dicts,
+        sindy_equations=coeff_est_dicts,
+        model=model,
+        input_features=input_features,
+        smooth_train=smooth_x,
+    )
+    MOD_LOG.info(f"Evaluating a model: {model}")
+    metrics = coeff_metrics(coeff_est_dicts, coeff_true_dicts)
+    if simulations:
+        sims: list[SINDyTrialUpdate] = []
+        integration_metric_list: list[dict[str, float | np.floating]] = []
+        for traj in trajectories:
+            sim = _simulate_test_data(model, traj.dt, traj.x_train_true)
+            sims.append(sim)
+            integration_metric_list.append(
+                integration_metrics(
+                    model,
+                    traj.x_train_true,
+                    traj.t_train,
+                    traj.x_train_true_dot,
+                )
+            )
+
+        agg_integration_metrics: dict[str, float | np.floating] = {}
+        for key in integration_metric_list[0].keys():
+            values = [m[key] for m in integration_metric_list]
+            agg_integration_metrics[key] = float(np.mean(values))
+        metrics.update(agg_integration_metrics)
+
+        trial_data = FullDynamicsTrialData(sims=sims, **trial_data.__dict__)
+    if display:
+        plot_ode_panel(trial_data)
+        for i, traj in enumerate(trajectories):
+            fig_composite, fig_by_coord_1d = plot_training_data(
+                traj.t_train,
+                traj.x_train,
+                traj.x_train_true,
+                x_smooth=smooth_x[i],
+                coord_names=input_features,
+            )
+            if simulations:
+                # Overlay test trajectory time series on the coordinate-wise figure
+                plot_test_trajectory(
+                    traj.x_train_true,
+                    sims[i].x_sim,
+                    traj.t_train,
+                    sims[i].t_sim,
+                    figs=(fig_composite, fig_by_coord_1d),
+                    coord_names=input_features,
+                )
+    if return_all:
+        return {"metrics": metrics, "data": trial_data, "main": metrics["main"]}
+    return metrics
+
+
+def plot_ode_panel(trial_data: DynamicsTrialData):
+    trial_data.model.print()
+    compare_coefficient_plots_from_dicts(
+        trial_data.sindy_equations,
+        trial_data.true_equations,
+        input_features=[_texify(feat) for feat in trial_data.input_features],
+    )
+    plt.show()
+
+
+def _texify(input: str) -> str:
+    if input[0] != "$":
+        input = "$" + input
+    if input[-1] != "$":
+        input = input + "$"
+    return input