NBK: Update notebooks to new SR3 threshold vs weight param

dynamicslab · Sep 17, 2024 · fa65917 · fa65917
1 parent 2249c3e
commit fa65917
Show file tree

Hide file tree

Showing 8 changed files with 268 additions and 274 deletions.
diff --git 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
@@ -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
@@ -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)"