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

{sindy_exp-0.2.2 → sindy_exp-0.3.0}/.github/workflows/main.yaml

@@ -17,7 +17,7 @@ jobs:
     - name: "Set up Python"
       uses: actions/setup-python@v3
       with:
-          python-version: "3.10"
+          python-version: "3.12"
     - name: run pre-commit
       run: |
         pip install pre-commit
@@ -30,7 +30,7 @@ jobs:
     - name: "Set up Python"
       uses: actions/setup-python@v3
       with:
-          python-version: "3.10"
+          python-version: "3.12"
     - name: install dependencies
       run: |
         pip install -e .[dev,jax]
@@ -45,7 +45,7 @@ jobs:
       fail-fast: false
       max-parallel: 4
       matrix:
-          python-version: ["3.10", "3.11", "3.12"]
+          python-version: ["3.12", "3.13"]
     steps:
     - uses: actions/checkout@v3
     - name: Set up Python ${{ matrix.python-version }}

{sindy_exp-0.2.2 → sindy_exp-0.3.0}/.github/workflows/release.yml

@@ -27,7 +27,7 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v6
         with:
-          python-version: "3.10"
+          python-version: "3.12"
       - name: Install Build
         run: |
           python -m pip install --upgrade pip

{sindy_exp-0.2.2 → sindy_exp-0.3.0}/.gitignore

@@ -1,7 +1,6 @@
 commit-msg
 todo
 scratch/
