diff --git a/examples/8_trapping_sindy_examples/example.ipynb b/examples/8_trapping_sindy_examples/example.ipynb
index bbe365d3..58af3da3 100644
--- a/examples/8_trapping_sindy_examples/example.ipynb
+++ b/examples/8_trapping_sindy_examples/example.ipynb
@@ -204,7 +204,7 @@
     ").y.T\n",
     "\n",
     "# define hyperparameters\n",
-    "threshold = 0.0\n",
+    "reg_weight_lam = 0.0\n",
     "eta = 1e5\n",
     "max_iter = 5000\n",
     "\n",
@@ -213,7 +213,7 @@
     "sindy_opt = ps.TrappingSR3(\n",
     "    _n_tgts=3,\n",
     "    _include_bias=True,\n",
-    "    threshold=threshold,\n",
+    "    reg_weight_lam=reg_weight_lam,\n",
     "    eta=eta,\n",
     "    max_iter=max_iter,\n",
     "    gamma=-1,\n",
@@ -397,11 +397,11 @@
     "# define hyperparameters\n",
     "eta = 1.0e8\n",
     "\n",
-    "# run trapping SINDy, reusing previous threshold, max_iter and constraints\n",
+    "# run trapping SINDy, reusing previous reg_weight_lam, max_iter and constraints\n",
     "sindy_opt = ps.TrappingSR3(\n",
     "    _n_tgts=3,\n",
     "    _include_bias=True,\n",
-    "    threshold=threshold,\n",
+    "    reg_weight_lam=reg_weight_lam,\n",
     "    eta=eta,\n",
     "    max_iter=max_iter,\n",
     "    verbose=True\n",
@@ -649,7 +649,7 @@
     "x_test = solve_ivp(lorenz, t_span, x0, t_eval=t, **integrator_keywords).y.T\n",
     "\n",
     "# define hyperparameters\n",
-    "threshold = 0\n",
+    "reg_weight_lam = 0\n",
     "max_iter = 5000\n",
     "eta = 1.0e3\n",
     "\n",
@@ -659,7 +659,7 @@
     "sindy_opt = ps.TrappingSR3(\n",
     "    _n_tgts=3,\n",
     "    _include_bias=True,\n",
-    "    threshold=threshold,\n",
+    "    reg_weight_lam=reg_weight_lam,\n",
     "    eta=eta,\n",
     "    alpha_m=alpha_m,\n",
     "    max_iter=max_iter,\n",
@@ -885,7 +885,7 @@
     "x_test = solve_ivp(mhd, t_span, x0, t_eval=t, **integrator_keywords).y.T\n",
     "\n",
     "# define hyperparameters\n",
-    "threshold = 0.0\n",
+    "reg_weight_lam = 0.0\n",
     "max_iter = 1000\n",
     "eta = 1.0e10\n",
     "alpha_m = 9.0e-1 * eta\n",
@@ -894,7 +894,7 @@
     "sindy_opt = ps.TrappingSR3(\n",
     "    _n_tgts=6,\n",
     "    _include_bias=True,\n",
-    "    threshold=threshold,\n",
+    "    reg_weight_lam=reg_weight_lam,\n",
     "    eta=eta,\n",
     "    max_iter=max_iter,\n",
     "    verbose=True,\n",
@@ -1325,16 +1325,16 @@
     "max_iter = 10000\n",
     "eta = 1.0\n",
     "\n",
-    "# don't need a threshold if eta is sufficiently small\n",
+    "# don't need a reg_weight_lam if eta is sufficiently small\n",
     "# which is good news because CVXPY is much slower\n",
-    "threshold = 0\n",
+    "reg_weight_lam = 0\n",
     "alpha_m = 1e-1 * eta\n",
     "\n",
     "# run trapping SINDy\n",
     "sindy_opt = ps.TrappingSR3(\n",
     "    _n_tgts=5,\n",
     "    _include_bias=False,\n",
-    "    threshold=threshold,\n",
+    "    reg_weight_lam=reg_weight_lam,\n",
     "    eta=eta,\n",
     "    alpha_m=alpha_m,\n",
     "    max_iter=max_iter,\n",
diff --git a/examples/8_trapping_sindy_examples/example.py b/examples/8_trapping_sindy_examples/example.py
index 3df101fc..c44e70c7 100644
--- a/examples/8_trapping_sindy_examples/example.py
+++ b/examples/8_trapping_sindy_examples/example.py
@@ -43,7 +43,7 @@
     sindy_library_no_bias,
     make_fits,
     make_lissajou,
-    check_stability,
+    check_local_stability,
     trapping_region,
     make_progress_plots,
     galerkin_model,
@@ -97,7 +97,7 @@
 ).y.T
 
 # define hyperparameters
-threshold = 0.0
+reg_weight_lam = 0.0
 eta = 1e5
 max_iter = 5000
 
@@ -106,7 +106,7 @@
 sindy_opt = ps.TrappingSR3(
     _n_tgts=3,
     _include_bias=True,
-    threshold=threshold,
+    reg_weight_lam=reg_weight_lam,
     eta=eta,
     max_iter=max_iter,
     gamma=-1,
@@ -133,7 +133,7 @@
 print("Frobenius Error = ", E_pred)
 mean_val = np.mean(x_test_pred, axis=0)
 mean_val = np.sqrt(np.sum(mean_val**2))
-check_stability(r, Xi, sindy_opt, mean_val)
+check_local_stability(Xi, sindy_opt, mean_val)
 
 # compute relative Frobenius error in the model coefficients
 terms = sindy_library.get_feature_names()
@@ -154,15 +154,6 @@
     ) / np.dot(xdot_test[i, :], xdot_test[i, :])
 print("Time-averaged derivative error = ", np.nanmean(deriv_error))
 
-# %% [markdown]
-# Awesome! The trapping algorithm gets exactly the right model and produces a negative definite matrix,
-# $$\mathbf{A}^S = \begin{bmatrix}
-#     -1.32 & 0 & 0 \\
-#     0 & -1.31 & 0 \\
-#     0 & 0 & -1
-#     \end{bmatrix},$$
-# i.e. it identifies $\epsilon \approx 1.3$ from above. Note that with different algorithm hyperparameters it will produce different $\epsilon$, since the algorithm only cares that the matrix is negative definite (i.e. only cares about the largest eigenvalue), not the precise value of $\epsilon$. Moreover, these eigenvalues can change as the algorithm converges further.
-
 # %% [markdown]
 #
 #
@@ -235,13 +226,14 @@
 # define hyperparameters
 eta = 1.0e8
 
