From 06b715d07aef9ebc2cebe44768ec5638ddd27065 Mon Sep 17 00:00:00 2001
From: Jake Stevens-Haas <37048747+Jacob-Stevens-Haas@users.noreply.github.com>
Date: Thu, 30 Nov 2023 01:29:15 +0000
Subject: [PATCH] CLN: Extract function to encapsulate coefficient-solving step
 of Trapping

This makes the algorithm in code look like algorithm in paper
---
 pysindy/optimizers/trapping_sr3.py | 41 +++++++++++++++++-------------
 1 file changed, 23 insertions(+), 18 deletions(-)

diff --git a/pysindy/optimizers/trapping_sr3.py b/pysindy/optimizers/trapping_sr3.py
index eea4aef36..51aa17fdf 100644
--- a/pysindy/optimizers/trapping_sr3.py
+++ b/pysindy/optimizers/trapping_sr3.py
@@ -423,6 +423,7 @@ def _reduce(self, x, y):
         self.history_ = []
         n_samples, n_features = x.shape
         n_tgts = y.shape[1]
+        var_len = n_features * n_tgts
         n_feat_expected = int((n_tgts**2 + 3 * n_tgts) / 2.0)
 
         # Only relevant if the stability term is turned on.
@@ -494,31 +495,19 @@ def _reduce(self, x, y):
             if self.relax_optim:
                 if self.threshold > 0.0:
                     # sparse relax_and_split
-                    xi, cost = self._create_var_and_part_cost(
-                        n_features * n_tgts, x_expanded, y
-                    )
-                    cost = cost + cp.sum_squares(Pmatrix @ xi - A.flatten()) / self.eta
-                    coef_sparse = self._update_coef_cvxpy(
-                        xi, cost, n_tgts * n_features, coef_prev, self.eps_solver
+                    coef_sparse = self._update_coef_sparse_rs(
+                        var_len, x_expanded, y, Pmatrix, A, coef_prev
                     )
                 else:
-                    pTp = np.dot(Pmatrix.T, Pmatrix)
-                    H = xTx + pTp / self.eta
-                    P_transpose_A = np.dot(Pmatrix.T, A.flatten())
-                    coef_sparse = self._solve_nonsparse_relax_and_split(
-                        H, xTy, P_transpose_A, coef_prev
+                    coef_sparse = self._update_coef_nonsparse_rs(
+                        Pmatrix, A, coef_prev, xTx, xTy
                     )
                 trap_prev_ctr, trap_ctr, A, tk_prev = self._solve_m_relax_and_split(
                     trap_prev_ctr, trap_ctr, A, coef_sparse, tk_prev
                 )
             else:
-                var_len = n_tgts * n_features
-                xi, cost = self._create_var_and_part_cost(var_len, x_expanded, y)
-                cost += (
-                    cp.lambda_max(cp.reshape(Pmatrix @ xi, (n_tgts, n_tgts))) / self.eta
-                )
-                coef_sparse = self._update_coef_cvxpy(
-                    xi, cost, var_len, coef_prev, self.eps_solver
+                coef_sparse = self._update_coef_direct(
+                    var_len, x_expanded, y, Pmatrix, coef_prev, n_tgts
                 )
                 trap_ctr = self._solve_m_direct(n_tgts, coef_sparse)
 
@@ -556,3 +545,19 @@ def _reduce(self, x, y):
             )
 
         self.coef_ = coef_sparse.T
+
+    def _update_coef_sparse_rs(self, var_len, x_expanded, y, Pmatrix, A, coef_prev):
+        xi, cost = self._create_var_and_part_cost(var_len, x_expanded, y)
+        cost = cost + cp.sum_squares(Pmatrix @ xi - A.flatten()) / self.eta
+        return self._update_coef_cvxpy(xi, cost, var_len, coef_prev, self.eps_solver)
+
+    def _update_coef_nonsparse_rs(self, Pmatrix, A, coef_prev, xTx, xTy):
+        pTp = np.dot(Pmatrix.T, Pmatrix)
+        H = xTx + pTp / self.eta
+        P_transpose_A = np.dot(Pmatrix.T, A.flatten())
+        return self._solve_nonsparse_relax_and_split(H, xTy, P_transpose_A, coef_prev)
+
+    def _update_coef_direct(self, var_len, x_expanded, y, Pmatrix, coef_prev, n_tgts):
+        xi, cost = self._create_var_and_part_cost(var_len, x_expanded, y)
+        cost += cp.lambda_max(cp.reshape(Pmatrix @ xi, (n_tgts, n_tgts))) / self.eta
+        return self._update_coef_cvxpy(xi, cost, var_len, coef_prev, self.eps_solver)