Merge pull request #20 from Jacob-Stevens-Haas/grid-locator

Jacob-Stevens-Haas · Mar 15, 2024 · e236e17 · e236e17
2 parents 7faa74d + d17275e
commit e236e17
Show file tree

Hide file tree

Showing 13 changed files with 1,236 additions and 1,040 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -92,7 +92,13 @@ addopts = '-m "not slow"'
 markers = ["slow"]
 
 [tool.mypy]
-files = ["src/gen_experiments/__init__.py", "src/gen_experiments/utils.py"]
+files = [
+  "src/gen_experiments/__init__.py",
+  "src/gen_experiments/utils.py",
+  "src/gen_experiments/gridsearch/typing.py",
+  "tests/test_all.py",
+  "tests/test_gridsearch.py",
+]
 
 [[tool.mypy.overrides]]
 module="auto_ks.*"

diff --git a/src/gen_experiments/__init__.py b/src/gen_experiments/__init__.py
@@ -5,7 +5,7 @@
 
 import numpy as np
 from numpy.typing import NDArray
-from pysindy import BaseDifferentiation, FiniteDifference, SINDy  # type: ignore
+from pysindy import BaseDifferentiation, FiniteDifference, SINDy
 
 from . import gridsearch, odes, pdes
 from .utils import SINDyTrialData, make_model  # noqa: F401

diff --git a/src/gen_experiments/config.py b/src/gen_experiments/config.py
@@ -1,11 +1,19 @@
+from collections.abc import Iterable
 from typing import TypeVar
 
 import numpy as np
 import pysindy as ps
 
 from gen_experiments.data import _signal_avg_power
+from gen_experiments.gridsearch.typing import (
+    GridLocator,
+    SeriesDef,
+    SeriesList,
+    SkinnySpecs,
+)
 from gen_experiments.plotting import _PlotPrefs
-from gen_experiments.utils import FullSINDyTrialData, NestedDict, SeriesDef, SeriesList
+from gen_experiments.typing import NestedDict
+from gen_experiments.utils import FullSINDyTrialData
 
 T = TypeVar("T")
 U = TypeVar("U")