-# run trapping SINDy, reusing previous threshold, max_iter and constraints
+# run trapping SINDy, reusing previous reg_weight_lam, max_iter and constraints
 sindy_opt = ps.TrappingSR3(
     _n_tgts=3,
     _include_bias=True,
-    threshold=threshold,
+    reg_weight_lam=reg_weight_lam,
     eta=eta,
     max_iter=max_iter,
+    verbose=True,
 )
 model = ps.SINDy(
     optimizer=sindy_opt,
@@ -270,7 +262,7 @@
 print("Frobenius error = ", E_pred)
 mean_val = np.mean(x_test_pred, axis=0)
 mean_val = np.sqrt(np.sum(mean_val**2))
-check_stability(r, Xi, sindy_opt, mean_val)
+check_local_stability(Xi, sindy_opt, mean_val)
 
 # compute relative Frobenius error in the model coefficients
 terms = sindy_library.get_feature_names()
@@ -371,7 +363,7 @@
 x_test = solve_ivp(lorenz, t_span, x0, t_eval=t, **integrator_keywords).y.T
 
 # define hyperparameters
-threshold = 0
+reg_weight_lam = 0
 max_iter = 5000
 eta = 1.0e3
 
@@ -381,7 +373,7 @@
 sindy_opt = ps.TrappingSR3(
     _n_tgts=3,
     _include_bias=True,
-    threshold=threshold,
+    reg_weight_lam=reg_weight_lam,
     eta=eta,
     alpha_m=alpha_m,
     max_iter=max_iter,
@@ -416,7 +408,7 @@
 
 mean_val = np.mean(x_test_pred, axis=0)
 mean_val = np.sqrt(np.sum(mean_val**2))
-check_stability(r, Xi, sindy_opt, mean_val)
+check_local_stability(Xi, sindy_opt, mean_val)
 E_pred = np.linalg.norm(x_test - x_test_pred) / np.linalg.norm(x_test)
 print("Frobenius error = ", E_pred)
 
@@ -501,9 +493,9 @@
 nu = 0.0  # viscosity
 mu = 0.0  # resistivity
 
-# define training and testing data
+# define training and testing data (low resolution)
 dt = 0.02
-T = 50
+T = 40
 t = np.arange(0, T + dt, dt)
 t_span = (t[0], t[-1])
 x0 = rng.random((6,)) - 0.5
@@ -512,16 +504,16 @@
 x_test = solve_ivp(mhd, t_span, x0, t_eval=t, **integrator_keywords).y.T
 
 # define hyperparameters
-threshold = 0.0
+reg_weight_lam = 0.0
 max_iter = 1000
 eta = 1.0e10
-alpha_m = 5.0e-1 * eta
+alpha_m = 9.0e-1 * eta
 
 
 sindy_opt = ps.TrappingSR3(
     _n_tgts=6,
     _include_bias=True,
-    threshold=threshold,
+    reg_weight_lam=reg_weight_lam,
     eta=eta,
     max_iter=max_iter,
     verbose=True,
@@ -551,7 +543,7 @@
 make_lissajou(r, x_train, x_test, x_train_pred, x_test_pred, "mhd")
 mean_val = np.mean(x_test_pred, axis=0)
 mean_val = np.sqrt(np.sum(mean_val**2))
-check_stability(r, Xi, sindy_opt, mean_val)
+check_local_stability(Xi, sindy_opt, mean_val)
 E_pred = np.linalg.norm(x_test - x_test_pred) / np.linalg.norm(x_test)
 print(E_pred)
 
@@ -788,20 +780,19 @@
 x_test = a
 
 # define hyperparameters
-max_iter = 5000
-eta = 1.0e2
+max_iter = 10000
+eta = 1.0
 
-# don't need a threshold if eta is sufficiently small
+# don't need a reg_weight_lam if eta is sufficiently small
 # which is good news because CVXPY is much slower
-threshold = 0
-alpha_m = 9e-1 * eta
-
+reg_weight_lam = 0
+alpha_m = 1e-1 * eta
 
 # run trapping SINDy
 sindy_opt = ps.TrappingSR3(
     _n_tgts=5,
     _include_bias=False,
-    threshold=threshold,
+    reg_weight_lam=reg_weight_lam,
     eta=eta,
     alpha_m=alpha_m,
     max_iter=max_iter,
@@ -822,7 +813,7 @@
 Q = np.tensordot(PQ_tensor, Xi, axes=([4, 3], [0, 1]))
 Q_sum = np.max(np.abs((Q + np.transpose(Q, [1, 2, 0]) + np.transpose(Q, [2, 0, 1]))))
 print("Max deviation from the constraints = ", Q_sum)
-if check_stability(r, Xi, sindy_opt, 1):
+if check_local_stability(Xi, sindy_opt, 1):
     x_train_pred = model.simulate(x_train[0, :], t, integrator_kws=integrator_keywords)
     x_test_pred = model.simulate(a0, t, integrator_kws=integrator_keywords)
     make_progress_plots(r, sindy_opt)
@@ -832,7 +823,7 @@
     make_lissajou(r, x_train, x_test, x_train_pred, x_test_pred, "VonKarman")
     mean_val = np.mean(x_test_pred, axis=0)
     mean_val = np.sqrt(np.sum(mean_val**2))
-    check_stability(r, Xi, sindy_opt, mean_val)
+    check_local_stability(Xi, sindy_opt, mean_val)
     A_guess = sindy_opt.A_history_[-1]
     m_guess = sindy_opt.m_history_[-1]
     E_pred = np.linalg.norm(x_test - x_test_pred) / np.linalg.norm(x_test)
diff --git a/examples/8_trapping_sindy_examples/example_dysts.ipynb b/examples/8_trapping_sindy_examples/example_dysts.ipynb
index 48a7f533..d196550b 100644
--- a/examples/8_trapping_sindy_examples/example_dysts.ipynb
+++ b/examples/8_trapping_sindy_examples/example_dysts.ipynb
@@ -517,7 +517,7 @@
     "from scipy.optimize import dual_annealing as anneal_algo\n",
     "\n",
     "# define hyperparameters\n",
-    "threshold = 0\n",
+    "reg_weight_lam = 0\n",
     "max_iter = 5000\n",
     "eta = 1.0e3\n",
     "alpha_m = 4e-2 * eta  # default is 1e-2 * eta so this speeds up the code here\n",
@@ -541,7 +541,7 @@
     "        method=\"global\",\n",
     "        _n_tgts=r,\n",
     "        _include_bias=True,\n",
-    "        threshold=threshold,\n",
+    "        reg_weight_lam=reg_weight_lam,\n",
     "        eta=eta,\n",
     "        alpha_m=alpha_m,\n",
     "        max_iter=max_iter,\n",
@@ -785,7 +785,7 @@
     }
    ],
    "source": [
-    "threshold = 0\n",
+    "reg_weight_lam = 0\n",
     "max_iter = 5000\n",
     "eta = 1.0e2\n",
     "alpha_m = 0.1 * eta  # default is 1e-2 * eta so this speeds up the code here\n",
@@ -809,7 +809,7 @@
     "        method=\"local\",\n",
     "        _n_tgts=r,\n",
     "        _include_bias=True,\n",
-    "        threshold=threshold,\n",
+    "        reg_weight_lam=reg_weight_lam,\n",
     "        eta=eta,\n",
     "        alpha_m=alpha_m,\n",
     "        max_iter=max_iter,\n",
@@ -989,7 +989,7 @@
     "x_train = a_dns[:, :r]\n",
     "t_train = t_dns.copy()\n",
     "\n",
-    "threshold = 0.0\n",
+    "reg_weight_lam = 0.0\n",
     "eta = 1.0e10\n",
     "alpha_m = 1e-3 * eta  # default is 1e-2 * eta so this speeds up the code here\n",
     "beta = 1e-4\n",
@@ -998,7 +998,7 @@
     "    method=\"local\",\n",
     "    _n_tgts=r,\n",
     "    _include_bias=True,\n",
-    "    threshold=threshold,\n",
+    "    reg_weight_lam=reg_weight_lam,\n",
     "    eta=eta,\n",
     "    alpha_m=alpha_m,\n",
     "    max_iter=max_iter,\n",
@@ -1019,7 +1019,7 @@
     "\n",
     "# Fit a baseline model -- this is almost always an unstable model!\n",
     "model_baseline = ps.SINDy(\n",
-    "    optimizer=ps.STLSQ(threshold=0.0),\n",
+    "    optimizer=ps.STLSQ(reg_weight_lam=0.0),\n",
     "    feature_library=ps.PolynomialLibrary(),\n",
     ")\n",
     "model_baseline.fit(x_train, t=t_train)"
diff --git a/examples/8_trapping_sindy_examples/example_dysts.py b/examples/8_trapping_sindy_examples/example_dysts.py
index 908c3a99..ef6570a1 100644
--- a/examples/8_trapping_sindy_examples/example_dysts.py
+++ b/examples/8_trapping_sindy_examples/example_dysts.py
@@ -204,7 +204,7 @@
 from scipy.optimize import dual_annealing as anneal_algo
 
 # define hyperparameters
-threshold = 0
+reg_weight_lam = 0
 max_iter = 5000
 eta = 1.0e3
 alpha_m = 4e-2 * eta  # default is 1e-2 * eta so this speeds up the code here
@@ -228,7 +228,7 @@
         method="global",
         _n_tgts=r,
         _include_bias=True,
-        threshold=threshold,
+        reg_weight_lam=reg_weight_lam,
         eta=eta,
         alpha_m=alpha_m,
         max_iter=max_iter,
@@ -309,7 +309,7 @@
 
 # %%
 # define hyperparameters
-threshold = 0
+reg_weight_lam = 0
 max_iter = 5000
 eta = 1.0e7
 alpha_m = 0.1 * eta  # default is 1e-2 * eta so this speeds up the code here
@@ -333,7 +333,7 @@
         method="local",
         _n_tgts=r,
         _include_bias=True,
-        threshold=threshold,
+        reg_weight_lam=reg_weight_lam,
         eta=eta,
         alpha_m=alpha_m,
         max_iter=max_iter,
@@ -457,7 +457,7 @@ def rhs(t, x):
 x_train = a_dns[:, :r]
 t_train = t_dns.copy()
 
-threshold = 0.0
+reg_weight_lam = 0.0
 eta = 1.0e10
 alpha_m = 1e-4 * eta  # default is 1e-2 * eta so this speeds up the code here
 beta = 1e-5
@@ -466,7 +466,7 @@ def rhs(t, x):
     method="local",
     _n_tgts=r,
     _include_bias=True,
-    threshold=threshold,
+    reg_weight_lam=reg_weight_lam,
     eta=eta,
     alpha_m=alpha_m,
     max_iter=max_iter,
@@ -487,7 +487,7 @@ def rhs(t, x):
 
 # Fit a baseline model -- this is almost always an unstable model!
 model_baseline = ps.SINDy(
-    optimizer=ps.STLSQ(threshold=0.0),
+    optimizer=ps.STLSQ(reg_weight_lam=0.0),
     feature_library=ps.PolynomialLibrary(),
 )
 model_baseline.fit(x_train, t=t_train)
diff --git a/examples/8_trapping_sindy_examples/trapping_extended.ipynb b/examples/8_trapping_sindy_examples/trapping_extended.ipynb
index 7f302b49..ff4a7991 100644
--- a/examples/8_trapping_sindy_examples/trapping_extended.ipynb
+++ b/examples/8_trapping_sindy_examples/trapping_extended.ipynb
@@ -159,7 +159,7 @@
     "x_test = solve_ivp(lorenz, t_span, x0, t_eval=t, **integrator_keywords).y.T\n",
     "\n",
     "# define hyperparameters\n",
-    "threshold = 0\n",
+    "reg_weight_lam = 0\n",
     "max_iter = 5000\n",
     "eta = 1.0e3\n",
     "alpha_m = 8e-1 * eta\n",
@@ -169,7 +169,7 @@
     "    method=\"global\",\n",
     "    _n_tgts=3,\n",
     "    _include_bias=True,\n",
-    "    threshold=threshold,\n",
+    "    reg_weight_lam=reg_weight_lam,\n",
     "    eta=eta,\n",
     "    max_iter=max_iter,\n",
     "    gamma=-1,\n",
@@ -393,14 +393,14 @@
     "eta = 1.0e2\n",
     "alpha = 1e-15\n",
     "beta = 1e20\n",
-    "threshold = 0\n",
+    "reg_weight_lam = 0\n",
     "\n",
     "# run trapping SINDy... no more constraints!\n",
     "sindy_opt = ps.TrappingSR3(\n",
     "    method=\"local\",\n",
     "    _n_tgts=3,\n",
     "    _include_bias=True,\n",
-    "    threshold=threshold,\n",
+    "    reg_weight_lam=reg_weight_lam,\n",
     "    eta=eta,\n",
     "    max_iter=max_iter,\n",
     "    gamma=-1,\n",
@@ -477,7 +477,7 @@
     "eta = 1.0e3\n",
     "alpha = 1e20\n",
     "beta = 1e-10\n",
-    "threshold = 0\n",
+    "reg_weight_lam = 0\n",
     "alpha_m = 0.9 * eta\n",
     "\n",
     "# run trapping SINDy... no more constraints!\n",
@@ -485,7 +485,7 @@
     "    method=\"local\",\n",
     "    _n_tgts=3,\n",
     "    _include_bias=True,\n",
-    "    threshold=threshold,\n",
+    "    reg_weight_lam=reg_weight_lam,\n",
     "    eta=eta,\n",
     "    alpha_m=alpha_m,\n",
     "    max_iter=max_iter,\n",
@@ -662,7 +662,7 @@
     "eta = 1.0e5\n",
     "alpha = 1e20\n",
     "beta = 1e-10\n",
-    "threshold = 5\n",
+    "reg_weight_lam = 5\n",
     "alpha_m = 0.9 * eta\n",
     "\n",
     "# run trapping SINDy... no more constraints!\n",
@@ -670,7 +670,7 @@
     "    method=\"local\",\n",
     "    _n_tgts=3,\n",
     "    _include_bias=True,\n",
-    "    threshold=threshold,\n",
+    "    reg_weight_lam=reg_weight_lam,\n",
     "    eta=eta,\n",
     "    alpha_m=alpha_m,\n",
     "    max_iter=max_iter,\n",
@@ -790,7 +790,7 @@
     "eta = 1.0e2\n",
     "alpha = 1e20\n",
     "beta = 1e-14\n",
-    "threshold = 0\n",
+    "reg_weight_lam = 0\n",
     "alpha_m = 0.1 * eta\n",
     "\n",
     "# run trapping SINDy... no more constraints!\n",
@@ -798,7 +798,7 @@
     "    method=\"local\",\n",
     "    _n_tgts=3,\n",
     "    _include_bias=True,\n",
-    "    threshold=threshold,\n",
+    "    reg_weight_lam=reg_weight_lam,\n",
     "    eta=eta,\n",
     "    alpha_m=alpha_m,\n",
     "    max_iter=max_iter,\n",
diff --git a/examples/8_trapping_sindy_examples/trapping_extended.py b/examples/8_trapping_sindy_examples/trapping_extended.py
index 7abf64a8..435bcd12 100644
--- a/examples/8_trapping_sindy_examples/trapping_extended.py
+++ b/examples/8_trapping_sindy_examples/trapping_extended.py
@@ -1,14 +1,14 @@
 # %% [markdown]
-# # Extended Trapping SINDy
-# By Mai Peng and Alan Kaptanoglu
+#  # Extended Trapping SINDy
+#  By Mai Peng, Alan Kaptanoglu and Jake Stevens-Haas
 #
-# A very common issue is that models identified by system identification methods typically have no guarantees that the models are numerically or physically stable. This can be addressed with heuristic, data-driven, or analytic closure models, but we have recently directly promoted globally stable models into the system identification itself (see the Example 8 Jupyter notebook). This is really nice but there are three potential caveats, (1) the regression is nonconvex and there a number of hyperparameters, so this method can be difficult to learn, and (2) in order to promote global stability, one needs an analytic result from stability theory, and the one we use applies only for quadratically nonlinear dynamics (typically fluid and plasma flows) with energy-preserving, quadratic, nonlinearities. Moreover, we have good reason to believe that (3) generic quadratically nonlinear models will always be globally *unbounded*, so for these situations we can also promote local Lyapunov stability of the origin using some variations of the original Trapping SINDy algorithm. That is the goal of this notebook -- to illustrate how various forms of global and local stability can be promoted explicitly in the SINDy method to obtain stable data-driven models.
+#  A very common issue is that models identified by system identification methods typically have no guarantees that the models are numerically or physically stable. This can be addressed with heuristic, data-driven, or analytic closure models, but we have recently directly promoted globally stable models into the system identification itself (see the Example 8 Jupyter notebook). This is really nice but there are three potential caveats, (1) the regression is nonconvex and there a number of hyperparameters, so this method can be difficult to learn, and (2) in order to promote global stability, one needs an analytic result from stability theory, and the one we use applies only for quadratically nonlinear dynamics (typically fluid and plasma flows) with energy-preserving, quadratic, nonlinearities. Moreover, we have good reason to believe that (3) generic quadratically nonlinear models will always be globally *unbounded*, so for these situations we can also promote local Lyapunov stability of the origin using some variations of the original Trapping SINDy algorithm. That is the goal of this notebook -- to illustrate how various forms of global and local stability can be promoted explicitly in the SINDy method to obtain stable data-driven models.
 #
-# For the following, we will consider dynamical models of the form
-# $$\dot{x}_i = C_i +  L_{ij}x_j + Q_{ijk}x_ix_j.$$
-# For global stability promotion, we will require that the totally symmetric part of the quadratic coefficients vanishes (without loss of generality, $Q_{ijk}$ is symmetric in the last two indices):
-# $$ Q_{ijk} + Q_{jik} + Q_{kij} = 0.$$
-# This equation can be implemented as a hard or soft constraint in the optimization. For dynamical models that do not satisfy this condition, we can still promote locally stable models that are stable even at very large distances of the origin. The following examples show different ways to relax this hard constraint.
+#  For the following, we will consider dynamical models of the form
+#  $$\dot{x}_i = C_i +  L_{ij}x_j + Q_{ijk}x_ix_j.$$
+#  For global stability promotion, we will require that the totally symmetric part of the quadratic coefficients vanishes (without loss of generality, $Q_{ijk}$ is symmetric in the last two indices):
+#  $$ Q_{ijk} + Q_{jik} + Q_{kij} = 0.$$
+#  This equation can be implemented as a hard or soft constraint in the optimization. For dynamical models that do not satisfy this condition, we can still promote locally stable models that are stable even at very large distances of the origin. The following examples show different ways to relax this hard constraint.
 # %%
 import warnings
 
@@ -29,7 +29,6 @@
     make_fits,
     obj_function,
     check_local_stability,
-    check_stability,
     make_trap_progress_plots,
 )
 
@@ -37,64 +36,64 @@
 
 # %% [markdown]
 # # Lorenz model
-# The Lorenz system originates from a simple fluid model of atmospheric dynamics from Lorenz et al. (1963).
-# This system is likely the most famous example of chaotic, nonlinear behavior despite the somewhat innocuous system of equations,
+#  The Lorenz system originates from a simple fluid model of atmospheric dynamics from Lorenz et al. (1963).
+#  This system is likely the most famous example of chaotic, nonlinear behavior despite the somewhat innocuous system of equations,
 #
-# $$
-# \begin{align}
-#     \frac{d}{dt}\begin{bmatrix}
-#     x \\
-#     y \\
-#     z \\
-#     \end{bmatrix} &= \begin{bmatrix}
-#     -\sigma & \sigma & 0 \\
-#     \rho & -1 & 0 \\
-#     0 & 0 & -\beta
-#     \end{bmatrix}
-#     \begin{bmatrix}
-#     x \\
-#     y \\
-#     z
-#     \end{bmatrix}
-#     +
-#     \begin{bmatrix}
-#     0 \\
-#     -xz \\
-#     xy
-#     \end{bmatrix}, \qquad
-#     \mathbf{A}^S = \begin{bmatrix}
-#     -\sigma & \frac{1}{2}(\rho+\sigma - m_3) & \frac{1}{2}m_2 \\
-#     \frac{1}{2}(\rho+\sigma - m_3) & -1 & 0 \\
-#     \frac{1}{2}m_2 & 0 & -\beta
-#     \end{bmatrix}.
-# \end{align}
-# $$
+#  $$
+#  \begin{align}
+#      \frac{d}{dt}\begin{bmatrix}
+#      x \\
+#      y \\
+#      z \\
+#      \end{bmatrix} &= \begin{bmatrix}
+#      -\sigma & \sigma & 0 \\
+#      \rho & -1 & 0 \\
+#      0 & 0 & -\beta
+#      \end{bmatrix}
+#      \begin{bmatrix}
+#      x \\
+#      y \\
+#      z
+#      \end{bmatrix}
+#      +
+#      \begin{bmatrix}
+#      0 \\
+#      -xz \\
+#      xy
+#      \end{bmatrix}, \qquad
+#      \mathbf{A}^S = \begin{bmatrix}
+#      -\sigma & \frac{1}{2}(\rho+\sigma - m_3) & \frac{1}{2}m_2 \\
+#      \frac{1}{2}(\rho+\sigma - m_3) & -1 & 0 \\
+#      \frac{1}{2}m_2 & 0 & -\beta
+#      \end{bmatrix}.
+#  \end{align}
+#  $$
 #
-# For Lorenz's choice of parameters, $\sigma = 10$, $\rho = 28$, $\beta  = 8/3$, this system is known to exhibit a stable attractor. For $\mathbf{m} = [0,m_2,\rho+\sigma]$ ($m_1$ does not contribute to $\mathbf{A}^S$ so we set it to zero),
+#  For Lorenz's choice of parameters, $\sigma = 10$, $\rho = 28$, $\beta  = 8/3$, this system is known to exhibit a stable attractor. For $\mathbf{m} = [0,m_2,\rho+\sigma]$ ($m_1$ does not contribute to $\mathbf{A}^S$ so we set it to zero),
 #
-# $$
-# \begin{align}
-#     \mathbf{A}^S &= \begin{bmatrix}
-#     -\sigma & 0 & \frac{1}{2}m_2 \\
-#     0 & -1 & 0 \\
-#     \frac{1}{2}m_2 & 0 & -\beta
-#     \end{bmatrix}, \qquad
-#     \lambda_1 = -1, \qquad \lambda_{\pm} = -\frac{1}{2}\left[\beta+\sigma \mp \sqrt{m_2^2 + (\beta-\sigma)^2}\right],
-# \end{align}
-# $$
+#  $$
+#  \begin{align}
+#      \mathbf{A}^S &= \begin{bmatrix}
+#      -\sigma & 0 & \frac{1}{2}m_2 \\
+#      0 & -1 & 0 \\
+#      \frac{1}{2}m_2 & 0 & -\beta
+#      \end{bmatrix}, \qquad
+#      \lambda_1 = -1, \qquad \lambda_{\pm} = -\frac{1}{2}\left[\beta+\sigma \mp \sqrt{m_2^2 + (\beta-\sigma)^2}\right],
+#  \end{align}
+#  $$
 #
-# so that if $\lambda_{\pm} < 0$, then $-2\sqrt{\sigma\beta} < m_2 < 2\sqrt{\sigma\beta}$.
-# Our algorithm can successfully identify the optimal $\mathbf{m}$, and can be used to identify the inequality bounds on $m_2$ for stability.
+#  so that if $\lambda_{\pm} < 0$, then $-2\sqrt{\sigma\beta} < m_2 < 2\sqrt{\sigma\beta}$.
+#  Our algorithm can successfully identify the optimal $\mathbf{m}$, and can be used to identify the inequality bounds on $m_2$ for stability.
 
 # %% [markdown]
 # ### Check global stability of the Lorenz model
-# The skew-symmetric models below are globally stable *if and only if* there exists a vector $\mathbf{m}$ such that following matrix is negative definite:
-# $$A^S_{ij} = L^S_{ij} + (Q_{ijk} + Q_{jik})m_k.$$
-# Note that if the quadratic tensor has zero totally symmetric part, this is equal to
-# $$A^S_{ij} = L^S_{ij} - Q_{kij}m_k.$$
-# A negative definite $\mathbf{A}^S$ turns out to also be necessary for models that do not quite satisfy the constraint on $Q_{jik}$, but in this case is not sufficient for global boundedness.
+#  The skew-symmetric models below are globally stable *if and only if* there exists a vector $\mathbf{m}$ such that following matrix is negative definite:
+#  $$A^S_{ij} = L^S_{ij} + (Q_{ijk} + Q_{jik})m_k.$$
+#  Note that if the quadratic tensor has zero totally symmetric part, this is equal to
+#  $$A^S_{ij} = L^S_{ij} - Q_{kij}m_k.$$
+#  A negative definite $\mathbf{A}^S$ turns out to also be necessary for models that do not quite satisfy the constraint on $Q_{jik}$, but in this case is not sufficient for global boundedness.
 #
-# A decent-enough algorithm for a nonlinear search for such a $\mathbf{m}$ that makes $A^S_{ij}$ negative definite is simulated annealing, and a simple interface is provided by scipy.
+#  A decent-enough algorithm for a nonlinear search for such a $\mathbf{m}$ that makes $A^S_{ij}$ negative definite is simulated annealing, and a simple interface is provided by scipy.
 
 # %%
 # define parameters
@@ -112,7 +111,7 @@
 x_test = solve_ivp(lorenz, t_span, x0, t_eval=t, **integrator_keywords).y.T
 
 # define hyperparameters
-threshold = 0
+reg_weight_lam = 0
 max_iter = 5000
 eta = 1.0e3
 alpha_m = 8e-1 * eta
@@ -122,7 +121,7 @@
     method="global",
     _n_tgts=3,
     _include_bias=True,
-    threshold=threshold,
+    reg_weight_lam=reg_weight_lam,
     eta=eta,
     max_iter=max_iter,
     gamma=-1,
@@ -142,8 +141,8 @@
 Lenergy = np.tensordot(PL_tensor, Xi, axes=([3, 2], [0, 1]))
 Q = np.tensordot(PQ_tensor, Xi, axes=([4, 3], [0, 1]))
 print(
-    "Maximum deviation of Qijk from having zero totally symmetric part: ",
-    np.max(np.abs((Q + np.transpose(Q, [1, 2, 0]) + np.transpose(Q, [2, 0, 1])))),
+    r"|tilde{H_0}|_F = ",
+    np.sqrt(np.sum((Q + np.transpose(Q, [1, 2, 0]) + np.transpose(Q, [2, 0, 1])) ** 2)),
 )
 
 # %%
@@ -185,7 +184,7 @@
 
 # %% [markdown]
 # ### Use simulated annealing
-# We are going to check if any $\mathbf{m}$ exists such that $\mathbf{A}^S$ is negative definite, using the identified coefficients, to verify again that our model is globally stable.
+#  We are going to check if any $\mathbf{m}$ exists such that $\mathbf{A}^S$ is negative definite, using the identified coefficients, to verify again that our model is globally stable.
 
 # %%
 # Import simulated annealing algorithm from scipy
@@ -223,29 +222,29 @@
 
 # %% [markdown]
 # ### Promoting locally stable models with estimates of the stability radius
-# So far, we have promoted globally stable models with trapping SINDy by enforcing the skew-symmetry structure in the nonlinearities as a hard constraint in the optimization problem:
-# $$\text{argmin}_{\mathbf{\xi},\mathbf m}\|\dot{\mathbf a} - \mathbf \Theta(\mathbf a) \mathbf{\xi}\|^2 + \gamma R(\mathbf \xi)  + \eta \lambda_1(\mathbf A^S)  \quad s.t. \quad Q_{ijk} + Q_{jik} + Q_{kji} = 0.$$
-# This problem is solved with a convex relaxation of the optimization.
+#  So far, we have promoted globally stable models with trapping SINDy by enforcing the skew-symmetry structure in the nonlinearities as a hard constraint in the optimization problem:
+#  $$\text{argmin}_{\mathbf{\xi},\mathbf m}\|\dot{\mathbf a} - \mathbf \Theta(\mathbf a) \mathbf{\xi}\|^2 + \gamma R(\mathbf \xi)  + \eta \lambda_1(\mathbf A^S)  \quad s.t. \quad Q_{ijk} + Q_{jik} + Q_{kji} = 0.$$
+#  This problem is solved with a convex relaxation of the optimization.
 #
-# Below, we relax the hard constraint to a soft constraint and instead solve
-# $$\text{argmin}_{\mathbf{\xi},\mathbf m}\|\dot{\mathbf a} - \mathbf \Theta(\mathbf a) \mathbf{\xi}\|^2 + \gamma R(\mathbf \xi)  + \eta \lambda_1(\mathbf A^S)  \quad s.t. \quad -\epsilon_Q \leq Q_{ijk} + Q_{jik} + Q_{kji} \leq \epsilon_Q.$$
-# This allows us to build locally Lyapunov stable models, and adjust the size of the local stability radius by varying $\epsilon_Q$. A conservative estimate of the local stability is:
-# $$\rho_+ = \frac{3|\lambda_{\text{max}}|}{2r^{\frac{3}{2}}\epsilon_Q} \left( 1 + \sqrt{1 - \frac{4r^{\frac{3}{2}}\epsilon_Q}{3\lambda^2_{\text{max}}(\textbf{A}_S)\|\mathbf{d}\|_2}} \right).$$
-# And the radius of the trapping region is given by:
-# $$\rho_- = \frac{3|\lambda_{\text{max}}|}{2r^{\frac{3}{2}}\epsilon_Q} \left( 1 - \sqrt{1 - \frac{4r^{\frac{3}{2}}\epsilon_Q}{3\lambda^2_{\text{max}}(\textbf{A}_S)\|\mathbf{d}\|_2}} \right).$$
+#  Below, we relax the hard constraint to a soft constraint and instead solve
+#  $$\text{argmin}_{\mathbf{\xi},\mathbf m}\|\dot{\mathbf a} - \mathbf \Theta(\mathbf a) \mathbf{\xi}\|^2 + \gamma R(\mathbf \xi)  + \eta \lambda_1(\mathbf A^S)  \quad s.t. \| Q_{ijk} + Q_{jik} + Q_{kji} \|_F \leq \epsilon_Q,$$
+#  where $\|\cdot\|_F$ denotes the Frobenius norm. This allows us to build locally Lyapunov stable models, and adjust the size of the local stability radius by varying $\epsilon_Q$. A conservative estimate of the local stability is:
+#  $$\rho_+ = \frac{3|\lambda_{\text{max}}|}{2\epsilon_Q} \left( 1 + \sqrt{1 - \frac{4\epsilon_Q}{3\lambda^2_{\text{max}}(\textbf{A}_S)\|\mathbf{d}\|_2}} \right).$$
+#  And the radius of the trapping region is given by:
+#  $$\rho_- = \frac{3|\lambda_{\text{max}}|}{2\epsilon_Q} \left( 1 - \sqrt{1 - \frac{4\epsilon_Q}{3\lambda^2_{\text{max}}(\textbf{A}_S)\|\mathbf{d}\|_2}} \right).$$
 #
-# In other words, there is a region $\rho_- < \|\mathbf{a}(t)\| < \rho_+$ such that the energy $K$ satisfies $K > 0$ and $\dot{K} < 0$, so that any trajectory with initial condition $\|\mathbf{a}_0\| < \rho_+$ will be bounded for all time. This is because it will fall towards the origin until at least it reaches $\rho_-$, and then it stays in the ball of radius $\rho_-$ for all time.
+#  In other words, there is a region $\rho_- < \|\mathbf{a}(t)\| < \rho_+$ such that the energy $K$ satisfies $K > 0$ and $\dot{K} < 0$, so that any trajectory with initial condition $\|\mathbf{a}_0\| < \rho_+$ will be bounded for all time. This is because it will fall towards the origin until at least it reaches $\rho_-$, and then it stays in the ball of radius $\rho_-$ for all time.
 
 # %% [markdown]
 # ### A better way to optimize
-# However, we find empirically that CVXPY struggles to solve the inequality-constrained problem adequately, and find much better performance by incorporating the constraint as a loss term in the objective.
-# Two other loss terms that can be used as alternatives to increase the size of the stability radius while avoiding extra constraints:
-# $$\alpha^{-1}\|Q_{ijk}\|$$
-# and
-# $$\beta^{-1}\|Q_{ijk} + Q_{jki} + Q_{kij}\|.$$
-# We can combine all of these options into the following unconstrained optimization problem:
-# $$argmin_{\mathbf{\xi},\mathbf m}\|\dot{\mathbf a} - \mathbf \Theta(\mathbf a) \mathbf{\xi}\|^2 + \gamma R(\mathbf \xi)  + \eta^{-1} \lambda_1(\mathbf A) + \alpha^{-1}\|Q_{ijk}\| + \beta^{-1}\|Q_{ijk} + Q_{jki} + Q_{kij}\|.$$
-# We now solve this problem for $\alpha \gg \beta$, $\alpha \ll \beta$, and $\alpha \sim \beta \sim 1.$
+#  However, we find empirically that CVXPY struggles to solve the inequality-constrained problem adequately, and find much better performance by incorporating the constraint as a loss term in the objective.
+#  Two other loss terms that can be used as alternatives to increase the size of the stability radius while avoiding extra constraints:
+#  $$\alpha^{-1}\|Q_{ijk}\|$$
+#  and
+#  $$\beta^{-1}\|Q_{ijk} + Q_{jki} + Q_{kij}\|.$$
+#  We can combine all of these options into the following unconstrained optimization problem:
+#  $$argmin_{\mathbf{\xi},\mathbf m}\|\dot{\mathbf a} - \mathbf \Theta(\mathbf a) \mathbf{\xi}\|^2 + \gamma R(\mathbf \xi)  + \eta^{-1} \lambda_1(\mathbf A) + \alpha^{-1}\|Q_{ijk}\| + \beta^{-1}\|Q_{ijk} + Q_{jki} + Q_{kij}\|.$$
+#  We now solve this problem for $\alpha \gg \beta$, $\alpha \ll \beta$, and $\alpha \sim \beta \sim 1.$
 
 # %% [markdown]
 # ### First case: $\alpha \gg 1$, $\beta \ll 1$, for which the model should just zero out all the quadratic nonlinear terms
@@ -253,16 +252,16 @@
 # %%
 max_iter = 500
 eta = 1.0e2
-alpha = 1e-20
+alpha = 1e-15
 beta = 1e20
-threshold = 0
+reg_weight_lam = 0
 
 # run trapping SINDy... no more constraints!
 sindy_opt = ps.TrappingSR3(
     method="local",
     _n_tgts=3,
     _include_bias=True,
-    threshold=threshold,
+    reg_weight_lam=reg_weight_lam,
     eta=eta,
     max_iter=max_iter,
     gamma=-1,
@@ -283,21 +282,21 @@
 Qenergy = np.tensordot(PQ_tensor, Xi, axes=([4, 3], [0, 1]))
 mean_val = np.mean(x_test_pred, axis=0)
 mean_val = np.sqrt(np.sum(mean_val**2))
-check_stability(r, Xi, sindy_opt, mean_val)
+Rm, R_ls = check_local_stability(Xi, sindy_opt, mean_val)
 
 # %% [markdown]
 # Indeed, we found that if $\alpha \gg 1$ large enough, the quadratic terms in the model are zeroed, which is bad news both for fitting the model and for applying the trapping theorem since the theorem relies on nontrivial quadratic contributions.
 
 # %% [markdown]
 # ### Second case: $\alpha \ll 1$, $\beta \gg 1$, which should reproduce the energy-preserving nonlinear constraint to high accuracy
-# This is a different strategy for stability -- don't make the model's quadratic nonlinearities weak, but make it so that the totally symmetric part of $Q_{ijk}$ is very small.
+#  This is a different strategy for stability -- don't make the model's quadratic nonlinearities weak, but make it so that the totally symmetric part of $Q_{ijk}$ is very small.
 
 # %%
-max_iter = 2000
+max_iter = 10000
 eta = 1.0e3
 alpha = 1e20
 beta = 1e-10
-threshold = 0
+reg_weight_lam = 0
 alpha_m = 0.9 * eta
 
 # run trapping SINDy... no more constraints!
@@ -305,7 +304,7 @@
     method="local",
     _n_tgts=3,
     _include_bias=True,
-    threshold=threshold,
+    reg_weight_lam=reg_weight_lam,
     eta=eta,
     alpha_m=alpha_m,
     max_iter=max_iter,
@@ -327,21 +326,22 @@
 Qenergy = np.tensordot(PQ_tensor, Xi, axes=([4, 3], [0, 1]))
 mean_val = np.mean(x_test_pred, axis=0)
 mean_val = np.sqrt(np.sum(mean_val**2))
-check_local_stability(Xi, sindy_opt, mean_val)
+R_m, R_ls = check_local_stability(Xi, sindy_opt, mean_val)
 Q = np.tensordot(sindy_opt.PQ_, Xi, axes=([4, 3], [0, 1]))
 
 # %% [markdown]
 # ### Plot how the two stability radii changes as the algorithm iterates
-# As the algorithm iterates, it is biasing the model to have a negative definite $\mathbf{A}^S$ matrix. Once this is true, we can estimate the local Lyapunov stability radius $\rho_+$ and the trapping region radius $\rho_-$.
+#  As the algorithm iterates, it is biasing the model to have a negative definite $\mathbf{A}^S$ matrix. Once this is true, we can estimate the local Lyapunov stability radius $\rho_+$ and the trapping region radius $\rho_-$.
 #
-# #### Note that with the soft constraint we can get the stability radius arbitrarily large here!
+#  #### Note that with the soft constraint we can get the stability radius arbitrarily large here!
 
 # %%
 rhos_minus, rhos_plus = make_trap_progress_plots(r, sindy_opt)
 Q = np.tensordot(sindy_opt.PQ_, Xi, axes=([4, 3], [0, 1]))
 Q_sum = np.max(np.abs((Q + np.transpose(Q, [1, 2, 0]) + np.transpose(Q, [2, 0, 1]))))
 print(
-    "Maximum deviation from having zero totally symmetric part: ", np.max(np.abs(Q_sum))
+    r"|tilde{H_0}|_F = ",
+    np.sqrt(np.sum((Q + np.transpose(Q, [1, 2, 0]) + np.transpose(Q, [2, 0, 1])) ** 2)),
 )
 
 # %%
@@ -374,22 +374,22 @@
 
 # %% [markdown]
 # ### Repeat $\alpha \gg 1$, $\beta \ll 1$ case with $\lambda > 0$
-# I find that solver will fail if eps_solver parameter is made too small (error tolerance of the CVXPY solver is very stringent)
+#  I find that solver will fail if eps_solver parameter is made too small (error tolerance of the CVXPY solver is very stringent)
 
 # %%
 max_iter = 100
-eta = 1.0e2
+eta = 1.0e5
 alpha = 1e20
-beta = 1e-20
-threshold = 5
-alpha_m = 9e-1 * eta
+beta = 1e-10
+reg_weight_lam = 5
+alpha_m = 0.9 * eta
 
 # run trapping SINDy... no more constraints!
 sindy_opt = ps.TrappingSR3(
     method="local",
     _n_tgts=3,
     _include_bias=True,
-    threshold=threshold,
+    reg_weight_lam=reg_weight_lam,
     eta=eta,
     alpha_m=alpha_m,
     max_iter=max_iter,
@@ -415,8 +415,8 @@
 check_local_stability(Xi, sindy_opt, mean_val)
 Q = np.tensordot(sindy_opt.PQ_, Xi, axes=([4, 3], [0, 1]))
 print(
-    "Maximum deviation from having zero totally symmetric part: ",
-    np.max(np.abs((Q + np.transpose(Q, [1, 2, 0]) + np.transpose(Q, [2, 0, 1])))),
+    r"|tilde{H_0}|_F = ",
+    np.sqrt(np.sum((Q + np.transpose(Q, [1, 2, 0]) + np.transpose(Q, [2, 0, 1])) ** 2)),
 )
 # make_progress_plots(r, sindy_opt)
 
@@ -431,27 +431,28 @@
 )
 
 # %% [markdown]
-# ### Now we add A LOT of noise to the Lorenz data and see if trapping extended algorithm improves robustness to noise.
+# ### Now we add a lot of noise to the Lorenz data and see if trapping extended algorithm improves robustness to noise.
 
 # %%
+np.random.seed(10)
 lorenz_noise = np.random.normal(
     0, mean_val / 4, x_train.shape
 )  # 25% noise added with zero mean
 x_train_noise = x_train + lorenz_noise
 
 max_iter = 10000
-eta = 1.0e-2
+eta = 1.0e2
 alpha = 1e20
 beta = 1e-14
-threshold = 0
-alpha_m = 0.9 * eta
+reg_weight_lam = 0
+alpha_m = 0.1 * eta
 
 # run trapping SINDy... no more constraints!
 sindy_opt = ps.TrappingSR3(
     method="local",
     _n_tgts=3,
     _include_bias=True,
-    threshold=threshold,
+    reg_weight_lam=reg_weight_lam,
     eta=eta,
     alpha_m=alpha_m,
     max_iter=max_iter,
@@ -476,10 +477,10 @@
 check_local_stability(Xi, sindy_opt, mean_val)
 Q = np.tensordot(sindy_opt.PQ_, Xi, axes=([4, 3], [0, 1]))
 print(
-    "Maximum deviation from having zero totally symmetric part: ",
-    np.max(np.abs((Q + np.transpose(Q, [1, 2, 0]) + np.transpose(Q, [2, 0, 1])))),
+    r"|tilde{H_0}|_F = ",
+    np.sqrt(np.sum((Q + np.transpose(Q, [1, 2, 0]) + np.transpose(Q, [2, 0, 1])) ** 2)),
 )
-make_trap_progress_plots(r, sindy_opt)
+# make_trap_progress_plots(r, sindy_opt)
 
 # Calculate the x_dot and x trajectories for train and test sets
 xdot_test = model.differentiate(x_test, t=t)
@@ -500,12 +501,6 @@
 print("Time-averaged derivative error = ", np.nanmean(deriv_error))
 
 # %%
-# Calculate the x_dot and x trajectories for train and test sets
-xdot_test = model.differentiate(x_test, t=t)
-xdot_test_pred = model.predict(x_test)
-x_train_pred = model.simulate(x_train[0, :], t, integrator_kws=integrator_keywords)
-x_test_pred = model.simulate(x_test[0, :], t, integrator_kws=integrator_keywords)
-
 # plotting and analysis
 make_fits(r, t, xdot_test, xdot_test_pred, x_test, x_test_pred, "lorenz")
 mean_val = np.mean(x_test_pred, axis=0)
@@ -532,7 +527,7 @@
 ax1.plot(x_train_noise[:, 0], x_train_noise[:, 1], x_train_noise[:, 2], "r-")
 ax1.plot(x_train_pred[:, 0], x_train_pred[:, 1], x_train_pred[:, 2], "k-")
 ax1.set(
-    xlabel="$x_0$", ylabel="$x_1$", zlabel="$x_2$", title="model simulation + 50% noise"
+    xlabel="$x_0$", ylabel="$x_1$", zlabel="$x_2$", title="model simulation + 25% noise"
 )
 
 ax2 = fig.add_subplot(122, projection="3d")
@@ -541,4 +536,5 @@
 ax2.set(
     xlabel="$x_0$", ylabel="$x_1$", zlabel="$x_2$", title="true simulation + prediction"
 )
+
 plt.show()
diff --git a/examples/8_trapping_sindy_examples/von_karman_trapping_extended.ipynb b/examples/8_trapping_sindy_examples/von_karman_trapping_extended.ipynb
index 246ee0e9..97735930 100644
--- a/examples/8_trapping_sindy_examples/von_karman_trapping_extended.ipynb
+++ b/examples/8_trapping_sindy_examples/von_karman_trapping_extended.ipynb
@@ -47,7 +47,6 @@
     "    obj_function,\n",
     "    sindy_library_no_bias,\n",
     "    check_local_stability,\n",
-    "    check_local_stability,\n",
     "    make_trap_progress_plots,\n",
     "    make_bar,\n",
     "    nx,\n",
@@ -540,16 +539,16 @@
     "max_iter = 10000\n",
     "eta = 1.0\n",
     "\n",
-    "# don't need a threshold if eta is sufficiently small\n",
+    "# don't need a reg_weight_lam if eta is sufficiently small\n",
     "# which is good news because CVXPY is much slower\n",
-    "threshold = 0\n",
+    "reg_weight_lam = 0\n",
     "alpha_m = 1e-1 * eta\n",
     "\n",
     "# run trapping SINDy\n",
     "sindy_opt = ps.TrappingSR3(\n",
     "    _n_tgts=5,\n",
     "    _include_bias=False,\n",
-    "    threshold=threshold,\n",
+    "    reg_weight_lam=reg_weight_lam,\n",
     "    eta=eta,\n",
     "    alpha_m=alpha_m,\n",
     "    max_iter=max_iter,\n",
@@ -585,7 +584,7 @@
    "source": [
     "max_iter = 2000\n",
     "eta = 1.0\n",
-    "threshold = 0\n",
+    "reg_weight_lam = 0\n",
     "alpha_m = 1e-2 * eta\n",
     "\n",
     "mod_matrix = enstrophy_integrals[1:, 1:]\n",
@@ -593,7 +592,7 @@
     "    method=\"global\",\n",
     "    _n_tgts=5,\n",
     "    _include_bias=False,\n",
-    "    threshold=threshold,\n",
+    "    reg_weight_lam=reg_weight_lam,\n",
     "    eta=eta,\n",
     "    alpha_m=alpha_m,\n",
     "    max_iter=max_iter,\n",
@@ -805,7 +804,7 @@
    "source": [
     "max_iter = 5000\n",
     "eta = 1.0e4\n",
-    "threshold = 0\n",
+    "reg_weight_lam = 0\n",
     "alpha_m = 1e-1 * eta\n",
     "# alpha = 1e10\n",
     "beta = 1e-10\n",
@@ -815,7 +814,7 @@
     "    method=\"local\",\n",
     "    _n_tgts=5,\n",
     "    _include_bias=False,\n",
-    "    threshold=threshold,\n",
+    "    reg_weight_lam=reg_weight_lam,\n",
     "    eta=eta,\n",
     "    gamma=-1,\n",
     "    alpha_m=alpha_m,\n",
@@ -1121,7 +1120,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.9"
+   "version": "3.10.12"
   }
  },
  "nbformat": 4,
diff --git a/examples/8_trapping_sindy_examples/von_karman_trapping_extended.py b/examples/8_trapping_sindy_examples/von_karman_trapping_extended.py
index 3d8a1b2f..305da966 100644
--- a/examples/8_trapping_sindy_examples/von_karman_trapping_extended.py
+++ b/examples/8_trapping_sindy_examples/von_karman_trapping_extended.py
@@ -1,16 +1,16 @@
 # %% [markdown]
 # # Variations of the Trapping theorem -- different Lyapunov functions
-# This example will be used to demonstrate that both the energy and enstrophy can be used as Lyapunov functions for promoting local and global stability with the trapping theorems used for quadratically nonlinear systems of ODEs.
+#  This example will be used to demonstrate that both the energy and enstrophy can be used as Lyapunov functions for promoting local and global stability with the trapping theorems used for quadratically nonlinear systems of ODEs.
 #
-# We will demonstrate this using a canonical 2D fluid flow (2D so that the enstrophy is conserved).
-# In many cases, the wake behind a bluff body is characterized by a periodic vortex shedding phenomenon known as a von Karman street.
-# The two-dimensional incompressible flow past a cylinder is a stereotypical example of such behavior.
-# The transient energy growth and saturation amplitude of this instability mode is of particular interest and has historically posed a significant modeling challenge.
-# Noack et al. (2003) used an 8-mode POD basis that was augmented with a ninth "shift mode" parameterizing a mean flow deformation. The 9-mode quadratic Galerkin model does resolve the transient dynamics, nonlinear stability mechanism, and post-transient oscillation, accurately reproducing all of the key physical features of the vortex street. In general, the totally symmetric structure in the quadratic coefficients $Q_{ijk}$ is weakly *broken* in models for the von Karman street, but it can be enforced to hold exactly.
+#  We will demonstrate this using a canonical 2D fluid flow (2D so that the enstrophy is conserved).
+#  In many cases, the wake behind a bluff body is characterized by a periodic vortex shedding phenomenon known as a von Karman street.
+#  The two-dimensional incompressible flow past a cylinder is a stereotypical example of such behavior.
+#  The transient energy growth and saturation amplitude of this instability mode is of particular interest and has historically posed a significant modeling challenge.
+#  Noack et al. (2003) used an 8-mode POD basis that was augmented with a ninth "shift mode" parameterizing a mean flow deformation. The 9-mode quadratic Galerkin model does resolve the transient dynamics, nonlinear stability mechanism, and post-transient oscillation, accurately reproducing all of the key physical features of the vortex street. In general, the totally symmetric structure in the quadratic coefficients $Q_{ijk}$ is weakly *broken* in models for the von Karman street, but it can be enforced to hold exactly.
 #
-# This is precisely what is done in Schlegel and Noack (2015), and in this perfectly-skew-symmetric case, the global stability of the quadratic model was proven with $m_9 = m_\text{shift} = \epsilon$, $\epsilon > 1$, and $m_i = 0$ for $i = \{1,...,8\}$. Indeed, we do similarly for data-driven models obtained with the SINDy method in the trapping_sindy_paper_examples.ipynb Jupyter notebook corresponding to the Trapping SINDy paper (https://journals.aps.org/prfluids/abstract/10.1103/PhysRevFluids.6.094401).
+#  This is precisely what is done in Schlegel and Noack (2015), and in this perfectly-skew-symmetric case, the global stability of the quadratic model was proven with $m_9 = m_\text{shift} = \epsilon$, $\epsilon > 1$, and $m_i = 0$ for $i = \{1,...,8\}$. Indeed, we do similarly for data-driven models obtained with the SINDy method in the trapping_sindy_paper_examples.ipynb Jupyter notebook corresponding to the Trapping SINDy paper (https://journals.aps.org/prfluids/abstract/10.1103/PhysRevFluids.6.094401).
 #
-# ### This notebook will show that both the energy and enstrophy can be used with the trapping theorem to promote global stability.
+#  ### This notebook will show that both the energy and enstrophy can be used with the trapping theorem to promote global stability.
 # %%
 import warnings
 
@@ -36,7 +36,6 @@
     get_vorticity,
     obj_function,
     sindy_library_no_bias,
-    check_stability,
     check_local_stability,
     make_trap_progress_plots,
     make_bar,
@@ -134,9 +133,9 @@
 
 # %% [markdown]
 # ### Compute the velocity, vorticity, and enstrophy
-# The first thing we will do is promote globally stable models via the trapping SINDy algorithm, as in the trapping_sindy_paper_examples.ipynb notebook. We can actually do this by calculating a stability matrix $\mathbf{A}^S$ with *either* the energy or the enstrophy (or any other positive definite, quadratic quantity that plays an important role in the dynamics). So we begin by constructing globally stable trapping models of the von Karman street in these two ways.
+#  The first thing we will do is promote globally stable models via the trapping SINDy algorithm, as in the trapping_sindy_paper_examples.ipynb notebook. We can actually do this by calculating a stability matrix $\mathbf{A}^S$ with *either* the energy or the enstrophy (or any other positive definite, quadratic quantity that plays an important role in the dynamics). So we begin by constructing globally stable trapping models of the von Karman street in these two ways.
 #
-# First, we load in the data and compute the energy and enstrophy.
+#  First, we load in the data and compute the energy and enstrophy.
 
 # %%
 # path to POD mode files
@@ -238,13 +237,13 @@
 
 # %% [markdown]
 # ### Check global stability of the POD-Galerkin models
-# Okay, so we have loaded in some DNS data from the von Karman Street and generated (analytic) 5D POD-Galerkin models for this system. The skew-symmetric models below are globally stable *if and only if* there exists a vector $\mathbf{m}$ such that following matrix is negative definite:
-# $$A^S_{ij} = L^S_{ij} + (Q_{ijk} + Q_{jik})m_k.$$
-# Note that if the quadratic term $Q_{ijk}$ has no totally-symmetric part this is equal to
-# $$A^S_{ij} = L^S_{ij} - Q_{kij}m_k.$$
-# A negative definite $\mathbf{A}^S$ turns out to also be necessary when $Q_{ijk}$ does have nontrivial totally symmetric component, but in this case is not sufficient for global boundedness and we can promote local stability as in the trapping_extended.ipynb notebook.
+#  Okay, so we have loaded in some DNS data from the von Karman Street and generated (analytic) 5D POD-Galerkin models for this system. The skew-symmetric models below are globally stable *if and only if* there exists a vector $\mathbf{m}$ such that following matrix is negative definite:
+#  $$A^S_{ij} = L^S_{ij} + (Q_{ijk} + Q_{jik})m_k.$$
+#  Note that if the quadratic term $Q_{ijk}$ has no totally-symmetric part this is equal to
+#  $$A^S_{ij} = L^S_{ij} - Q_{kij}m_k.$$
+#  A negative definite $\mathbf{A}^S$ turns out to also be necessary when $Q_{ijk}$ does have nontrivial totally symmetric component, but in this case is not sufficient for global boundedness and we can promote local stability as in the trapping_extended.ipynb notebook.
 #
-# Next we check with a simple nonlinear algorithm (simulated annealing) that our analytic models can be shown to be globally stable (there is an $\mathbf{m}$ such that $\mathbf{A}^S$ is negative definite) using both the energy or the enstrophy to construct $\mathbf{A}^S$.
+#  Next we check with a simple nonlinear algorithm (simulated annealing) that our analytic models can be shown to be globally stable (there is an $\mathbf{m}$ such that $\mathbf{A}^S$ is negative definite) using both the energy or the enstrophy to construct $\mathbf{A}^S$.
 
 # %%
 # Import simulated annealing algorithm from scipy
@@ -252,17 +251,22 @@
 
 # Search between -5000, 5000 for each component of m
 boundvals = np.zeros((r, 2))
-boundmax = 5000
-boundmin = -5000
+boundmax = 10
+boundmin = -10
 boundvals[:, 0] = boundmin
 boundvals[:, 1] = boundmax
 
 Ls_enstrophy = 0.5 * (L_enstrophy + L_enstrophy.T)
+Ls = 0.5 * (galerkin5["L"] + galerkin5["L"].T)
 # Run simulated annealing for the enstrophy eigenbasis
 algo_sol = anneal_algo(
     obj_function,
     bounds=boundvals,
-    args=(Ls_enstrophy, galerkin5["Q_ep"], P_enstrophy),
+    args=(
+        galerkin5["L"],
+        galerkin5["Q_ep"],
+        P_enstrophy,
+    ),  # Factors of P_enstrophy taken care of in obj_func
     maxiter=200,
 )
 opt_m = algo_sol.x
@@ -277,9 +281,11 @@
 )
 
 # Repeat using the energy
-Ls = 0.5 * (galerkin5["L"] + galerkin5["L"].T)
 algo_sol = anneal_algo(
-    obj_function, bounds=boundvals, args=(Ls, galerkin5["Q_ep"], P_energy), maxiter=1000
+    obj_function,
+    bounds=boundvals,
+    args=(galerkin5["L"], galerkin5["Q_ep"], P_energy),
+    maxiter=1000,
 )
 opt_m = algo_sol.x
 opt_energy = algo_sol.fun
@@ -294,15 +300,16 @@
 
 # %% [markdown]
 # #### We have proven that both the models, with totally symmetric quadratic tensors, are globally stable.
-# We now fix the coefficients of $\mathbf{m}$ except for the direction of the shift-mode, and scan this value to see how the largest eigenvalue of $\mathbf{A}^S$ (and corresponding trapping region radius) varies.
+#  We now fix the coefficients of $\mathbf{m}$ except for the direction of the shift-mode, and scan this value to see how the largest eigenvalue of $\mathbf{A}^S$ (and corresponding trapping region radius) varies.
 
 # %%
-N = 5000
-alphas = np.linspace(0, 10000, N)
+N = 500
+alphas = np.linspace(0, 20, N)
+r = 5
 m = np.zeros(r)
-m[:-1] = [1.05133223e-01, 3.65030827e-02, 5.29009012e01, 3.14215584e01]
+# m[:-1] = [4.40596938e-04, 2.68248063e-04, 1.44130186e-01, 4.11438479e-02]
 m_enstrophy = np.zeros(r)
-m_enstrophy[:-1] = [3.87083390e-01, 5.23018605e-01, 2.13126168e01, 8.95957413e01]
+# m_enstrophy[:-1] = [5.99535702e-04, -1.32512857e-04,  1.15691624e-01,  1.80317529e-02]
 
 obj_energy = np.zeros(N)
 obj_enstrophy = np.zeros(N)
@@ -319,14 +326,17 @@
     m_enstrophy[-1] = alpha
     obj_energy[i] = obj_function(m, Ls, galerkin5["Q_ep"], P_energy)
     obj_enstrophy[i] = obj_function(
-        m_enstrophy, Ls_enstrophy, galerkin5["Q_ep"], P_enstrophy
+        m_enstrophy, galerkin5["L"], galerkin5["Q_ep"], P_enstrophy
     )
     d_energy = np.dot(galerkin5["L"], m) + np.dot(
         np.tensordot(galerkin5["Q_ep"], m, axes=([2], [0])), m
     )
-    d_enstrophy = np.dot(L_enstrophy, m_enstrophy) + np.dot(
+    lsv, sing_vals, rsv = np.linalg.svd(P_enstrophy)
+    P_rt = lsv @ np.diag(np.sqrt(sing_vals)) @ rsv
+    d_enstrophy = np.dot(Ls, m_enstrophy) + np.dot(
         np.tensordot(galerkin5["Q_ep"], m_enstrophy, axes=([2], [0])), m_enstrophy
     )
+    d_enstrophy = P_rt @ d_enstrophy
     Rm_energy[i] = np.linalg.norm(d_energy) / np.abs(obj_energy[i])
     Rm_enstrophy[i] = np.linalg.norm(d_enstrophy) / np.abs(obj_enstrophy[i])
 
@@ -334,24 +344,28 @@
 plt.plot(alphas, obj_energy * 1e2, label=r"$\lambda_1^{energy}$ x $10^2$")
 plt.plot(
     alphas[obj_energy < 0],
-    Rm_energy[obj_energy < 0] / 1.0e4,
-    label="$R_m^{energy}$ x $10^{-4}$",
+    Rm_energy[obj_energy < 0],
+    label="$R_m^{energy}$",
 )
-plt.plot(alphas, obj_enstrophy, label=r"$\lambda_1^{enstrophy}$")
+plt.plot(alphas, obj_enstrophy * 1e2, label=r"$\lambda_1^{enstrophy}$ x $10^2$")
 plt.plot(
     alphas[obj_enstrophy < 0],
-    Rm_enstrophy[obj_enstrophy < 0] / 1.0e4,
-    label="$R_m^{enstrophy}$ x $10^{-4}$",
+    Rm_enstrophy[obj_enstrophy < 0],
+    label="$R_m^{enstrophy}$",
 )
 plt.legend(fontsize=12, loc="upper right", framealpha=1.0)
 plt.ylim(-5, 50)
-plt.xlim(0, 10000)
+# plt.xlim(0, 10000)
 plt.grid(True)
 plt.xlabel(r"$m_{shift}$", fontsize=18)
 plt.xticks(fontsize=12)
 plt.yticks(fontsize=12)
 plt.show()
 
+# %% [markdown]
+# ### Data-driven 5D models for the von Karman street using Trapping SINDy
+#  We have investigated a number of 5D (analytic) POD-Galerkin models and shown that the analytic, constrained models are globally stable. Now we show that trapping SINDy can be used with the energy or enstrophy to build data-driven models that are also globally stable.
+
 # %%
 galerkin_ep = (
     gQ
@@ -373,22 +387,21 @@
 ) / 6.0
 print(np.max(abs(galerkin_ep)))
 
-# %% [markdown]
-# ### Data-driven 5D models for the von Karman street using Trapping SINDy
-# We have investigated a number of 5D (analytic) POD-Galerkin models and shown that the analytic, constrained models are globally stable. Now we show that trapping SINDy can be used with the energy or enstrophy to build data-driven models that are also globally stable.
-
 # %%
-max_iter = 5000
-eta = 1.0e2
-threshold = 0
-alpha_m = 9e-1 * eta
+# define hyperparameters
+max_iter = 10000
+eta = 1.0
+
+# don't need a reg_weight_lam if eta is sufficiently small
+# which is good news because CVXPY is much slower
+reg_weight_lam = 0
+alpha_m = 1e-1 * eta
 
 # run trapping SINDy
 sindy_opt = ps.TrappingSR3(
-    method="global",
     _n_tgts=5,
     _include_bias=False,
-    threshold=threshold,
+    reg_weight_lam=reg_weight_lam,
     eta=eta,
     alpha_m=alpha_m,
     max_iter=max_iter,
@@ -397,36 +410,32 @@
 model = ps.SINDy(
     optimizer=sindy_opt,
     feature_library=sindy_library_no_bias,
+    differentiation_method=ps.FiniteDifference(drop_endpoints=True),
 )
 model.fit(a, t=t)
 Xi = model.coefficients().T
 PL_tensor = sindy_opt.PL_
 PQ_tensor = sindy_opt.PQ_
-Lenergy = np.tensordot(PL_tensor, Xi, axes=([3, 2], [0, 1]))
-Qenergy = np.tensordot(PQ_tensor, Xi, axes=([4, 3], [0, 1]))
-mean_val = np.mean(a, axis=0)
-mean_val = np.sqrt(np.sum(mean_val**2))
-check_stability(r, Xi, sindy_opt, mean_val)
-energy_model = model
-Q = np.tensordot(sindy_opt.PQ_, Xi, axes=([4, 3], [0, 1]))
+L = np.tensordot(PL_tensor, Xi, axes=([3, 2], [0, 1]))
+Q = np.tensordot(PQ_tensor, Xi, axes=([4, 3], [0, 1]))
 Q_sum = np.max(np.abs((Q + np.transpose(Q, [1, 2, 0]) + np.transpose(Q, [2, 0, 1]))))
-print("Maximum deviation from the constraints = ", Q_sum)
+print("Max deviation from the constraints = ", Q_sum)
 
 # %% [markdown]
 # The previous model finds an $\mathbf{m}$ such that $\mathbf{A}^S$ is negative definite, while also fitting the data. Now we can repeat but in the eigenbasis of enstrophy. If the enstrophy is $H = \mathbf{y}^T\mathcal{P}\mathbf{A}^S \mathbf{y}$, now we are searching for an $\mathbf{m}$ such that $\mathcal{P}\mathbf{A}^S$ is negative definite.
 
 # %%
-max_iter = 5000
-eta = 1.0e4
-threshold = 0
-alpha_m = 1e-1 * eta
+max_iter = 2000
+eta = 1.0
+reg_weight_lam = 0
+alpha_m = 1e-2 * eta
 
 mod_matrix = enstrophy_integrals[1:, 1:]
 sindy_opt = ps.TrappingSR3(
     method="global",
     _n_tgts=5,
     _include_bias=False,
-    threshold=threshold,
+    reg_weight_lam=reg_weight_lam,
     eta=eta,
     alpha_m=alpha_m,
     max_iter=max_iter,
@@ -443,29 +452,30 @@
 Qenstrophy = np.tensordot(PQ_tensor, Xi, axes=([4, 3], [0, 1]))
 mean_val = np.mean(a, axis=0)
 mean_val = np.sqrt(np.sum(mean_val**2))
-check_stability(r, Xi, sindy_opt, mean_val, mod_matrix)
+check_local_stability(Xi, sindy_opt, mean_val)
 enstrophy_model = model
-Q = np.tensordot(sindy_opt.PQ_, Xi, axes=([4, 3], [0, 1]))
-Q_sum = np.max(np.abs((Q + np.transpose(Q, [1, 2, 0]) + np.transpose(Q, [2, 0, 1]))))
-print("Maximum deviation from the constraints = ", Q_sum)
+# Q = np.tensordot(sindy_opt.PQ_, Xi, axes=([4, 3], [0, 1]))
+# Q_sum = np.max(np.abs((Q + np.transpose(Q, [1, 2, 0]) + np.transpose(Q, [2, 0, 1]))))
+# print("Maximum deviation from the constraints = ", Q_sum)
+make_progress_plots(r, sindy_opt)
 
 # %% [markdown]
 # Enstrophy model was successful!
-# #### We have built two data-driven models, one using the energy as a Lyapunov function for trapping SINDy, and the other using the enstrophy. Now we compare the distribution of coefficients in each model.
+#  #### We have built two data-driven models, one using the energy as a Lyapunov function for trapping SINDy, and the other using the enstrophy. Now we compare the distribution of coefficients in each model.
 
 # %%
-make_bar(galerkin5, Lenergy, Qenergy, Lenstrophy, Qenstrophy)
+make_bar(galerkin5, L, Q, Lenstrophy, Qenstrophy)
 
 # %% [markdown]
 # ### Compare the models
-# Below, we compare the 5D POD-Galerkin with trapping SINDy models obtained by considering the energy or by the enstrophy. Both trapping models outperform the POD-Galerkin model.
+#  Below, we compare the 5D POD-Galerkin with trapping SINDy models obtained by considering the energy or by the enstrophy. Both trapping models outperform the POD-Galerkin model.
 
 # %%
 # Interpolate onto better time base
 t_traj = np.linspace(t[0], t[-1], len(t) * 1)
 
 # simulate trapping SINDy results
-xtraj_energy = energy_model.simulate(a0, t_traj)
+xtraj_energy = model.simulate(a0, t_traj)
 xtraj_enstrophy = enstrophy_model.simulate(a0, t_traj)
 
 # simulate and plot 5D von Karman trajectory results
@@ -571,18 +581,17 @@
 # %%
 max_iter = 5000
 eta = 1.0e4
-threshold = 0
-alpha_m = 1e-2 * eta
+reg_weight_lam = 0
+alpha_m = 1e-1 * eta
 # alpha = 1e10
-beta = 1e-12
-
+beta = 1e-10
 
 mod_matrix = enstrophy_integrals[1:, 1:]
 sindy_opt = ps.TrappingSR3(
     method="local",
     _n_tgts=5,
     _include_bias=False,
-    threshold=threshold,
+    reg_weight_lam=reg_weight_lam,
     eta=eta,
     gamma=-1,
     alpha_m=alpha_m,
@@ -598,34 +607,33 @@
 )
 model.fit(a, t=t)
 Xi = model.coefficients().T
+PL_tensor = sindy_opt.PL_
+PQ_tensor = sindy_opt.PQ_
 Lenstrophy = np.tensordot(PL_tensor, Xi, axes=([3, 2], [0, 1]))
 Qenstrophy = np.tensordot(PQ_tensor, Xi, axes=([4, 3], [0, 1]))
 mean_val = np.mean(a, axis=0)
 mean_val = np.sqrt(np.sum(mean_val**2))
-check_stability(r, Xi, sindy_opt, mean_val, mod_matrix)
-
-check_local_stability(Xi, sindy_opt, mean_val, mod_matrix)
+check_local_stability(Xi, sindy_opt, mean_val)
 enstrophy_model = model
 Q = np.tensordot(sindy_opt.PQ_, Xi, axes=([4, 3], [0, 1]))
 Q = np.tensordot(mod_matrix, Q, axes=([1], [0]))
 Q_sum = np.max(np.abs((Q + np.transpose(Q, [1, 2, 0]) + np.transpose(Q, [2, 0, 1]))))
 print("Maximum deviation from the constraints = ", Q_sum)
-rhos_minus, rhos_plus = make_trap_progress_plots(r, sindy_opt, mod_matrix)
+rhos_minus, rhos_plus = make_trap_progress_plots(r, sindy_opt)
 
 # %%
 make_progress_plots(r, sindy_opt)
 
 # %%
-# Interpolate onto better time base
 fs = 18
-t_traj = np.linspace(t[0], t[-1], len(t) * 2)
+t_traj = np.linspace(t[0], t[-1], len(t))  # * 2)
 fig, ax = plt.subplots(3, 3, subplot_kw={"projection": "3d"}, figsize=(16, 16))
 a0s = [(np.random.rand(5) - 0.5) * 20 for i in range(2)]
 for a0 in a0s:
     print(a0)
 
     # simulate trapping SINDy results
-    xtraj_energy = energy_model.simulate(a0, t_traj)
+    xtraj_energy = model.simulate(a0, t_traj)
     xtraj_enstrophy = enstrophy_model.simulate(a0, t_traj)
 
     # simulate and plot 5D von Karman trajectory results
@@ -746,7 +754,7 @@
         ax[i, 2].set_xlim(-xlim, xlim)
         ax[i, 2].set_ylim(-ylim, ylim)
         ax[i, 2].set_zlim(zlim_minus, zlim)
-plt.savefig("enstrophy_results.pdf")
+# plt.savefig("enstrophy_results.pdf")
 
 # %%
 # define energies of the DNS, and both the 5D and 9D models