sindy-exp 0.2.1__py3-none-any.whl → 0.3.0__py3-none-any.whl

sindy_exp/__init__.py

@@ -7,14 +7,16 @@ from ._plotting import (
     plot_test_trajectory,
     plot_training_data,
 )
-from ._typing import DynamicsTrialData, ProbData
+from ._typing import DynamicsTrialData, FullDynamicsTrialData, ProbData, SimProbData
 from ._utils import coeff_metrics, integration_metrics, pred_metrics
 
 __all__ = [
     "gen_data",
     "fit_eval",
     "ProbData",
+    "SimProbData",
     "DynamicsTrialData",
+    "FullDynamicsTrialData",
     "coeff_metrics",
     "pred_metrics",
     "integration_metrics",

sindy_exp/_data.py

@@ -1,15 +1,17 @@
+from importlib import resources
 from logging import getLogger
-from typing import Any, Callable, Optional, cast
+from typing import Callable, Optional, cast
 
 import dysts.flows
 import dysts.systems
 import numpy as np
 import scipy
+import sympy as sp
+from dysts.base import DynSys
 
 from ._dysts_to_sympy import dynsys_to_sympy
-from ._odes import SHO, CubicHO, Hopf, Kinematics, LotkaVolterra, VanDerPol
 from ._plotting import plot_training_data
-from ._typing import Float1D, ProbData
+from ._typing import ExperimentResult, Float1D, SimProbData
 from ._utils import _sympy_expr_to_feat_coeff
 
 try:
@@ -21,21 +23,12 @@ except ImportError:
 
 INTEGRATOR_KEYWORDS = {"rtol": 1e-12, "method": "LSODA", "atol": 1e-12}
 MOD_LOG = getLogger(__name__)
+LOCAL_DYNAMICS_PATH = resources.files("sindy_exp").joinpath("addl_attractors.json")
 
 ODE_CLASSES = {
     klass.lower(): getattr(dysts.flows, klass)
     for klass in dysts.systems.get_attractor_list()
 }
-ODE_CLASSES.update(
-    {
-        "lotkavolterra": LotkaVolterra,
-        "sho": SHO,
-        "cubicho": CubicHO,
-        "hopf": Hopf,
-        "vanderpol": VanDerPol,
-        "kinematics": Kinematics,
-    }
-)
 
 
 def gen_data(
@@ -49,7 +42,7 @@ def gen_data(
     t_end: float = 10,
     display: bool = False,
     array_namespace: str = "numpy",
-) -> dict[str, Any]:
+) -> ExperimentResult[tuple[list[SimProbData], list[dict[sp.Expr, float]]]]:
     """Generate random training and test data
 
     An Experiment step according to the mitosis experiment runner.
@@ -82,7 +75,7 @@ def gen_data(
     coeff_true = _sympy_expr_to_feat_coeff(sp_expr)
     rhsfunc = lambda t, X: dyst_sys.rhs(X, t)  # noqa: E731
     try:
-        x0_center = dyst_sys.ic
+        x0_center = cast(Float1D, dyst_sys.ic)
     except KeyError:
         x0_center = np.zeros((len(input_features)), dtype=np.float64)
     try:
@@ -95,7 +88,7 @@ def gen_data(
         noise_abs = 0.1
 
     MOD_LOG.info(f"Generating {n_trajectories} trajectories of f{system}")
-    prob_data_list: list[ProbData] = []
+    prob_data_list: list[SimProbData] = []
     if array_namespace == "numpy":
         feature_names = [feat.name for feat in input_features]
         for _ in range(n_trajectories):
@@ -115,20 +108,20 @@ def gen_data(
             prob_data_list.append(prob)
     elif array_namespace == "jax":
         try:
-            globals()["_gen_data_jax"]
-        except KeyError:
+            jax  # type: ignore
+        except NameError:
             raise ImportError(
                 "jax data generation requested but diffrax or sympy2jax not"
                 " installed"
             )
-        this_seed = jax.random.PRNGKey(seed)
+        this_seed = jax.random.PRNGKey(seed)  # type: ignore
         for _ in range(n_trajectories):
-            this_seed, _ = jax.random.split(this_seed)
-            prob = _gen_data_jax(
+            this_seed, _ = jax.random.split(this_seed)  # type: ignore
+            prob = _gen_data_jax(  # type: ignore
                 sp_expr,
                 input_features,
                 this_seed,
-                x0_center=x0_center,
+                x0_center=x0_center,  # type: ignore # numpy->jax
                 nonnegative=nonnegative,
                 ic_stdev=ic_stdev,
                 noise_abs=noise_abs,
@@ -143,10 +136,12 @@ def gen_data(
         )
     if display and prob_data_list:
         sample = prob_data_list[0]
+        assert sample.x_train_true is not None  # typing
         figs = plot_training_data(sample.t_train, sample.x_train, sample.x_train_true)
         figs[0].suptitle("Sample Trajectory")
+
     return {
-        "data": {"trajectories": prob_data_list, "coeff_true": coeff_true},
+        "data": (prob_data_list, coeff_true),
         "main": f"{n_trajectories} trajectories of {rhsfunc}",
         "metrics": {"rel_noise": noise_rel, "abs_noise": noise_abs},
     }
@@ -163,7 +158,7 @@ def _gen_data(
     nonnegative: bool,
     dt: float,
     t_end: float,
-) -> ProbData:
+) -> SimProbData:
     rng = np.random.default_rng(seed)
     t_train = np.arange(0, t_end, dt)
     t_train_span = (t_train[0], t_train[-1])
@@ -187,8 +182,9 @@ def _gen_data(
         noise_abs = np.sqrt(_signal_avg_power(x_train) * noise_rel)
     x_train = x_train + cast(float, noise_abs) * rng.standard_normal(x_train.shape)
 
-    return ProbData(
-        dt, t_train, x_train, x_train_true, x_train_true_dot, input_features
+    assert noise_abs is not None  # typing
+    return SimProbData(
+        t_train, x_train, input_features, x_train_true, x_train_true_dot, noise_abs
     )
 
 
@@ -200,3 +196,116 @@ def _max_amplitude(signal: np.ndarray, axis: int) -> float:
 
 def _signal_avg_power(signal: np.ndarray) -> float:
     return np.square(signal).mean()
+
+
+def _register_dyst(klass: type[DynSys]) -> type[DynSys]:
+    """Register a custom dysts DynSys class for use in sindy_exp data generation."""
+    ODE_CLASSES[klass.__name__.lower()] = klass
+    return klass
+
+
+@_register_dyst
+class LotkaVolterra(DynSys):
+    """Lotka-Volterra (predator-prey) dynamical system."""
+
+    nonnegative = True
+
+    def __init__(self):
+        super().__init__(metadata_path=LOCAL_DYNAMICS_PATH)  # type: ignore # dysts
+
+    @staticmethod
+    def _rhs(  # type: ignore # dysts
+        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])
+
+
+@_register_dyst
+class Hopf(DynSys):
+    """Hopf normal form dynamical system."""
+
+    def __init__(self):
+        super().__init__(metadata_path=LOCAL_DYNAMICS_PATH)  # type: ignore # dysts
+
+    @staticmethod
+    def _rhs(x, y, t: float, mu, omega, A) -> np.ndarray:  # type: ignore # dysts
+        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])
+
+
+@_register_dyst
+class SHO(DynSys):
+    """Linear damped simple harmonic oscillator"""
+
+    def __init__(self):
+        super().__init__(metadata_path=LOCAL_DYNAMICS_PATH)  # type: ignore # dysts
+
+    @staticmethod
+    def _rhs(x, y, t: float, a, b, c, d) -> np.ndarray:  # type: ignore # dysts
+        dxdt = a * x + b * y
+        dydt = c * x + d * y
+        return np.array([dxdt, dydt])
+
+
+@_register_dyst
+class CubicHO(DynSys):
+    """Cubic damped harmonic oscillator."""
+
+    def __init__(self):
+        super().__init__(metadata_path=LOCAL_DYNAMICS_PATH)  # type: ignore # dysts
+
+    @staticmethod
+    def _rhs(x, y, t: float, a, b, c, d) -> np.ndarray:  # type: ignore # dysts
+        dxdt = a * x**3 + b * y**3
+        dydt = c * x**3 + d * y**3
+        return np.array([dxdt, dydt])
+
+
+@_register_dyst
+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)  # type: ignore # dysts
+
+    @staticmethod
+    def _rhs(x, x_dot, t: float, mu) -> np.ndarray:  # type: ignore # dysts
+        dxdt = x_dot
+        dx2dt2 = mu * (1 - x**2) * x_dot - x
+        return np.array([dxdt, dx2dt2])
+
+
+@_register_dyst
+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)  # type: ignore # dysts
+
+    @staticmethod
+    def _rhs(x, v, t: float, a) -> np.ndarray:  # type: ignore # dysts
+        dxdt = v
+        dvdt = a
+        return np.array([dxdt, dvdt])

sindy_exp/_diffrax_solver.py

@@ -6,7 +6,7 @@ import jax.numpy as jnp
 import sympy2jax
 from sympy import Expr, Symbol
 
-from ._typing import ProbData
+from ._typing import SimProbData
 
 jax.config.update("jax_enable_x64", True)
 
@@ -22,7 +22,7 @@ def _gen_data_jax(
     nonnegative: bool,
     dt: float,
     t_end: float,
-) -> ProbData:
+) -> SimProbData:
     rhstree = sympy2jax.SymbolicModule(exprs)
 
     def ode_sys(t, state, args):
@@ -71,6 +71,8 @@ def _gen_data_jax(
     if noise_abs is None:
         assert noise_rel is not None  # force type narrowing
         noise_abs = float(jnp.sqrt(_signal_avg_power(x_train_true)) * noise_rel)
+    else:
+        noise_rel = noise_abs / float(jnp.sqrt(_signal_avg_power(x_train_true)))
 
     x_train = x_train_true + jax.random.normal(key, x_train_true.shape) * noise_abs
 
@@ -78,27 +80,16 @@ def _gen_data_jax(
     x_train_true_dot = jnp.array([ode_sys(0, xi, None) for xi in x_train_true])
 
     stringy_features = [sym.name for sym in input_features]
-    return ProbData(
-        dt, t_train, x_train, x_train_true, x_train_true_dot, stringy_features, sol
+    return SimProbData(
+        t_train,  # type: ignore # jax->numpy
+        x_train,  # type: ignore # jax->numpy
+        stringy_features,
+        x_train_true,  # type: ignore # jax->numpy
+        x_train_true_dot,  # type: ignore # jax->numpy
+        noise_abs,
+        sol,
     )
 
 
 def _signal_avg_power(signal: jax.Array) -> jax.Array:
     return jnp.square(signal).mean()
-
-
-## %  # noqa:E266
-if __name__ == "__main__":
-    # Debug example
-    from sindy_exp._data import gen_data
-
-    data_dict = gen_data(
-        "valliselnino",
-        seed=50,
-        n_trajectories=1,
-        ic_stdev=3,
-        noise_rel=0.1,
-        display=True,
-        array_namespace="jax",
-    )
-    print(data_dict["input_features"])

sindy_exp/_odes.py

@@ -1,12 +1,9 @@
-from importlib import resources
 from logging import getLogger
-from typing import Callable, TypeVar
+from typing import Any, Literal, TypeVar, cast, overload
 
 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,
@@ -15,8 +12,9 @@ from ._plotting import (
 )
 from ._typing import (
     DynamicsTrialData,
+    ExperimentResult,
     FullDynamicsTrialData,
-    ProbData,
+    SimProbData,
     SINDyTrialUpdate,
     _BaseSINDy,
 )
@@ -41,154 +39,46 @@ metric_ordering = {
 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
-    """
+@overload
+def fit_eval(
+    data: tuple[list[SimProbData], list[dict[sp.Expr, float]]],
+    model: _BaseSINDy,
+    simulations: Literal[False],
+    display: bool,
+) -> ExperimentResult[DynamicsTrialData]: ...
 
-    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])
+@overload
+def fit_eval(
+    data: tuple[list[SimProbData], list[dict[sp.Expr, float]]],
+    model: _BaseSINDy,
+    simulations: Literal[True],
+    display: bool,
+) -> ExperimentResult[FullDynamicsTrialData]: ...
 
 
 def fit_eval(
-    data: tuple[list[ProbData], list[dict[sp.Expr, float]]],
-    model: _BaseSINDy,
+    data: tuple[list[SimProbData], list[dict[sp.Expr, float]]],
+    model: Any,
     simulations: bool = True,
     display: bool = True,
-    return_all: bool = True,
-) -> dict | tuple[dict, DynamicsTrialData | FullDynamicsTrialData]:
+) -> ExperimentResult:
     """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
+            a list of SimProbData 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.
     """