-*.png
 debugme*.py
 trials/*
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sindy-exp
-Version: 0.2.2
+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/README.md

@@ -0,0 +1,51 @@
+# Overview
+
+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
+
+Install with `pip install sindy-exp` or `pip install sindy-exp[jax]`.
+
+Generate data
+
+    data = sindy_exp.data.gen_data("lorenz", num_trajectories=5, t_end=10.0, dt=0.01)["data]
+
+Evaluate your SINDy-like model with:
+
+    sindy_exp.odes.fit_eval(model, data)
+
+![Coefficient plots](images/coeff.png)
+
+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 & 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)
+
+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
+
+* 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.2.2 → sindy_exp-0.3.0}/pyproject.toml

@@ -7,7 +7,7 @@ name = "sindy-exp"
 dynamic = ["version"]
 description = "A basic library for constructing dynamics experiments"
 readme = "README.md"
-requires-python = ">=3.10"
+requires-python = ">=3.12"
 license = {file = "LICENSE"}
 keywords = ["Machine Learning", "Science", "Mathematics", "Experiments"]
 authors = [
@@ -29,6 +29,7 @@ dependencies = [
   "scipy",
   "sympy",
   "dysts",
+  "scikit-learn",
 ]
 
 [project.optional-dependencies]
@@ -94,25 +95,20 @@ markers = ["slow"]
 [tool.mypy]
 files = [
   "src/sindy_exp/__init__.py",
+  "src/sindy_exp/_data.py",
   "src/sindy_exp/_utils.py",
   "src/sindy_exp/_diffrax_solver.py",
   "src/sindy_exp/_odes.py",
+  "src/sindy_exp/_plotting.py",
   "src/sindy_exp/_typing.py",
   "tests/test_all.py",
 ]
+warn_unused_configs = true
 
 [[tool.mypy.overrides]]
-module="sklearn.*"
-ignore_missing_imports=true
+module = ["sindy_exp.*"]
+disable_error_code = ["import-untyped"]
 
 [[tool.mypy.overrides]]
-module="pysindy.*"
-ignore_missing_imports=true
-
-[[tool.mypy.overrides]]
-module="sympy.*"
-ignore_missing_imports=true
-
-[[tool.mypy.overrides]]
-module="scipy.*"
-ignore_missing_imports=true
+module = ["pysindy.*", "sympy.*"]
+ignore_missing_imports = true

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

@@ -7,13 +7,14 @@ from ._plotting import (
     plot_test_trajectory,
     plot_training_data,
 )
-from ._typing import DynamicsTrialData, FullDynamicsTrialData, 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",

{sindy_exp-0.2.2 → sindy_exp-0.3.0}/src/sindy_exp/_data.py

@@ -11,7 +11,7 @@ from dysts.base import DynSys
 
 from ._dysts_to_sympy import dynsys_to_sympy
 from ._plotting import plot_training_data
-from ._typing import ExperimentResult, Float1D, ProbData
+from ._typing import ExperimentResult, Float1D, SimProbData
 from ._utils import _sympy_expr_to_feat_coeff
 
 try:
@@ -42,7 +42,7 @@ def gen_data(
     t_end: float = 10,
     display: bool = False,
     array_namespace: str = "numpy",
-) -> ExperimentResult[tuple[list[ProbData], list[dict[sp.Expr, float]]]]:
+) -> ExperimentResult[tuple[list[SimProbData], list[dict[sp.Expr, float]]]]:
     """Generate random training and test data
 
     An Experiment step according to the mitosis experiment runner.
@@ -75,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:
@@ -88,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):
@@ -108,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,
@@ -136,8 +136,10 @@ 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": (prob_data_list, coeff_true),
         "main": f"{n_trajectories} trajectories of {rhsfunc}",
@@ -156,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])
@@ -180,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
     )
 
 
@@ -208,10 +211,12 @@ class LotkaVolterra(DynSys):
     nonnegative = True
 
     def __init__(self):
-        super().__init__(metadata_path=LOCAL_DYNAMICS_PATH)
+        super().__init__(metadata_path=LOCAL_DYNAMICS_PATH)  # type: ignore # dysts
 
     @staticmethod
-    def _rhs(x, y, t: float, alpha, beta, gamma, delta) -> np.ndarray:
+    def _rhs(  # type: ignore # dysts
+        x, y, t: float, alpha, beta, gamma, delta
+    ) -> np.ndarray:
         """LV dynamics
 
         Args:
@@ -233,10 +238,10 @@ class Hopf(DynSys):
     """Hopf normal form dynamical system."""
 
     def __init__(self):
-        super().__init__(metadata_path=LOCAL_DYNAMICS_PATH)
+        super().__init__(metadata_path=LOCAL_DYNAMICS_PATH)  # type: ignore # dysts
 
     @staticmethod
-    def _rhs(x, y, t: float, mu, omega, A) -> np.ndarray:
+    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])
@@ -247,10 +252,10 @@ class SHO(DynSys):
     """Linear damped simple harmonic oscillator"""
 
     def __init__(self):
-        super().__init__(metadata_path=LOCAL_DYNAMICS_PATH)
+        super().__init__(metadata_path=LOCAL_DYNAMICS_PATH)  # type: ignore # dysts
 
     @staticmethod
-    def _rhs(x, y, t: float, a, b, c, d) -> np.ndarray:
+    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])
@@ -261,10 +266,10 @@ class CubicHO(DynSys):
     """Cubic damped harmonic oscillator."""
 
     def __init__(self):
-        super().__init__(metadata_path=LOCAL_DYNAMICS_PATH)
+        super().__init__(metadata_path=LOCAL_DYNAMICS_PATH)  # type: ignore # dysts
 
     @staticmethod
-    def _rhs(x, y, t: float, a, b, c, d) -> np.ndarray:
+    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])
@@ -279,10 +284,10 @@ class VanDerPol(DynSys):
     """
 
     def __init__(self):
-        super().__init__(metadata_path=LOCAL_DYNAMICS_PATH)
+        super().__init__(metadata_path=LOCAL_DYNAMICS_PATH)  # type: ignore # dysts
 
     @staticmethod
-    def _rhs(x, x_dot, t: float, mu) -> np.ndarray:
+    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])
@@ -297,10 +302,10 @@ class Kinematics(DynSys):
     """
 
     def __init__(self):
-        super().__init__(metadata_path=LOCAL_DYNAMICS_PATH)
+        super().__init__(metadata_path=LOCAL_DYNAMICS_PATH)  # type: ignore # dysts
 
     @staticmethod
-    def _rhs(x, v, t: float, a) -> np.ndarray:
+    def _rhs(x, v, t: float, a) -> np.ndarray:  # type: ignore # dysts
         dxdt = v
         dvdt = a
         return np.array([dxdt, dvdt])

{sindy_exp-0.2.2 → sindy_exp-0.3.0}/src/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,8 +80,14 @@ 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,
     )
 
 

{sindy_exp-0.2.2 → sindy_exp-0.3.0}/src/sindy_exp/_odes.py

@@ -1,5 +1,5 @@
 from logging import getLogger
-from typing import Any, Callable, Literal, TypeVar, cast, overload
+from typing import Any, Literal, TypeVar, cast, overload
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -14,7 +14,7 @@ from ._typing import (
     DynamicsTrialData,
     ExperimentResult,
     FullDynamicsTrialData,
-    ProbData,
+    SimProbData,
     SINDyTrialUpdate,
     _BaseSINDy,
 )
@@ -41,33 +41,9 @@ DType = TypeVar("DType", bound=np.dtype)
 MOD_LOG = getLogger(__name__)
 
 
-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
-
-
 @overload
 def fit_eval(
-    data: tuple[list[ProbData], list[dict[sp.Expr, float]]],
+    data: tuple[list[SimProbData], list[dict[sp.Expr, float]]],
     model: _BaseSINDy,
     simulations: Literal[False],
     display: bool,
@@ -76,7 +52,7 @@ def fit_eval(
 
 @overload
 def fit_eval(
-    data: tuple[list[ProbData], list[dict[sp.Expr, float]]],
+    data: tuple[list[SimProbData], list[dict[sp.Expr, float]]],
     model: _BaseSINDy,
     simulations: Literal[True],
     display: bool,
@@ -84,7 +60,7 @@ def fit_eval(
 
 
 def fit_eval(
-    data: tuple[list[ProbData], list[dict[sp.Expr, float]]],
+    data: tuple[list[SimProbData], list[dict[sp.Expr, float]]],
     model: Any,
     simulations: bool = True,
     display: bool = True,
@@ -93,7 +69,7 @@ def fit_eval(
 
     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.
@@ -101,6 +77,8 @@ def fit_eval(
         display: Whether to generate plots as part of evaluation.
     """
     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
 
@@ -143,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(
@@ -154,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)
 

{sindy_exp-0.2.2 → sindy_exp-0.3.0}/src/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-0.2.2 → sindy_exp-0.3.0}/src/sindy_exp/_typing.py

@@ -1,10 +1,10 @@
 from collections import defaultdict
+from collections.abc import Mapping
 from dataclasses import dataclass
 from typing import (
     Any,
     Callable,
     Literal,
-    NamedTuple,
     Optional,
     Protocol,
     TypedDict,
@@ -31,9 +31,9 @@ TrajectoryType = TypeVar("TrajectoryType", list[np.ndarray], np.ndarray)
 class ExperimentResult[T](TypedDict):
     """Results from a SINDy ODE experiment."""
 
-    metrics: float
+    metrics: Mapping[str, float | None]
     data: T
-    main: float
+    main: object
 
 
 class _BaseSINDy(Protocol):
@@ -71,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
 
 
@@ -147,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-0.2.2 → sindy_exp-0.3.0}/src/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
@@ -202,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(
@@ -284,7 +284,7 @@ def unionize_coeff_dicts(
 
 
 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
 
@@ -292,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.2 → sindy_exp-0.3.0}/src/sindy_exp.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sindy-exp
-Version: 0.2.2
+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.2.2 → sindy_exp-0.3.0}/src/sindy_exp.egg-info/SOURCES.txt

@@ -7,6 +7,9 @@ pyproject.toml
 setup.cfg
 .github/workflows/main.yaml
 .github/workflows/release.yml
+images/1d.png
+images/coeff.png
+images/composite.png
 src/sindy_exp/__init__.py
 src/sindy_exp/_data.py
 src/sindy_exp/_diffrax_solver.py

{sindy_exp-0.2.2 → sindy_exp-0.3.0}/src/sindy_exp.egg-info/requires.txt

@@ -4,6 +4,7 @@ seaborn
 scipy
 sympy
 dysts
+scikit-learn
 
 [dev]
 ipykernel

sindy_exp-0.2.2/README.md

@@ -1,45 +0,0 @@
-# Dynamics Experiments
-
-A library for constructing dynamics experiments.
-This includes data generation and plotting/evaluation.
-
-## 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`
-
-Generate data
-
-    data = sindy_exp.data.gen_data("lorenz", num_trajectories=5, t_end=10.0, dt=0.01)["data]
-
-Evaluate your SINDy-like model with:
-
-    sindy_exp.odes.fit_eval(model, data)
-
-![Coefficient plots](images/coeff.png)
-
-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
-
-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.
-Each element is a dictionary mapping a term (Sympy expression) to its coefficient.
-
-## 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`.
-
-![3d plot](images/composite.png)
-![1d plot](images/1d.png)