sindy-exp 0.2.1__tar.gz → 0.2.2__tar.gz

{sindy_exp-0.2.1 → sindy_exp-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sindy-exp
-Version: 0.2.1
+Version: 0.2.2
 Summary: A basic library for constructing dynamics experiments
 Author-email: Jake Stevens-Haas <jacob.stevens.haas@gmail.com>
 License: MIT License

{sindy_exp-0.2.1 → sindy_exp-0.2.2}/pyproject.toml

@@ -95,6 +95,9 @@ markers = ["slow"]
 files = [
   "src/sindy_exp/__init__.py",
   "src/sindy_exp/_utils.py",
+  "src/sindy_exp/_diffrax_solver.py",
+  "src/sindy_exp/_odes.py",
+  "src/sindy_exp/_typing.py",
   "tests/test_all.py",
 ]
 

{sindy_exp-0.2.1 → sindy_exp-0.2.2}/src/sindy_exp/__init__.py

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

{sindy_exp-0.2.1 → sindy_exp-0.2.2}/src/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, ProbData
 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[ProbData], list[dict[sp.Expr, float]]]]:
     """Generate random training and test data
 
     An Experiment step according to the mitosis experiment runner.
@@ -146,7 +139,7 @@ def gen_data(
         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},
     }
@@ -200,3 +193,114 @@ 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)
+
+    @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])
+
+
+@_register_dyst
+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])
+
+
+@_register_dyst
+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])
+
+
+@_register_dyst
+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])
+
+
+@_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)
+
+    @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])
+
+
+@_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)
+
+    @staticmethod
+    def _rhs(x, v, t: float, a) -> np.ndarray:
+        dxdt = v
+        dvdt = a
+        return np.array([dxdt, dvdt])

{sindy_exp-0.2.1 → sindy_exp-0.2.2}/src/sindy_exp/_diffrax_solver.py

@@ -85,20 +85,3 @@ def _gen_data_jax(
 
 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-0.2.1 → sindy_exp-0.2.2}/src/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, Callable, 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,6 +12,7 @@ from ._plotting import (
 )
 from ._typing import (
     DynamicsTrialData,
+    ExperimentResult,
     FullDynamicsTrialData,
     ProbData,
     SINDyTrialUpdate,
@@ -41,7 +39,6 @@ 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(
@@ -68,112 +65,30 @@ def _add_forcing(
     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[ProbData], 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[ProbData], 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,
+    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:
@@ -184,11 +99,8 @@ def fit_eval(
         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)
     trajectories, true_equations = data
     input_features = trajectories[0].input_features
 
@@ -198,12 +110,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 +128,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,
@@ -264,9 +181,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-0.2.1 → sindy_exp-0.2.2}/src/sindy_exp/_typing.py

@@ -7,6 +7,7 @@ from typing import (
     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: float
+    data: T
+    main: float
+
+
 class _BaseSINDy(Protocol):
     optimizer: Any
     feature_library: Any

{sindy_exp-0.2.1 → sindy_exp-0.2.2}/src/sindy_exp/_utils.py

@@ -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()
@@ -279,51 +283,6 @@ 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
 ) -> SINDyTrialUpdate:

{sindy_exp-0.2.1 → sindy_exp-0.2.2}/src/sindy_exp.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sindy-exp
-Version: 0.2.1
+Version: 0.2.2
 Summary: A basic library for constructing dynamics experiments
 Author-email: Jake Stevens-Haas <jacob.stevens.haas@gmail.com>
 License: MIT License

{sindy_exp-0.2.1 → sindy_exp-0.2.2}/src/sindy_exp.egg-info/SOURCES.txt

@@ -16,6 +16,7 @@ src/sindy_exp/_plotting.py
 src/sindy_exp/_typing.py
 src/sindy_exp/_utils.py
 src/sindy_exp/addl_attractors.json
+src/sindy_exp/py.typed
 src/sindy_exp.egg-info/PKG-INFO
 src/sindy_exp.egg-info/SOURCES.txt
 src/sindy_exp.egg-info/dependency_links.txt

{sindy_exp-0.2.1 → sindy_exp-0.2.2}/tests/test_all.py

@@ -180,7 +180,7 @@ def test_gen_data(rhs_name, array_namespace, jax_cpu_only):
     result = gen_data(
         rhs_name, t_end=0.1, noise_abs=0.01, seed=42, array_namespace=array_namespace
     )["data"]
-    trajectories = result["trajectories"]
+    trajectories = result[0]
     assert len(trajectories) == 1
     traj = trajectories[0]
     assert traj.x_train.shape == traj.x_train_true_dot.shape