-
+    model = cast(_BaseSINDy, model)
+    for trajectory in data[0]:
+        assert trajectory.x_train_true is not None
     trajectories, true_equations = data
     input_features = trajectories[0].input_features
 
@@ -198,12 +88,16 @@ def fit_eval(
 
     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):
+
+    # Special workaround for pysindy's legacy WeakPDELibrary
+    if hasattr(model.feature_library, "K"):
         # 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(
+
+    # Special workaround for pysindy's bad (soon to be legacy) differentiation API
+    if hasattr(model, "differentiation_method") and hasattr(
         model.differentiation_method, "smoothed_x_"
     ):
         smooth_x = []
@@ -212,6 +106,7 @@ def fit_eval(
             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,
@@ -226,7 +121,7 @@ def fit_eval(
         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)
+            sim = _simulate_test_data(model, traj.t_train, traj.x_train_true)
             sims.append(sim)
             integration_metric_list.append(
                 integration_metrics(
@@ -237,9 +132,9 @@ def fit_eval(
                 )
             )
 
-        agg_integration_metrics: dict[str, float | np.floating] = {}
+        agg_integration_metrics: dict[str, float] = {}
         for key in integration_metric_list[0].keys():
-            values = [m[key] for m in integration_metric_list]
+            values = cast(list[float], [m[key] for m in integration_metric_list])
             agg_integration_metrics[key] = float(np.mean(values))
         metrics.update(agg_integration_metrics)
 
@@ -264,9 +159,8 @@ def fit_eval(
                     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
+
+    return {"metrics": metrics, "data": trial_data, "main": metrics["main"]}
 
 
 def plot_ode_panel(trial_data: DynamicsTrialData):

sindy_exp/_plotting.py

@@ -57,7 +57,7 @@ def plot_coefficients(
     feature_names: Sequence[str],
     ax: Axes,
     **heatmap_kws,
-) -> None:
+) -> Axes:
     """Plot a set of dynamical system coefficients in a heatmap.
 
     Args:
@@ -162,6 +162,7 @@ def _compare_coefficient_plots_impl(
                 1, 2, figsize=(1.9 * n_cols, 8), sharey=True, sharex=True
             )
             fig.tight_layout()
+            assert axs is not None  # type narrowing
 
         vmax = signed_root(max_val)
 
@@ -275,7 +276,12 @@ def _plot_training_trajectory(
     """
     if x_train.shape[1] == 2:
         ax.plot(
-            x_true[:, 0], x_true[:, 1], ".", label="True", color=COLOR.TRUE, **PLOT_KWS
+            x_true[:, 0],
+            x_true[:, 1],
+            ".",
+            label="True",
+            color=COLOR.TRUE,
+            **PLOT_KWS,  # type: ignore[arg-type]
         )
         ax.plot(
             x_train[:, 0],
@@ -283,7 +289,7 @@ def _plot_training_trajectory(
             ".",
             label="Measured",
             color=COLOR.MEAS,
-            **PLOT_KWS,
+            **PLOT_KWS,  # type: ignore[arg-type]
         )
         if (
             x_smooth is not None
@@ -295,7 +301,7 @@ def _plot_training_trajectory(
                 ".",
                 label="Smoothed",
                 color=COLOR.EST,
-                **PLOT_KWS,
+                **PLOT_KWS,  # type: ignore[arg-type]
             )
         if labels:
             ax.set(xlabel="$x_0$", ylabel="$x_1$")
@@ -308,7 +314,7 @@ def _plot_training_trajectory(
             x_true[:, 2],
             color=COLOR.TRUE,
             label="True values",
-            **PLOT_KWS,
+            **PLOT_KWS,  # type: ignore[arg-type]
         )
 
         ax.plot(

sindy_exp/_typing.py

@@ -1,12 +1,13 @@
 from collections import defaultdict
+from collections.abc import Mapping
 from dataclasses import dataclass
 from typing import (
     Any,
     Callable,
     Literal,
-    NamedTuple,
     Optional,
     Protocol,
+    TypedDict,
     TypeVar,
     overload,
 )
@@ -27,6 +28,14 @@ FloatND = np.ndarray[Shape, np.dtype[np.floating[NBitBase]]]
 TrajectoryType = TypeVar("TrajectoryType", list[np.ndarray], np.ndarray)
 
 
+class ExperimentResult[T](TypedDict):
+    """Results from a SINDy ODE experiment."""
+
+    metrics: Mapping[str, float | None]
+    data: T
+    main: object
+
+
 class _BaseSINDy(Protocol):
     optimizer: Any
     feature_library: Any
@@ -62,23 +71,38 @@ class _BaseSINDy(Protocol):
         self, precision: int, fmt: Literal["sympy"]
     ) -> list[dict[Expr, float]]: ...
 
+    @overload
+    def print(self, **kwargs) -> None: ...
+
+    @overload
     def print(self, precision: int, **kwargs) -> None: ...
 
     def get_feature_names(self) -> list[str]: ...
 
 
-class ProbData(NamedTuple):
-    """Data bundle for a single trajectory.
+@dataclass
+class ProbData:
+    """Represents a single trajectory's data.
 
-    Represents a trajectory's training data and associated metadata.
+    For measured data, only t_train, x_train, and input_features are required.
     """
 
-    dt: float
     t_train: Float1D
     x_train: Float2D
+    input_features: list[str]
+
+
+@dataclass
+class SimProbData(ProbData):
+    """For simulated data, the noiseless trajectory is known.
+
+    Optionally includes the integrator solution object for evaluating
+    at other points.
+    """
+
     x_train_true: Float2D
     x_train_true_dot: Float2D
-    input_features: list[str]
+    noise_abs: float
     integrator: Optional[Any] = None  # diffrax.Solution
 
 
@@ -138,7 +162,7 @@ class NestedDict(defaultdict):
 
 @dataclass
 class DynamicsTrialData:
-    trajectories: list[ProbData]
+    trajectories: list[SimProbData]
     true_equations: list[dict[sp.Expr, float]]
     sindy_equations: list[dict[sp.Expr, float]]
     model: _BaseSINDy

sindy_exp/_utils.py

@@ -149,7 +149,7 @@ def opt_lookup(kind):
 def coeff_metrics(
     coeff_est_dicts: list[dict[sp.Expr, float]],
     coeff_true_dicts: list[dict[sp.Expr, float]],
-) -> dict[str, float | np.floating]:
+) -> dict[str, float]:
     """Compute coefficient metrics from aligned coefficient dictionaries.
 
     Both arguments are expected to be lists of coefficient dictionaries sharing
@@ -182,14 +182,18 @@ def coeff_metrics(
             coefficients[row_ind, col_ind] = est_row[feat]
 
     metrics: dict[str, float | np.floating] = {}
-    metrics["coeff_precision"] = sklearn.metrics.precision_score(
-        coeff_true.flatten() != 0, coefficients.flatten() != 0
+    metrics["coeff_precision"] = float(
+        sklearn.metrics.precision_score(
+            coeff_true.flatten() != 0, coefficients.flatten() != 0
+        )
     )
-    metrics["coeff_recall"] = sklearn.metrics.recall_score(
-        coeff_true.flatten() != 0, coefficients.flatten() != 0
+    metrics["coeff_recall"] = float(
+        sklearn.metrics.recall_score(
+            coeff_true.flatten() != 0, coefficients.flatten() != 0
+        )
     )
-    metrics["coeff_f1"] = sklearn.metrics.f1_score(
-        coeff_true.flatten() != 0, coefficients.flatten() != 0
+    metrics["coeff_f1"] = float(
+        sklearn.metrics.f1_score(coeff_true.flatten() != 0, coefficients.flatten() != 0)
     )
     metrics["coeff_mse"] = sklearn.metrics.mean_squared_error(
         coeff_true.flatten(), coefficients.flatten()
@@ -198,7 +202,7 @@ def coeff_metrics(
         coeff_true.flatten(), coefficients.flatten()
     )
     metrics["main"] = metrics["coeff_f1"]
-    return metrics
+    return {k: float(v) for k, v in metrics.items()}
 
 
 def pred_metrics(
@@ -279,53 +283,8 @@ def unionize_coeff_dicts(
     return true_aligned, est_aligned
 
 
-def make_model(
-    input_features: list[str],
-    dt: float,
-    diff_params: dict | ps.BaseDifferentiation,
-    feat_params: dict | ps.feature_library.base.BaseFeatureLibrary,
-    opt_params: dict | ps.BaseOptimizer,
-) -> ps.SINDy:
-    """Build a model with object parameters dictionaries
-
-    e.g. {"kind": "finitedifference"} instead of FiniteDifference()
-    """
-
-    def finalize_param(lookup_func, pdict, lookup_key):
-        try:
-            cls_name = pdict.pop(lookup_key)
-        except AttributeError:
-            cls_name = pdict.vals.pop(lookup_key)
-            pdict = pdict.vals
-
-        param_cls = lookup_func(cls_name)
-        param_final = param_cls(**pdict)
-        pdict[lookup_key] = cls_name
-        return param_final
-
-    if isinstance(diff_params, ps.BaseDifferentiation):
-        diff = diff_params
-    else:
-        diff = finalize_param(diff_lookup, diff_params, "diffcls")
-    if isinstance(feat_params, ps.feature_library.base.BaseFeatureLibrary):
-        features = feat_params
-    else:
-        features = finalize_param(feature_lookup, feat_params, "featcls")
-    if isinstance(opt_params, ps.BaseOptimizer):
-        opt = opt_params
-    else:
-        opt = finalize_param(opt_lookup, opt_params, "optcls")
-    return ps.SINDy(
-        differentiation_method=diff,
-        optimizer=opt,
-        t_default=dt,  # type: ignore
-        feature_library=features,
-        feature_names=input_features,
-    )
-
-
 def _simulate_test_data(
-    model: _BaseSINDy, dt: float, x_test: Float2D
+    model: _BaseSINDy, t_test: Float1D, x_test: Float2D
 ) -> SINDyTrialUpdate:
     """Add simulation data to grid_data
 
@@ -333,7 +292,6 @@ def _simulate_test_data(
     Returns:
         Complete GridPointData
     """
-    t_test = cast(Float1D, np.arange(0, len(x_test) * dt, step=dt))
     t_sim = t_test
     try:
 

{sindy_exp-0.2.1.dist-info → sindy_exp-0.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sindy-exp
-Version: 0.2.1
+Version: 0.3.0
 Summary: A basic library for constructing dynamics experiments
 Author-email: Jake Stevens-Haas <jacob.stevens.haas@gmail.com>
 License: MIT License
@@ -34,7 +34,7 @@ Classifier: Intended Audience :: Science/Research
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Natural Language :: English
 Classifier: Operating System :: POSIX :: Linux
-Requires-Python: >=3.10
+Requires-Python: >=3.12
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: matplotlib
@@ -43,6 +43,7 @@ Requires-Dist: seaborn
 Requires-Dist: scipy
 Requires-Dist: sympy
 Requires-Dist: dysts
+Requires-Dist: scikit-learn
 Provides-Extra: jax
 Requires-Dist: jax[cuda12]; extra == "jax"
 Requires-Dist: diffrax; extra == "jax"
@@ -64,14 +65,20 @@ Requires-Dist: tomli; extra == "dev"
 Requires-Dist: pysindy>=2.1.0; extra == "dev"
 Dynamic: license-file
 
-# Dynamics Experiments
+# Overview
 
-A library for constructing dynamics experiments.
-This includes data generation and plotting/evaluation.
+A library for constructing dynamics experiments from the dynamics models in the `dysts` package.
+This includes data generation and model evaluation.
+The first contribution is the static typing of trajectory data (`ProbData`) that, I believe, provides the necessary information to be useful in evaluating a wide variety of dynamics/time-series learning methods.
+The second contribution is the collection of utility functions for designing dynamics learning experiments.
+The third contribution is the collection of such experiments for evaluating dynamics/time-series learning models that meet the `BaseSINDy` API.
+
+It aims to (a) be amenable to both `numpy` and `jax` arrays, (b) be usable by any dynamics/time-series learning models that meet the `BaseSINDy` or scikit-time API.
+Internally, this package is used/will be used in benchmarking pysindy runtime/memory usage and choosing default hyperparameters.
 
 ## Getting started
 
-It's not yet on PyPI, so install it with `pip install sindy_exp @ git+https://github.com/Jacob-Stevens-Haas/sindy-experiments`
+Install with `pip install sindy-exp` or `pip install sindy-exp[jax]`.
 
 Generate data
 
@@ -86,26 +93,26 @@ Evaluate your SINDy-like model with:
 A list of available ODE systems can be found in `ODE_CLASSES`, which includes most
 of the systems from the [dysts package](https://pypi.org/project/dysts/) as well as some non-chaotic systems.
 
-## ODE representation
+## ODE & Data Model
+
+Generated or measured data has the dataclass type `ProbData` or `SimProbData`, respectively,
+to indicate whether it includes ground truth information and a noise level.
+If the data is generated in jax, it will have an integrator that can later be used to evaluate the true data on collocation points.
 
 We deal primarily with autonomous ODE systems of the form:
 
     dx/dt = sum_i f_i(x)
 
-Thus, we represent ODE systems as a list of right-hand side expressions.
+We represent ODE systems as a list of right-hand side expressions.
 Each element is a dictionary mapping a term (Sympy expression) to its coefficient.
+Thus, the rhs of an ODE is of type: `list[dict[sympy.Expr, float]]`
 
 ## Other useful imports, compatibility, and extensions
 
-This is built to be compatible with dynamics learning models that follow the
-pysindy _BaseSINDy interface.
-The experiments are also built to be compatible with the `mitosis` tool,
-an experiment runner.
-To integrate your own experiments or data generation in a way that is compatible,
-see the `ProbData` and `DynamicsTrialData` classes.
-For plotting tools, see `plot_coefficients`, `compare_coefficient_plots_from_dicts`,
-`plot_test_trajectory`, `plot_training_data`, and `COLOR`.
-For metrics, see `coeff_metrics`, `pred_metrics`, and `integration_metrics`.
+* The experiments are built to be compatible with the `mitosis` tool, an experiment runner.  Mitosis is not a dependency, however, to allow using other experiment runners.
+* To integrate your own experiments or data generation in a way that is compatible, see the `ProbData`, `SimProbData`, `DynamicsTrialData`, and `FullDynamicsTrialData` classes.
+* For plotting tools, see `plot_coefficients`, `compare_coefficient_plots_from_dicts`, `plot_test_trajectory`, `plot_training_data`, and `COLOR`.
+* For evaluation of models, see `coeff_metrics`, `pred_metrics`, and `integration_metrics`.
 
 ![3d plot](images/composite.png)
 ![1d plot](images/1d.png)

sindy_exp-0.3.0.dist-info/RECORD

@@ -0,0 +1,15 @@
+sindy_exp/__init__.py,sha256=DJuJM0QVe7zmpOHY--7U8m4MCU8XkL_WKGD1H4skSM0,773
+sindy_exp/_data.py,sha256=GwGAnLvEgxRjrwGLTe8RC6GAIOocJtDBey1Dhj8Gyk4,9744
+sindy_exp/_diffrax_solver.py,sha256=HvhmUT9u9A06c6XRSgewdbKl0JEVB312UMmI_Lb4DVo,2634
+sindy_exp/_dysts_to_sympy.py,sha256=d_rvnfayOmFcGn4bZRJCfNGFO6yS1mw2QmBbOdWZwxg,15654
+sindy_exp/_odes.py,sha256=E7BdvFoPI7kxNh0Q9WMpE7MsKGbke5zDjX5Ml31amkQ,5927
+sindy_exp/_plotting.py,sha256=mfsDz_693XQS6tMZ8iQluRJqZMcmxyoZeIMP4DfJX2s,17688
+sindy_exp/_typing.py,sha256=hxIxBx52oKMjd1LA8PItwrXlP-PtNGLoZm6oIMZi18I,4536
+sindy_exp/_utils.py,sha256=RRFtk1iFWeo0Rpw1vYhoNSoABLJcRmGo1sIgTZ9lnWM,11299
+sindy_exp/addl_attractors.json,sha256=KXoHWekFoa4KctjLCqcj_BpLBhXV0zlYrpgxV-uObwE,2928
+sindy_exp/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+sindy_exp-0.3.0.dist-info/licenses/LICENSE,sha256=ubi77tIG3RVrqo0Z8cK91D4KZePQs-W1J-vJ-LkVOmE,1075
+sindy_exp-0.3.0.dist-info/METADATA,sha256=qTE2umRXW5K9mCtTIt2u-vvtht92rxII90K94ZYaPOE,5669
+sindy_exp-0.3.0.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+sindy_exp-0.3.0.dist-info/top_level.txt,sha256=0-tKKdmxHG3IRccz463rOb6xTsVJD-v9c8zSDpTRr5E,10
+sindy_exp-0.3.0.dist-info/RECORD,,

sindy_exp-0.2.1.dist-info/RECORD

@@ -1,14 +0,0 @@
-sindy_exp/__init__.py,sha256=F1bJz9Gzk2nPB6DGjEa0qZcaQPTb-Yhh0ZnO9TcQci0,689
-sindy_exp/_data.py,sha256=X1WdFZ53FeqG1fwtnxQnCHUQqKtGSAvv3URWdcSavE8,6511
-sindy_exp/_diffrax_solver.py,sha256=YmKF-U1fkxW4z4AbXW_qm1sX6eEy1BGzKbqL63dxo2M,2705
-sindy_exp/_dysts_to_sympy.py,sha256=d_rvnfayOmFcGn4bZRJCfNGFO6yS1mw2QmBbOdWZwxg,15654
-sindy_exp/_odes.py,sha256=RN97nh1Y_6DQCLPfRTyCCJq3exxBq9pfVHVPTxr0na4,8830
-sindy_exp/_plotting.py,sha256=dpcqAXKzb0mSVl0p2WyMadVoGhQi43oL6ZqsbuheEuk,17470
-sindy_exp/_typing.py,sha256=hIOj270YkoITi_qJ8xTO1EaBeoVzcDidIGMXuC0QHzM,4046
-sindy_exp/_utils.py,sha256=RnMoNV_N7ubWVMbgdiJaRGo437DvWUcOIn4fP2rhwJI,12717
-sindy_exp/addl_attractors.json,sha256=KXoHWekFoa4KctjLCqcj_BpLBhXV0zlYrpgxV-uObwE,2928
-sindy_exp-0.2.1.dist-info/licenses/LICENSE,sha256=ubi77tIG3RVrqo0Z8cK91D4KZePQs-W1J-vJ-LkVOmE,1075
-sindy_exp-0.2.1.dist-info/METADATA,sha256=deZyEMmZibk24ZMU1Fc7s_qD-ibRwRSjfoNqW1enqpk,4514
-sindy_exp-0.2.1.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-sindy_exp-0.2.1.dist-info/top_level.txt,sha256=0-tKKdmxHG3IRccz463rOb6xTsVJD-v9c8zSDpTRr5E,10
-sindy_exp-0.2.1.dist-info/RECORD,,