@@ -41,77 +49,43 @@ def addn(x):
 
 
 plot_prefs = {
-    "test": _PlotPrefs(True, False, ({"sim_params.t_end": 10},)),
+    "test": _PlotPrefs(),
     "test-absrel": _PlotPrefs(
-        True, _convert_abs_rel_noise, ({"sim_params.noise_abs": 1},)
+        True, _convert_abs_rel_noise, GridLocator(..., {("sim_params.noise_abs", (1,))})
     ),
     "test-absrel2": _PlotPrefs(
         True,
         _convert_abs_rel_noise,
-        (
-            {"sim_params.noise_abs": 0.1},
-            {"sim_params.noise_abs": 0.5},
-            {"sim_params.noise_abs": 1},
-            {"sim_params.noise_abs": 2},
-            {"sim_params.noise_abs": 4},
-            {"sim_params.noise_abs": 8},
+        GridLocator(
+            ...,
+            (..., ...),
+            (
+                {"sim_params.noise_abs": 0.1},
+                {"sim_params.noise_abs": 0.5},
+                {"sim_params.noise_abs": 1},
+                {"sim_params.noise_abs": 2},
+                {"sim_params.noise_abs": 4},
+                {"sim_params.noise_abs": 8},
+            ),
         ),
     ),
-    "test-absrel3": _PlotPrefs(
+    "absrel-newloc": _PlotPrefs(
         True,
         _convert_abs_rel_noise,
-        (
-            {
-                "sim_params.noise_abs": 1,
-                "diff_params.smoother_kws.window_length": 15,
-            },
-            {"sim_params.noise_abs": 1, "diff_params.meas_var": 1},
-            {"sim_params.noise_abs": 1, "diff_params.alpha": 1e-2},
+        GridLocator(
+            ["coeff_mse", "coeff_f1"],
+            (..., (2, 3, 4)),
+            (
+                {"diff_params.kind": "kalman", "diff_params.alpha": None},
+                {
+                    "diff_params.kind": "kalman",
+                    "diff_params.alpha": lambda a: isinstance(a, int),
+                },
+                {"diff_params.kind": "trend_filtered"},
+                {"diff_params.diffcls": "SmoothedFiniteDifference"},
+            ),
         ),
     ),
-    "test-absrel4": _PlotPrefs(
-        True,
-        _convert_abs_rel_noise,
-        (
-            {
-                "sim_params.noise_abs": 1,
-                "diff_params.smoother_kws.window_length": 15,
-            },
-            {"sim_params.noise_abs": 1, "diff_params.meas_var": 1},
-            {"sim_params.noise_abs": 1, "diff_params.alpha": 1e0},
-            {
-                "sim_params.noise_abs": 2,
-                "diff_params.smoother_kws.window_length": 15,
-            },
-            {"sim_params.noise_abs": 2, "diff_params.meas_var": 4},
-            {"sim_params.noise_abs": 2, "diff_params.alpha": 1e-1},
-        ),
-    ),
-    "test-absrel5": _PlotPrefs(
-        True,
-        _convert_abs_rel_noise,
-        (
-            {
-                "sim_params.noise_abs": 1,
-                "diff_params.diffcls": "SmoothedFiniteDifference",
-            },
-            {"sim_params.noise_abs": 1, "diff_params.kind": "kalman"},
-            {"sim_params.noise_abs": 1, "diff_params.kind": "trend_filtered"},
-            {
-                "sim_params.noise_abs": 2,
-                "diff_params.diffcls": "SmoothedFiniteDifference",
-            },
-            {"sim_params.noise_abs": 2, "diff_params.kind": "kalman"},
-            {"sim_params.noise_abs": 2, "diff_params.kind": "trend_filtered"},
-            {
-                "sim_params.noise_abs": 4,
-                "diff_params.diffcls": "SmoothedFiniteDifference",
-            },
-            {"sim_params.noise_abs": 4, "diff_params.kind": "kalman"},
-            {"sim_params.noise_abs": 4, "diff_params.kind": "trend_filtered"},
-        ),
-        {(0, 2), (3, 2), (0, 3), (3, 3), (0, 4), (3, 4)},
-    ),
 }
 sim_params = {
     "test": ND({"n_trajectories": 2}),
@@ -301,7 +275,7 @@ def addn(x):
     "duration-absnoise": ["sim_params.t_end", "sim_params.noise_abs"],
     "rel_noise": ["sim_params.t_end", "sim_params.noise_rel"],
 }
-grid_vals = {
+grid_vals: dict[str, list[Iterable]] = {
     "test": [[5, 10, 15, 20]],
     "abs_noise": [[0.1, 0.5, 1, 2, 4, 8]],
     "abs_noise-kalman": [[0.1, 0.5, 1, 2, 4, 8], [0.1, 0.5, 1, 2, 4, 8]],
@@ -319,7 +293,7 @@ def addn(x):
     "lorenzk": ["plot", "plot", "max"],
     "plot2": ["plot", "plot"],
 }
-diff_series = {
+diff_series: dict[str, SeriesDef] = {
     "kalman1": SeriesDef(
         "Kalman",
         diff_params["kalman"],
@@ -375,7 +349,7 @@ def addn(x):
         [[5, 8, 12, 15]],
     ),
 }
-series_params = {
+series_params: dict[str, SeriesList] = {
     "test": SeriesList(
         "diff_params",
         "Differentiation Method",
@@ -434,7 +408,7 @@ def addn(x):
 }
 
 
-skinny_specs = {
+skinny_specs: dict[str, SkinnySpecs] = {
     "exp3": (
         ("sim_params.noise_abs", "diff_params.meas_var"),
         ((identity,), (identity,)),

diff --git a/src/gen_experiments/data.py b/src/gen_experiments/data.py
@@ -7,7 +7,8 @@
 import numpy as np
 import scipy
 
-from gen_experiments.utils import Float1D, Float2D, GridsearchResultDetails
+from gen_experiments.gridsearch.typing import GridsearchResultDetails
+from gen_experiments.utils import Float1D, Float2D
 
 INTEGRATOR_KEYWORDS = {"rtol": 1e-12, "method": "LSODA", "atol": 1e-12}
 TRIALS_FOLDER = Path(__file__).parent.absolute() / "trials"
@@ -56,7 +57,7 @@ def gen_data(
         noise_abs = 0.1
     rng = np.random.default_rng(seed)
     if x0_center is None:
-        x0_center = np.zeros((n_coord))
+        x0_center = np.zeros((n_coord), dtype=np.float_)
     t_train = np.arange(0, t_end, dt, dtype=np.float_)
     t_train_span = (t_train[0], t_train[-1])
     if nonnegative:

diff --git a/src/gen_experiments/debug.py b/src/gen_experiments/debug.py