feat(API): export make_model()

Jacob-Stevens-Haas · Feb 19, 2024 · fa78b92 · fa78b92
1 parent 2557f81
commit fa78b92
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/src/gen_experiments/__init__.py b/src/gen_experiments/__init__.py
@@ -8,7 +8,7 @@
 from pysindy import BaseDifferentiation, FiniteDifference, SINDy  # type: ignore
 
 from . import gridsearch, odes, pdes
-from .utils import SINDyTrialData
+from .utils import SINDyTrialData, make_model  # noqa: F401
 
 this_module = importlib.import_module(__name__)
 BORING_ARRAY = np.ones((2, 2), dtype=float)

diff --git a/src/gen_experiments/utils.py b/src/gen_experiments/utils.py
@@ -122,6 +122,8 @@ def opt_lookup(kind):
         return ps.SR3
     elif normalized_kind == "miosr":
         return ps.MIOSR
+    elif normalized_kind == "trap":
+        return ps.TrappingSR3
     elif normalized_kind == "ensemble":
         return ps.EnsembleOptimizer
     else:
@@ -207,7 +209,7 @@ def unionize_coeff_matrices(
     return true_coeff_mat, new_est_coeff, model_features
 
 
-def _make_model(
+def make_model(
     input_features: list[str],
     dt: float,
     diff_params: dict,