From f63b4c31fe210f2df1b3c866783f34fc5d1a2264 Mon Sep 17 00:00:00 2001
From: Jake <37048747+Jacob-Stevens-Haas@users.noreply.github.com>
Date: Tue, 5 Dec 2023 17:42:43 +0000
Subject: [PATCH 01/16] ENH: Add function to calculate number of polynomial
 features

This is much faster than fitting the features, and can be used
by trapping to replace some constraint index calculations.

Also:
Add docstring/type annotations
Rename a variable so it doesn't shadow newly-imported name
---
 pysindy/feature_library/polynomial_library.py | 66 ++++++++++++++++---
 test/test_feature_library.py                  |  4 ++
 2 files changed, 62 insertions(+), 8 deletions(-)

diff --git a/pysindy/feature_library/polynomial_library.py b/pysindy/feature_library/polynomial_library.py
index 75dbf5637..30710041a 100644
--- a/pysindy/feature_library/polynomial_library.py
+++ b/pysindy/feature_library/polynomial_library.py
@@ -1,7 +1,9 @@
 from itertools import chain
+from math import comb
 from typing import Iterator
 
 import numpy as np
+from numpy.typing import NDArray
 from scipy import sparse
 from sklearn.preprocessing import PolynomialFeatures
 from sklearn.utils.validation import check_is_fitted
@@ -73,8 +75,22 @@ def __init__(
 
     @staticmethod
     def _combinations(
-        n_features, degree, include_interaction, interaction_only, include_bias
-    ) -> Iterator[tuple]:
+        n_features: int,
+        degree: int,
+        include_interaction: bool,
+        interaction_only: bool,
+        include_bias: bool,
+    ) -> Iterator[tuple[int, ...]]:
+        """
+        Create selection tuples of input indexes for each polynomail term
+
+        Selection tuple iterates the input indexes present in a single term.
+        For example, (x+y+1)^2 would be in iterator of the tuples:
+        (), (0,), (1,), (0, 0), (0, 1), (1, 1)
+        1    x     y      x^2     x*y     y^2
+
+        Order of terms is preserved regardless of include_interation.
+        """
         if not include_interaction:
             return chain(
                 [()] if include_bias else [],
@@ -93,7 +109,12 @@ def _combinations(
         )
 
     @property
-    def powers_(self):
+    def powers_(self) -> NDArray[np.dtype("int")]:
+        """
+        The exponents of the polynomial as an array of shape
+        (n_features_out, n_features_in), where each item is the exponent of the
+        jth input variable in the ith polynomial term.
+        """
         check_is_fitted(self)
         combinations = self._combinations(
             n_features=self.n_features_in_,
@@ -208,10 +229,10 @@ def transform(self, x_full):
             )
             if sparse.isspmatrix(x):
                 columns = []
-                for comb in combinations:
-                    if comb:
+                for combo in combinations:
+                    if combo:
                         out_col = 1
-                        for col_idx in comb:
+                        for col_idx in combo:
                             out_col = x[..., col_idx].multiply(out_col)
                         columns.append(out_col)
                     else:
@@ -227,7 +248,36 @@ def transform(self, x_full):
                     ),
                     x.__dict__,
                 )
-                for i, comb in enumerate(combinations):
-                    xp[..., i] = x[..., comb].prod(-1)
+                for i, combo in enumerate(combinations):
+                    xp[..., i] = x[..., combo].prod(-1)
             xp_full = xp_full + [xp]
         return xp_full
+
+
+def n_poly_features(
+    n_in_feat: int,
+    degree: int,
+    include_bias: bool = False,
+    include_interation: bool = True,
+    interaction_only: bool = False,
+) -> int:
+    """Calculate number of polynomial features
+
+    Args:
+        n_in_feat: number of input features, e.g. 3 for x, y, z
+        degree: polynomial degree, e.g. 2 for quadratic
+        include_bias: whether to include a constant term
+        include_interaction: whether to include terms mixing multiple inputs
+        interaction_only: whether to omit terms of x_m * x_n^p for p > 1
+    """
+    if not include_interation and interaction_only:
+        raise ValueError("Cannot set interaction only if include_interaction is False")
+    n_feat = include_bias
+    if not include_interation:
+        return n_feat + n_in_feat * degree
+    for deg in range(1, degree + 1):
+        if interaction_only:
+            n_feat += comb(n_in_feat, deg)
+        else:
+            n_feat += comb(n_in_feat + deg - 1, deg)
+    return n_feat
diff --git a/test/test_feature_library.py b/test/test_feature_library.py
index 8e98b1a0d..8bf27c580 100644
--- a/test/test_feature_library.py
+++ b/test/test_feature_library.py
@@ -21,6 +21,7 @@
 from pysindy.feature_library import TensoredLibrary
 from pysindy.feature_library import WeakPDELibrary
 from pysindy.feature_library.base import BaseFeatureLibrary
+from pysindy.feature_library.polynomial_library import n_poly_features
 from pysindy.optimizers import SINDyPI
 from pysindy.optimizers import STLSQ
 
@@ -880,3 +881,6 @@ def test_polynomial_combinations(include_interaction, interaction_only, bias, ex
     )
     result = tuple(sorted(list(combos)))
     assert result == expected
+    assert len(expected) == n_poly_features(
+        2, 2, bias, include_interaction, interaction_only
+    )

From 9c8c5d8e0d79a4d70d16cb5f470783a4b03f40b3 Mon Sep 17 00:00:00 2001
From: Jake <37048747+Jacob-Stevens-Haas@users.noreply.github.com>
Date: Fri, 8 Dec 2023 16:47:07 +0000
Subject: [PATCH 02/16] ENH: add _make_constraints to TrappingSR3

There are several major differences from the version in example 8:
* Coupled with PolynomialLibrary logic using PolynomialLibrary.powers_.
Previously, this coupling was implicit and actually tied to the specific order
in Example 8, not the order users would expect.
* Instead of handling pre-flattened 2D array, constraints are built with
target and feature axes separately, letting the caller handle formatting
the constraints as "target" or "feature".  This gives us flexibility to
later clean up the constraint logic outside this function without needing to
return here and removes the need for striding calculations like:
n_tgts * (n_tgts + k - 1) + i + n_tgts * int(j * (2 * n_tgts - j - 3) / 2.0)
Since constraints are built directly from the terms and exponents in the
library, its much easier to read and debug the code.
* Renamed "r" as "n_tgts", etc.  Single-letter index variables are
removed.  When they occur in homogenous iterators, they are concatenated with
what they're iterating (e.g. tgt_i, tgt_j)
* Broke out the different kinds of constraints into helper functions to make
it easier to test.

This allows follow on work to auto-attach constraints to a problem.
---
 pysindy/optimizers/trapping_sr3.py | 94 ++++++++++++++++++++++++++++++
 1 file changed, 94 insertions(+)

diff --git a/pysindy/optimizers/trapping_sr3.py b/pysindy/optimizers/trapping_sr3.py
index dad0aa8c9..def9d71ee 100644
--- a/pysindy/optimizers/trapping_sr3.py
+++ b/pysindy/optimizers/trapping_sr3.py
@@ -1,12 +1,19 @@
 import warnings
+from itertools import combinations as combo_nr
+from itertools import combinations_with_replacement as combo_wr
+from math import comb
 from typing import Tuple
 
 import cvxpy as cp
 import numpy as np
+from numpy import intp
+from numpy.typing import NDArray
 from scipy.linalg import cho_factor
 from scipy.linalg import cho_solve
 from sklearn.exceptions import ConvergenceWarning
 
+from ..feature_library.polynomial_library import n_poly_features
+from ..feature_library.polynomial_library import PolynomialLibrary
 from ..utils import reorder_constraints
 from .sr3 import SR3
 
@@ -813,3 +820,90 @@ def _reduce(self, x, y):
 
         self.coef_ = coef_sparse.T
         self.objective_history = objective_history
+
+
+def _make_constraints(n_tgts: int):
+    """Create constraints for the Quadratic terms in TrappingSR3.
+
+    These are the constraints from equation 5 of the Trapping SINDy paper.
+
+    Args:
+        n_tgts: number of coordinates or modes for which you're fitting an ODE.
+
+    Returns:
+        A tuple of the constraint zeros, and a constraint matrix to multiply
+        by the coefficient matrix of Polynomial terms. Number of constraints is
+        ``n_tgts + 2 * math.comb(n_tgts, 2) + math.comb(n_tgts, 3)``.
+        Constraint matrix is of shape ``(n_constraint, n_feature, n_tgt)``.
+        To get "feature" order constraints, use
+        ``np.reshape(constraint_matrix, (n_constraints, -1))``.
+        To get "target" order constraints, transpose axis 1 and 2 before
+        reshaping.
+    """
+    n_terms = n_poly_features(n_tgts, degree=2, include_bias=False)
+    lib = PolynomialLibrary(2, include_bias=False).fit(np.zeros((1, n_tgts)))
+    terms = [(t_ind, exps) for t_ind, exps in enumerate(lib.powers_)]
+
+    # index of tgt -> index of its pure quadratic term
+    pure_terms = {np.argmax(exps): t_ind for t_ind, exps in terms if max(exps) == 2}
+    # two indexes of tgts -> index of their mixed quadratic term
+    mixed_terms = {
+        frozenset(np.argwhere(exponent == 1).flatten()): t_ind
+        for t_ind, exponent in terms
+        if max(exponent) == 1 and sum(exponent) == 2
+    }
+    constraint_mat = np.vstack(
+        (
+            _pure_constraints(n_tgts, n_terms, pure_terms),
+            _antisymm_double_constraint(n_tgts, n_terms, pure_terms, mixed_terms),
+            _antisymm_triple_constraints(n_tgts, n_terms, mixed_terms),
+        )
+    )
+
+    return np.zeros(len(constraint_mat)), constraint_mat
+
+
+def _pure_constraints(
+    n_tgts: int, n_terms: int, pure_terms: dict[intp, int]
+) -> NDArray[np.dtype("float")]:
+    """Set constraints for coefficients adorning terms like a_i^3 = 0"""
+    constraint_mat = np.zeros((n_tgts, n_terms, n_tgts))
+    for constr_ind, (tgt_ind, term_ind) in zip(range(n_tgts), pure_terms.items()):
+        constraint_mat[constr_ind, term_ind, tgt_ind] = 1.0
+    return constraint_mat
+
+
+def _antisymm_double_constraint(
+    n_tgts: int,
+    n_terms: int,
+    pure_terms: dict[intp, int],
+    mixed_terms: dict[frozenset[intp] : int],
+) -> NDArray[np.dtype("float")]:
+    """Set constraints for coefficients adorning terms like a_i^2 * a_j=0"""
+    constraint_mat_1 = np.zeros((comb(n_tgts, 2), n_terms, n_tgts))  # a_i^2 * a_j
+    constraint_mat_2 = np.zeros((comb(n_tgts, 2), n_terms, n_tgts))  # a_i * a_j^2
+    for constr_ind, ((tgt_i, tgt_j), mix_term) in zip(
+        range(n_tgts), mixed_terms.items()
+    ):
+        constraint_mat_1[constr_ind, mix_term, tgt_i] = 1.0
+        constraint_mat_1[constr_ind, pure_terms[tgt_i], tgt_j] = 1.0
+        constraint_mat_2[constr_ind, mix_term, tgt_j] = 1.0
+        constraint_mat_2[constr_ind, pure_terms[tgt_j], tgt_i] = 1.0
+
+    return np.concatenate((constraint_mat_1, constraint_mat_2), axis=0)
+
+
+def _antisymm_triple_constraints(
+    n_tgts: int, n_terms: int, mixed_terms: dict[frozenset[intp] : int]
+) -> NDArray[np.dtype("float")]:
+    constraint_mat = np.zeros((comb(n_tgts, 3), n_terms, n_tgts))  # a_ia_ja_k
+
+    def find_symm_term(a, b):
+        return mixed_terms[frozenset({a, b})]
+
+    for constr_ind, (tgt_i, tgt_j, tgt_k) in enumerate(combo_nr(range(n_tgts), 3)):
+        constraint_mat[constr_ind, find_symm_term(tgt_j, tgt_k), tgt_i] = 1
+        constraint_mat[constr_ind, find_symm_term(tgt_k, tgt_i), tgt_j] = 1
+        constraint_mat[constr_ind, find_symm_term(tgt_i, tgt_j), tgt_k] = 1
+
+    return constraint_mat

From 37935e15f209ad0b0f000c5139d0c58c12453d5b Mon Sep 17 00:00:00 2001
From: Jake <37048747+Jacob-Stevens-Haas@users.noreply.github.com>
Date: Sun, 10 Dec 2023 17:53:58 +0000
Subject: [PATCH 03/16] TST: Add Trapping _make_constraints test

---
 test/test_optimizers.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/test/test_optimizers.py b/test/test_optimizers.py
index c69ce9823..385d809c7 100644
--- a/test/test_optimizers.py
+++ b/test/test_optimizers.py
@@ -30,6 +30,8 @@
 from pysindy.optimizers import TrappingSR3
 from pysindy.optimizers import WrappedOptimizer
 from pysindy.optimizers.stlsq import _remove_and_decrement
+from pysindy.optimizers.trapping_sr3 import _antisymm_triple_constraints
+from pysindy.optimizers.trapping_sr3 import _make_constraints
 from pysindy.utils import supports_multiple_targets
 from pysindy.utils.odes import enzyme
 
@@ -1129,3 +1131,18 @@ def test_remove_and_decrement():
         existing_vals=existing_vals, vals_to_remove=vals_to_remove
     )
     np.testing.assert_array_equal(expected, result)
+
+
+def test_trapping_constraints():
+    # x, y, x^2, xy, y^2
+    constraint_rhs, constraint_lhs = _make_constraints(2)
+    stable_coefs = np.array([[0, 0, 0, 1, -1], [0, 0, -1, 1, 0]])
+    result = np.tensordot(constraint_lhs, stable_coefs, ((1, 2), (1, 0)))
+    np.testing.assert_array_equal(constraint_rhs, result)
+
+    # xy, xz, yz
+    stable_coefs = np.array([[0, 0, -1], [0, 0.5, 0], [0.5, 0, 0]])
+    mixed_terms = {frozenset((0, 1)): 0, frozenset((0, 2)): 1, frozenset((1, 2)): 2}
+    constraint_lhs = _antisymm_triple_constraints(3, 3, mixed_terms)
+    result = np.tensordot(constraint_lhs, stable_coefs, ((1, 2), (1, 0)))
+    assert result[0] == 0

From 191b8c45b7a2421a4247c0ae2981f1296a5d8de6 Mon Sep 17 00:00:00 2001
From: Jake <37048747+Jacob-Stevens-Haas@users.noreply.github.com>
Date: Sun, 10 Dec 2023 18:41:00 +0000
Subject: [PATCH 04/16] ENH: Allow all polynomial library library options in
 trapping

(other than degree, which has to be 2)
---
 pysindy/optimizers/trapping_sr3.py | 6 ++++--
 test/test_optimizers.py            | 2 +-
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/pysindy/optimizers/trapping_sr3.py b/pysindy/optimizers/trapping_sr3.py
index def9d71ee..1d8410bcb 100644
--- a/pysindy/optimizers/trapping_sr3.py
+++ b/pysindy/optimizers/trapping_sr3.py
@@ -822,13 +822,15 @@ def _reduce(self, x, y):
         self.objective_history = objective_history
 
 
-def _make_constraints(n_tgts: int):
+def _make_constraints(n_tgts: int, **kwargs):
     """Create constraints for the Quadratic terms in TrappingSR3.
 
     These are the constraints from equation 5 of the Trapping SINDy paper.
 
     Args:
         n_tgts: number of coordinates or modes for which you're fitting an ODE.
+        kwargs: Keyword arguments to PolynomialLibrary such as
+            ``include_bias``.
 
     Returns:
         A tuple of the constraint zeros, and a constraint matrix to multiply
@@ -841,7 +843,7 @@ def _make_constraints(n_tgts: int):
         reshaping.
     """
     n_terms = n_poly_features(n_tgts, degree=2, include_bias=False)
-    lib = PolynomialLibrary(2, include_bias=False).fit(np.zeros((1, n_tgts)))
+    lib = PolynomialLibrary(2, **kwargs).fit(np.zeros((1, n_tgts)))
     terms = [(t_ind, exps) for t_ind, exps in enumerate(lib.powers_)]
 
     # index of tgt -> index of its pure quadratic term
diff --git a/test/test_optimizers.py b/test/test_optimizers.py
index 385d809c7..ac65049e2 100644
--- a/test/test_optimizers.py
+++ b/test/test_optimizers.py
@@ -1135,7 +1135,7 @@ def test_remove_and_decrement():
 
 def test_trapping_constraints():
     # x, y, x^2, xy, y^2
-    constraint_rhs, constraint_lhs = _make_constraints(2)
+    constraint_rhs, constraint_lhs = _make_constraints(2, include_bias=False)
     stable_coefs = np.array([[0, 0, 0, 1, -1], [0, 0, -1, 1, 0]])
     result = np.tensordot(constraint_lhs, stable_coefs, ((1, 2), (1, 0)))
     np.testing.assert_array_equal(constraint_rhs, result)

From 989281c19dcb2be461bcf9dbf36b74cef97a47c7 Mon Sep 17 00:00:00 2001
From: Jake <37048747+Jacob-Stevens-Haas@users.noreply.github.com>
Date: Tue, 2 Jan 2024 16:45:04 +0000
Subject: [PATCH 05/16] CLN: remove unused import

---
 pysindy/optimizers/trapping_sr3.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/pysindy/optimizers/trapping_sr3.py b/pysindy/optimizers/trapping_sr3.py
index 1d8410bcb..e9816481e 100644
--- a/pysindy/optimizers/trapping_sr3.py
+++ b/pysindy/optimizers/trapping_sr3.py
@@ -1,6 +1,5 @@
 import warnings
 from itertools import combinations as combo_nr
-from itertools import combinations_with_replacement as combo_wr
 from math import comb
 from typing import Tuple
 

From 4e99fb41230cbc9ad7d3d92e45e3acc89c89264e Mon Sep 17 00:00:00 2001
From: Jake <37048747+Jacob-Stevens-Haas@users.noreply.github.com>
Date: Mon, 8 Jan 2024 13:19:21 +0000
Subject: [PATCH 06/16] BUG(_make_constraints): Request correct num of terms if
 include_bias

---
 pysindy/optimizers/trapping_sr3.py | 3 ++-
 test/test_optimizers.py            | 9 +++++++--
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/pysindy/optimizers/trapping_sr3.py b/pysindy/optimizers/trapping_sr3.py
index 993b2dc9d..3b250bd79 100644
--- a/pysindy/optimizers/trapping_sr3.py
+++ b/pysindy/optimizers/trapping_sr3.py
@@ -589,7 +589,8 @@ def _make_constraints(n_tgts: int, **kwargs):
         To get "target" order constraints, transpose axis 1 and 2 before
         reshaping.
     """
-    n_terms = n_poly_features(n_tgts, degree=2, include_bias=False)
+    include_bias = kwargs.get("include_bias", False)
+    n_terms = n_poly_features(n_tgts, degree=2, include_bias=include_bias)
     lib = PolynomialLibrary(2, **kwargs).fit(np.zeros((1, n_tgts)))
     terms = [(t_ind, exps) for t_ind, exps in enumerate(lib.powers_)]
 
diff --git a/test/test_optimizers.py b/test/test_optimizers.py
index 5cc6bd4c5..045657aef 100644
--- a/test/test_optimizers.py
+++ b/test/test_optimizers.py
@@ -1131,13 +1131,18 @@ def test_remove_and_decrement():
     np.testing.assert_array_equal(expected, result)
 
 
-def test_trapping_constraints():
+@pytest.mark.parametrize("include_bias", (True, False))
+def test_trapping_constraints(include_bias):
     # x, y, x^2, xy, y^2
-    constraint_rhs, constraint_lhs = _make_constraints(2, include_bias=False)
+    constraint_rhs, constraint_lhs = _make_constraints(2, include_bias=include_bias)
     stable_coefs = np.array([[0, 0, 0, 1, -1], [0, 0, -1, 1, 0]])
+    if include_bias:
+        stable_coefs = np.concatenate(([[0], [0]], stable_coefs), axis=1)
     result = np.tensordot(constraint_lhs, stable_coefs, ((1, 2), (1, 0)))
     np.testing.assert_array_equal(constraint_rhs, result)
 
+
+def test_trapping_mixed_only():
     # xy, xz, yz
     stable_coefs = np.array([[0, 0, -1], [0, 0.5, 0], [0.5, 0, 0]])
     mixed_terms = {frozenset((0, 1)): 0, frozenset((0, 2)): 1, frozenset((1, 2)): 2}

From ff7a2f0919b7e1c23fc52657faf729cee3022979 Mon Sep 17 00:00:00 2001
From: Jake <37048747+Jacob-Stevens-Haas@users.noreply.github.com>
Date: Mon, 8 Jan 2024 16:42:11 +0000
Subject: [PATCH 07/16] TST(trapping): Remove trapping w/cubic library

---
 test/test_optimizers.py | 27 ---------------------------
 1 file changed, 27 deletions(-)

diff --git a/test/test_optimizers.py b/test/test_optimizers.py
index 045657aef..0d7fc5ee1 100644
--- a/test/test_optimizers.py
+++ b/test/test_optimizers.py
@@ -504,33 +504,6 @@ def test_trapping_sr3_quadratic_library(params, trapping_sr3_params, quadratic_l
         assert np.allclose((opt.coef_.flatten())[0], 0.0, atol=1e-5)
 
 
-def test_trapping_cubic_library():
-    x = np.random.standard_normal((10, 3))
-    library_functions = [
-        lambda x: x,
-        lambda x, y: x * y,
-        lambda x: x**2,
-        lambda x, y, z: x * y * z,
-        lambda x, y: x**2 * y,
-        lambda x: x**3,
-    ]
-    library_function_names = [
-        lambda x: str(x),
-        lambda x, y: "{} * {}".format(x, y),
-        lambda x: "{}^2".format(x),
-        lambda x, y, z: "{} * {} * {}".format(x, y, z),
-        lambda x, y: "{}^2 * {}".format(x, y),
-        lambda x: "{}^3".format(x),
-    ]
-    sindy_library = CustomLibrary(
-        library_functions=library_functions, function_names=library_function_names
-    )
-    opt = TrappingSR3()
-    model = SINDy(optimizer=opt, feature_library=sindy_library)
-    model.fit(x)
-    check_is_fitted(model)
-
-
 @pytest.mark.parametrize(
     "error, optimizer, params",
     [

From 40d0b3f307f4ce0f970f67d19eeaaf152448b243 Mon Sep 17 00:00:00 2001
From: Jake <37048747+Jacob-Stevens-Haas@users.noreply.github.com>
Date: Mon, 8 Jan 2024 16:45:08 +0000
Subject: [PATCH 08/16] TST(trapping): remove unused quadratic_library fixutre

---
 test/conftest.py        | 17 -----------------
 test/test_optimizers.py |  2 +-
 2 files changed, 1 insertion(+), 18 deletions(-)

diff --git a/test/conftest.py b/test/conftest.py
index 5ea49e6b8..1b89d1d92 100644
--- a/test/conftest.py
+++ b/test/conftest.py
@@ -315,23 +315,6 @@ def custom_library_bias():
     )
 
 
-@pytest.fixture
-def quadratic_library():
-    library_functions = [
-        lambda x: x,
-        lambda x, y: x * y,
-        lambda x: x**2,
-    ]
-    function_names = [
-        lambda x: str(x),
-        lambda x, y: "{} * {}".format(x, y),
-        lambda x: "{}^2".format(x),
-    ]
-    return CustomLibrary(
-        library_functions=library_functions, function_names=function_names
-    )
-
-
 @pytest.fixture
 def generalized_library():
     tensor_array = [[1, 1]]
diff --git a/test/test_optimizers.py b/test/test_optimizers.py
index 0d7fc5ee1..ef0e9cfd7 100644
--- a/test/test_optimizers.py
+++ b/test/test_optimizers.py
@@ -473,7 +473,7 @@ def test_stable_linear_sr3_linear_library():
         dict(thresholder="weighted_l2", thresholds=1e-5 * np.ones((1, 2))),
     ],
 )
-def test_trapping_sr3_quadratic_library(params, trapping_sr3_params, quadratic_library):
+def test_trapping_sr3_quadratic_library(params, trapping_sr3_params):
     t = np.arange(0, 1, 0.1)
     x = np.exp(-t).reshape((-1, 1))
     x_dot = -x

From 1ea6c7e10447bb4b66a0501420cf8b58822fa7a1 Mon Sep 17 00:00:00 2001
From: Jake <37048747+Jacob-Stevens-Haas@users.noreply.github.com>
Date: Mon, 8 Jan 2024 17:31:52 +0000
Subject: [PATCH 09/16] feat(trapping): Create trapping constraints
 automatically

Previously, TrappingSR3 could only use constraints passed to it, and only then a
limited set of constraints.  It also didn't apply the trapping constraints
automatically, because constraints were required at __init__, and actually
shaping them requires knowledge about the number of features, typically not
known until fit (unless the user is a developer who knows how the feature
libraries work internally :wink:)

WIP, Spawned issue #452
---
 pysindy/optimizers/constrained_sr3.py |  4 +-
 pysindy/optimizers/trapping_sr3.py    | 47 ++++++++++++++++++----
 test/conftest.py                      |  3 +-
 test/test_optimizers.py               | 57 +++++++++++++++++++--------
 4 files changed, 84 insertions(+), 27 deletions(-)

diff --git a/pysindy/optimizers/constrained_sr3.py b/pysindy/optimizers/constrained_sr3.py
index dbc771a6b..0760aba28 100644
--- a/pysindy/optimizers/constrained_sr3.py
+++ b/pysindy/optimizers/constrained_sr3.py
@@ -72,8 +72,8 @@ class ConstrainedSR3(SR3):
 
     constraint_lhs : numpy ndarray, optional (default None)
         Shape should be (n_constraints, n_features * n_targets),
-        The left hand side matrix C of Cw <= d.
-        There should be one row per constraint.
+        The left hand side matrix C of Cw <= d (Or Cw = d for equality
+        constraints). There should be one row per constraint.
 
     constraint_rhs : numpy ndarray, shape (n_constraints,), optional (default None)
         The right hand side vector d of Cw <= d.
diff --git a/pysindy/optimizers/trapping_sr3.py b/pysindy/optimizers/trapping_sr3.py
index 3b250bd79..0183e8b7a 100644
--- a/pysindy/optimizers/trapping_sr3.py
+++ b/pysindy/optimizers/trapping_sr3.py
@@ -71,10 +71,6 @@ class TrappingSR3(ConstrainedSR3):
         If relax_optim = True, use the relax-and-split method. If False,
         try a direct minimization on the largest eigenvalue.
 
-    inequality_constraints : bool, optional (default False)
-        If True, relax_optim must be false or relax_optim = True
-        AND threshold != 0, so that the CVXPY methods are used.
-
     alpha_A :
         Determines the step size in the prox-gradient descent over A.
         For convergence, need alpha_A <= eta, so default
@@ -166,10 +162,11 @@ class TrappingSR3(ConstrainedSR3):
     def __init__(
         self,
         *,
+        _n_tgts: int = None,
+        _include_bias: bool = True,
         eta: Union[float, None] = None,
         eps_solver: float = 1e-7,
         relax_optim: bool = True,
-        inequality_constraints=False,
         alpha_A: Union[float, None] = None,
         alpha_m: Union[float, None] = None,
         gamma: float = -0.1,
@@ -180,10 +177,46 @@ def __init__(
         A0: Union[NDArray, None] = None,
         **kwargs,
     ):
-        super().__init__(thresholder=thresholder, **kwargs)
+        # n_tgts, constraints, etc are data-dependent parameters and belong in
+        # _reduce/fit ().  The following is a hack until we refactor how
+        # constraints are applied in ConstrainedSR3 and MIOSR
+        self._include_bias = _include_bias
+        self._n_tgts = _n_tgts
+        if _n_tgts is None:
+            warnings.warn(
+                "Trapping Optimizer initialized without _n_tgts.  It will likely"
+                " be unable to fit data"
+            )
+            _n_tgts = 1
+        constraint_separation_index = kwargs.get("constraint_separation_index", 0)
+        constraint_rhs, constraint_lhs = _make_constraints(
+            _n_tgts, include_bias=_include_bias
+        )
+        constraint_order = kwargs.get("constraint_order", "feature")
+        if constraint_order == "target":
+            constraint_lhs = np.reshape(np.transpose(constraint_lhs, [1, 0, 2]))
+        constraint_lhs = np.reshape(constraint_lhs, (constraint_lhs.shape[0], -1))
+        try:
+            constraint_lhs = np.concatenate(
+                (kwargs.pop("constraint_lhs"), constraint_lhs), 0
+            )
+            constraint_rhs = np.concatenate(
+                (kwargs.pop("constraint_rhs"), constraint_rhs), 0
+            )
+        except KeyError:
+            pass
+
+        super().__init__(
+            constraint_lhs=constraint_lhs,
+            constraint_rhs=constraint_rhs,
+            constraint_separation_index=constraint_separation_index,
+            constraint_order=constraint_order,
+            equality_constraints=True,
+            thresholder=thresholder,
+            **kwargs,
+        )
         self.eps_solver = eps_solver
         self.relax_optim = relax_optim
-        self.inequality_constraints = inequality_constraints
         self.m0 = m0
         self.A0 = A0
         self.alpha_A = alpha_A
diff --git a/test/conftest.py b/test/conftest.py
index 1b89d1d92..c899b8af2 100644
--- a/test/conftest.py
+++ b/test/conftest.py
@@ -410,7 +410,8 @@ def data_linear_oscillator_corrupted():
 @pytest.fixture(scope="session")
 def data_linear_combination():
     t = np.linspace(0, 5, 100)
-    x = np.stack((np.exp(t), np.sin(t), np.cos(t)), axis=-1)
+    lib = PolynomialLibrary(2)
+    x = lib.fit_transform(t)
     y = np.stack((x[:, 0] + x[:, 1], x[:, 1] + x[:, 2]), axis=-1)
 
     return x, y
diff --git a/test/test_optimizers.py b/test/test_optimizers.py
index ef0e9cfd7..839e05484 100644
--- a/test/test_optimizers.py
+++ b/test/test_optimizers.py
@@ -4,6 +4,7 @@
 import numpy as np
 import pytest
 from numpy.linalg import norm
+from numpy.typing import NDArray
 from scipy.integrate import solve_ivp
 from sklearn.base import BaseEstimator
 from sklearn.exceptions import ConvergenceWarning
@@ -18,6 +19,7 @@
 from pysindy import SINDy
 from pysindy.feature_library import CustomLibrary
 from pysindy.feature_library import SINDyPILibrary
+from pysindy.optimizers import BaseOptimizer
 from pysindy.optimizers import ConstrainedSR3
 from pysindy.optimizers import EnsembleOptimizer
 from pysindy.optimizers import FROLS
@@ -67,6 +69,18 @@ def predict(self, x):
         return x
 
 
+def _align_optimizer_and_1dfeatures(
+    opt: BaseOptimizer, features: NDArray
+) -> tuple[BaseOptimizer, NDArray]:
+    # This is a hack until constraints are moved from init to fit
+    if isinstance(opt, TrappingSR3):
+        opt = TrappingSR3(_n_tgts=1, _include_bias=True)
+        features = np.hstack([features, features, features])
+    else:
+        features = features
+    return opt, features
+
+
 @pytest.mark.parametrize(
     "cls, support",
     [
@@ -100,17 +114,19 @@ def data(request):
         SR3(),
         ConstrainedSR3(),
         StableLinearSR3(),
-        TrappingSR3(),
+        TrappingSR3(_n_tgts=1),
         Lasso(fit_intercept=False),
         ElasticNet(fit_intercept=False),
         DummyLinearModel(),
         MIOSR(),
     ],
+    ids=lambda param: type(param),
 )
 def test_fit(data_derivative_1d, optimizer):
     x, x_dot = data_derivative_1d
     if len(x.shape) == 1:
         x = x.reshape(-1, 1)
+    optimizer, x = _align_optimizer_and_1dfeatures(optimizer, x)
     opt = WrappedOptimizer(optimizer, unbias=False)
     opt.fit(x, x_dot)
 
@@ -167,12 +183,12 @@ def test_alternate_parameters(data_derivative_1d, kwargs):
     ],
 )
 def test_sample_weight_optimizers(data_1d, optimizer):
-    x, t = data_1d
-
+    y, t = data_1d
+    opt = optimizer()
+    opt, x = _align_optimizer_and_1dfeatures(opt, y)
     sample_weight = np.ones(x[:, 0].shape)
     sample_weight[::2] = 0
-    opt = optimizer()
-    opt.fit(x, x, sample_weight=sample_weight)
+    opt.fit(x, y, sample_weight=sample_weight)
     check_is_fitted(opt)
 
 
@@ -222,12 +238,12 @@ def test_sr3_bad_parameters(optimizer, params):
 )
 def test_trapping_bad_parameters(params):
     with pytest.raises(ValueError):
-        TrappingSR3(**params)
+        TrappingSR3(_n_tgts=1, **params)
 
 
 def test_trapping_objective_print():
     # test error in verbose print logic when max_iter < 10
-    opt = TrappingSR3(max_iter=2, verbose=True)
+    opt = TrappingSR3(_n_tgts=1, max_iter=2, verbose=True)
     arr = np.ones(1)
     opt._objective(arr, arr, arr, arr, arr, 1)
 
@@ -481,7 +497,7 @@ def test_trapping_sr3_quadratic_library(params, trapping_sr3_params):
 
     params.update(trapping_sr3_params)
 
-    opt = TrappingSR3(**params)
+    opt = TrappingSR3(_n_tgts=1, _include_bias=False, **params)
     opt.fit(features, x_dot)
     assert opt.PL_.shape == (1, 1, 1, 2)
     assert opt.PQ_.shape == (1, 1, 1, 1, 2)
@@ -494,7 +510,7 @@ def test_trapping_sr3_quadratic_library(params, trapping_sr3_params):
     params["constraint_rhs"] = np.zeros(p)
     params["constraint_lhs"] = np.eye(p, r * N)
 
-    opt = TrappingSR3(**params)
+    opt = TrappingSR3(_n_tgts=1, _include_bias=False, **params)
     opt.fit(features, x_dot)
     assert opt.PL_.shape == (1, 1, 1, 2)
     assert opt.PQ_.shape == (1, 1, 1, 1, 2)
@@ -631,7 +647,7 @@ def test_constrained_sr3_prox_functions(data_derivative_1d, thresholder):
         (SR3, {"trimming_fraction": 0.1}),
         (ConstrainedSR3, {"constraint_lhs": [1], "constraint_rhs": [1]}),
         (ConstrainedSR3, {"trimming_fraction": 0.1}),
-        (TrappingSR3, {"constraint_lhs": [1], "constraint_rhs": [1]}),
+        (TrappingSR3, {"_n_tgts": 1, "constraint_lhs": [1], "constraint_rhs": [1]}),
         (StableLinearSR3, {"constraint_lhs": [1], "constraint_rhs": [1]}),
         (StableLinearSR3, {"trimming_fraction": 0.1}),
         (SINDyPI, {}),
@@ -741,7 +757,7 @@ def test_sr3_enable_trimming(optimizer, data_linear_oscillator_corrupted):
         SR3(max_iter=1),
         ConstrainedSR3(max_iter=1),
         StableLinearSR3(max_iter=1),
-        TrappingSR3(max_iter=1),
+        TrappingSR3(_n_tgts=1, max_iter=1),
     ],
 )
 def test_fit_warn(data_derivative_1d, optimizer):
@@ -755,7 +771,11 @@ def test_fit_warn(data_derivative_1d, optimizer):
 
 @pytest.mark.parametrize(
     "optimizer",
-    [(ConstrainedSR3, {"max_iter": 80}), (TrappingSR3, {"max_iter": 100}), (MIOSR, {})],
+    [
+        (ConstrainedSR3, {"max_iter": 80}),
+        (TrappingSR3, {"_n_tgts": 5, "max_iter": 100}),
+        (MIOSR, {}),
+    ],
 )
 @pytest.mark.parametrize("target_value", [0, -1, 3])
 def test_row_format_constraints(data_linear_combination, optimizer, target_value):
@@ -966,6 +986,7 @@ def test_normalize_columns(data_derivative_1d, optimizer):
     if len(x.shape) == 1:
         x = x.reshape(-1, 1)
     opt = optimizer(normalize_columns=True)
+    opt, x = _align_optimizer_and_1dfeatures(opt, x)
     opt.fit(x, x_dot)
     check_is_fitted(opt)
     assert opt.complexity >= 0
@@ -1027,9 +1048,11 @@ def test_ssr_criteria(data_lorenz):
     ],
 )
 def test_optimizers_verbose(data_1d, optimizer):
-    x, _ = data_1d
+    y, _ = data_1d
     opt = optimizer(verbose=True)
-    opt.fit(x, x)
+    opt, x = _align_optimizer_and_1dfeatures(opt, y)
+    opt.verbose = True
+    opt.fit(x, y)
     check_is_fitted(opt)
 
 
@@ -1043,10 +1066,10 @@ def test_optimizers_verbose(data_1d, optimizer):
     ],
 )
 def test_optimizers_verbose_cvxpy(data_1d, optimizer):
-    x, _ = data_1d
-
+    y, _ = data_1d
     opt = optimizer(verbose_cvxpy=True)
-    opt.fit(x, x)
+    opt, x = _align_optimizer_and_1dfeatures(opt, y)
+    opt.fit(x, y)
     check_is_fitted(opt)
 
 

From 89aeef3d363bbfc4b46c9f075e369dbbee7a12fe Mon Sep 17 00:00:00 2001
From: Jake <37048747+Jacob-Stevens-Haas@users.noreply.github.com>
Date: Mon, 8 Jan 2024 20:57:46 +0000
Subject: [PATCH 10/16] BUG(axes): Make concat out param work

---
 pysindy/utils/axes.py   | 7 +++++--
 test/utils/test_axes.py | 7 +++++++
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/pysindy/utils/axes.py b/pysindy/utils/axes.py
index bad10d55c..ad0f79040 100644
--- a/pysindy/utils/axes.py
+++ b/pysindy/utils/axes.py
@@ -117,13 +117,16 @@ def decorator(func):
 
 
 @implements(np.concatenate)
-def concatenate(arrays, axis=0):
+def concatenate(arrays, axis=0, out=None, dtype=None, casting="same_kind"):
     parents = [np.asarray(obj) for obj in arrays]
     ax_list = [obj.__dict__ for obj in arrays if isinstance(obj, AxesArray)]
     for ax1, ax2 in zip(ax_list[:-1], ax_list[1:]):
         if ax1 != ax2:
             raise TypeError("Concatenating >1 AxesArray with incompatible axes")
-    return AxesArray(np.concatenate(parents, axis), axes=ax_list[0])
+    result = np.concatenate(parents, axis, out=out, dtype=dtype, casting=casting)
+    if isinstance(out, AxesArray):
+        out.__dict__ = ax_list[0]
+    return AxesArray(result, axes=ax_list[0])
 
 
 def comprehend_axes(x):
diff --git a/test/utils/test_axes.py b/test/utils/test_axes.py
index b1a38e6f4..e5d9a8385 100644
--- a/test/utils/test_axes.py
+++ b/test/utils/test_axes.py
@@ -7,6 +7,13 @@
 from pysindy import AxesArray
 
 
+def test_concat_out():
+    arr = AxesArray(np.arange(3).reshape(1, 3), {"ax_a": 0, "ax_b": 1})
+    arr_out = np.empty((2, 3)).view(AxesArray)
+    result = np.concatenate((arr, arr), axis=0, out=arr_out)
+    assert_equal(result, arr_out)
+
+
 def test_reduce_mean_noinf_recursion():
     arr = AxesArray(np.array([[1]]), {})
     np.mean(arr, axis=0)

From 4feea2327b94fef8acd7fb285edf6bf142f691b4 Mon Sep 17 00:00:00 2001
From: Jake <37048747+Jacob-Stevens-Haas@users.noreply.github.com>
Date: Tue, 9 Jan 2024 16:13:07 +0000
Subject: [PATCH 11/16] TST: Fix shape of linear_combination (1d->2d)

---
 test/conftest.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/test/conftest.py b/test/conftest.py
index c899b8af2..8af932c7a 100644
--- a/test/conftest.py
+++ b/test/conftest.py
@@ -409,10 +409,10 @@ def data_linear_oscillator_corrupted():
 
 @pytest.fixture(scope="session")
 def data_linear_combination():
-    t = np.linspace(0, 5, 100)
-    lib = PolynomialLibrary(2)
-    x = lib.fit_transform(t)
-    y = np.stack((x[:, 0] + x[:, 1], x[:, 1] + x[:, 2]), axis=-1)
+    t = np.linspace(0, 3, 100).reshape((-1, 1))
+    lib = PolynomialLibrary(2, include_bias=False)
+    x = lib.fit_transform(np.hstack([t, -2 * t, 3 * t]))
+    y = np.stack((x[:, 0] + x[:, 1], x[:, 1] + x[:, 2], x[:, 0] + x[:, 2]), axis=-1)
 
     return x, y
 

From c7bb2a2aeb5a1ad74d503657db2de5e1bb4027c5 Mon Sep 17 00:00:00 2001
From: Jake <37048747+Jacob-Stevens-Haas@users.noreply.github.com>
Date: Wed, 10 Jan 2024 01:36:40 +0000
Subject: [PATCH 12/16] feat(utils): Deprecate target format constraints

And simplify reorder_constraints
---
 pysindy/utils/base.py | 26 ++++++++++----------------
 1 file changed, 10 insertions(+), 16 deletions(-)

diff --git a/pysindy/utils/base.py b/pysindy/utils/base.py
index cc3a846ab..114916e7f 100644
--- a/pysindy/utils/base.py
+++ b/pysindy/utils/base.py
@@ -1,3 +1,4 @@
+import warnings
 from itertools import repeat
 from typing import Sequence
 
@@ -136,24 +137,17 @@ def drop_nan_samples(x, y):
     return x, y
 
 
-def reorder_constraints(c, n_features, output_order="row"):
-    """Reorder constraint matrix."""
-    ret = c.copy()
-
-    if ret.ndim == 1:
-        ret = ret.reshape(1, -1)
-
-    n_targets = ret.shape[1] // n_features
-    shape = (n_targets, n_features)
-
-    if output_order == "row":
-        for i in range(ret.shape[0]):
-            ret[i] = ret[i].reshape(shape).flatten(order="F")
+def reorder_constraints(arr, n_features, output_order="feature"):
+    """Switch between 'feature' and 'target' constraint order."""
+    warnings.warn("Target format constraints are deprecated.", stacklevel=2)
+    n_constraints = arr.shape[0] if arr.ndim > 1 else 1
+    n_tgt = arr.size // n_features // n_constraints
+    if output_order == "feature":
+        starting_shape = (n_constraints, n_tgt, n_features)
     else:
-        for i in range(ret.shape[0]):
-            ret[i] = ret[i].reshape(shape, order="F").flatten()
+        starting_shape = (n_constraints, n_features, n_tgt)
 
-    return ret
+    return arr.reshape(starting_shape).transpose([0, 2, 1]).reshape((n_constraints, -1))
 
 
 def prox_l0(x, threshold):

From 5599abf3ea6a6e4e2541433b76a25dd6b349dd0a Mon Sep 17 00:00:00 2001
From: Jake <37048747+Jacob-Stevens-Haas@users.noreply.github.com>
Date: Wed, 10 Jan 2024 03:30:36 +0000
Subject: [PATCH 13/16] test(reformat_constraints): Make target and feature
 format clear

Read the fucking diff

backport-to: constraint-order
---
 test/utils/test_utils.py | 44 +++++++++++++++++++++++++++-------------
 1 file changed, 30 insertions(+), 14 deletions(-)

diff --git a/test/utils/test_utils.py b/test/utils/test_utils.py
index 54578c43d..4c4d2daf2 100644
--- a/test/utils/test_utils.py
+++ b/test/utils/test_utils.py
@@ -7,28 +7,44 @@
 
 
 def test_reorder_constraints_1D():
-    target_order = np.arange(6)
-    row_order = np.array([0, 3, 1, 4, 2, 5])
     n_feats = 3
-
-    np.testing.assert_array_equal(
-        reorder_constraints(target_order, n_feats).flatten(), row_order
+    n_tgts = 2
+    target_order = np.array(
+        [f"t{i}f{j}" for i in range(n_tgts) for j in range(n_feats)]
     )
-    np.testing.assert_array_equal(
-        reorder_constraints(row_order, n_feats, output_order="target").flatten(),
-        target_order,
+    feature_order = np.array(
+        [f"t{i}f{j}" for j in range(n_feats) for i in range(n_tgts)]
     )
 
+    result = reorder_constraints(target_order, n_feats, output_order="feature")
+    np.testing.assert_array_equal(result.flatten(), feature_order)
+
+    result = reorder_constraints(feature_order, n_feats, output_order="target")
+    np.testing.assert_array_equal(result.flatten(), target_order)
+
 
 def test_reorder_constraints_2D():
-    target_order = np.arange(12).reshape((2, 6))
-    row_order = np.array([[0, 3, 1, 4, 2, 5], [6, 9, 7, 10, 8, 11]])
     n_feats = 3
-
-    np.testing.assert_array_equal(reorder_constraints(target_order, n_feats), row_order)
-    np.testing.assert_array_equal(
-        reorder_constraints(row_order, n_feats, output_order="target"), target_order
+    n_tgts = 2
+    n_const = 2
+    target_order = np.array(
+        [
+            [f"c{k}t{i}f{j}" for i in range(n_tgts) for j in range(n_feats)]
+            for k in range(n_const)
+        ]
     )
+    feature_order = np.array(
+        [
+            [f"c{k}t{i}f{j}" for j in range(n_feats) for i in range(n_tgts)]
+            for k in range(n_const)
+        ]
+    )
+
+    result = reorder_constraints(target_order, n_feats, output_order="feature")
+    np.testing.assert_array_equal(result, feature_order)
+
+    result = reorder_constraints(feature_order, n_feats, output_order="target")
+    np.testing.assert_array_equal(result, target_order)
 
 
 def test_validate_controls():

From 85d3e6d312aa2361ccfd09e7ef782d14d0dc5c2f Mon Sep 17 00:00:00 2001
From: Jake <37048747+Jacob-Stevens-Haas@users.noreply.github.com>
Date: Wed, 10 Jan 2024 03:41:40 +0000
Subject: [PATCH 14/16] feat(trapping): Enable merging trapping and user
 constraints

Also:
* (constrained_sr3)clean up creating cvxpy constraints
* (constrained_sr3)Clone the cvxpy problem in case of except statements,
because prob.solve() mutates prob in a mathematically significant way
---
 pysindy/optimizers/constrained_sr3.py | 63 +++++++++++++++------------
 pysindy/optimizers/trapping_sr3.py    | 33 ++++++++++----
 test/test_optimizers.py               | 54 +++++++++++++++--------
 3 files changed, 95 insertions(+), 55 deletions(-)

diff --git a/pysindy/optimizers/constrained_sr3.py b/pysindy/optimizers/constrained_sr3.py
index 0760aba28..300bce897 100644
--- a/pysindy/optimizers/constrained_sr3.py
+++ b/pysindy/optimizers/constrained_sr3.py
@@ -1,4 +1,6 @@
 import warnings
+from copy import deepcopy
+from typing import Optional
 from typing import Tuple
 
 try:
@@ -168,7 +170,7 @@ def __init__(
         thresholds=None,
         equality_constraints=False,
         inequality_constraints=False,
-        constraint_separation_index=0,
+        constraint_separation_index: Optional[bool] = None,
         verbose=False,
         verbose_cvxpy=False,
         unbias=False,
@@ -194,7 +196,7 @@ def __init__(
         self.constraint_lhs = constraint_lhs
         self.constraint_rhs = constraint_rhs
         self.constraint_order = constraint_order
-        self.use_constraints = (constraint_lhs is not None) and (
+        self.use_constraints = (constraint_lhs is not None) or (
             constraint_rhs is not None
         )
 
@@ -208,7 +210,7 @@ def __init__(
                 " but user did not specify if the constraints were equality or"
                 " inequality constraints. Assuming equality constraints."
             )
-            self.equality_constraints = True
+            equality_constraints = True
 
         if self.use_constraints:
             if constraint_order not in ("feature", "target"):
@@ -243,6 +245,16 @@ def __init__(
             )
         self.inequality_constraints = inequality_constraints
         self.equality_constraints = equality_constraints
+        if self.use_constraints and constraint_separation_index is None:
+            if self.inequality_constraints and not self.equality_constraints:
+                constraint_separation_index = len(constraint_lhs)
+            elif self.equality_constraints and not self.inequality_constraints:
+                constraint_separation_index = 0
+            else:
+                raise ValueError(
+                    "If passing both inequality and equality constraints, must specify"
+                    " constraint_separation_index."
+                )
         self.constraint_separation_index = constraint_separation_index
 
     def _update_full_coef_constraints(self, H, x_transpose_y, coef_sparse):
@@ -276,30 +288,20 @@ def _create_var_and_part_cost(
 
     def _update_coef_cvxpy(self, xi, cost, var_len, coef_prev, tol):
         if self.use_constraints:
-            if self.inequality_constraints and self.equality_constraints:
-                # Process inequality constraints then equality constraints
-                prob = cp.Problem(
-                    cp.Minimize(cost),
-                    [
-                        self.constraint_lhs[: self.constraint_separation_index, :] @ xi
-                        <= self.constraint_rhs[: self.constraint_separation_index],
-                        self.constraint_lhs[self.constraint_separation_index :, :] @ xi
-                        == self.constraint_rhs[self.constraint_separation_index :],
-                    ],
+            constraints = []
+            if self.equality_constraints:
+                constraints.append(
+                    self.constraint_lhs[self.constraint_separation_index :, :] @ xi
+                    == self.constraint_rhs[self.constraint_separation_index :],
                 )
-            elif self.inequality_constraints:
-                prob = cp.Problem(
-                    cp.Minimize(cost),
-                    [self.constraint_lhs @ xi <= self.constraint_rhs],
-                )
-            else:
-                prob = cp.Problem(
-                    cp.Minimize(cost),
-                    [self.constraint_lhs @ xi == self.constraint_rhs],
+            if self.inequality_constraints:
+                constraints.append(
+                    self.constraint_lhs[: self.constraint_separation_index, :] @ xi
+                    <= self.constraint_rhs[: self.constraint_separation_index]
                 )
-        else:
-            prob = cp.Problem(cp.Minimize(cost))
+            prob = cp.Problem(cp.Minimize(cost), constraints)
 
+        prob_clone = deepcopy(prob)
         # default solver is SCS/OSQP here but switches to ECOS for L2
         try:
             prob.solve(
@@ -313,13 +315,20 @@ def _update_coef_cvxpy(self, xi, cost, var_len, coef_prev, tol):
         # similar semantic changes for the other variables.
         except (TypeError, ValueError):
             try:
+                prob = prob_clone
                 prob.solve(max_iters=self.max_iter, verbose=self.verbose_cvxpy)
+                xi = prob.variables()[0]
             except cp.error.SolverError:
-                print("Solver failed, setting coefs to zeros")
+                warnings.warn("Solver failed, setting coefs to zeros")
                 xi.value = np.zeros(var_len)
         except cp.error.SolverError:
-            print("Solver failed, setting coefs to zeros")
-            xi.value = np.zeros(var_len)
+            try:
+                prob = prob_clone
+                prob.solve(max_iter=self.max_iter, verbose=self.verbose_cvxpy)
+                xi = prob.variables()[0]
+            except cp.error.SolverError:
+                warnings.warn("Solver failed, setting coefs to zeros")
+                xi.value = np.zeros(var_len)
 
         if xi.value is None:
             warnings.warn(
diff --git a/pysindy/optimizers/trapping_sr3.py b/pysindy/optimizers/trapping_sr3.py
index 0183e8b7a..415ce88c1 100644
--- a/pysindy/optimizers/trapping_sr3.py
+++ b/pysindy/optimizers/trapping_sr3.py
@@ -163,7 +163,8 @@ def __init__(
         self,
         *,
         _n_tgts: int = None,
-        _include_bias: bool = True,
+        _include_bias: bool = False,
+        _interaction_only: bool = False,
         eta: Union[float, None] = None,
         eps_solver: float = 1e-7,
         relax_optim: bool = True,
@@ -181,6 +182,7 @@ def __init__(
         # _reduce/fit ().  The following is a hack until we refactor how
         # constraints are applied in ConstrainedSR3 and MIOSR
         self._include_bias = _include_bias
+        self._interaction_only = _interaction_only
         self._n_tgts = _n_tgts
         if _n_tgts is None:
             warnings.warn(
@@ -188,13 +190,22 @@ def __init__(
                 " be unable to fit data"
             )
             _n_tgts = 1
-        constraint_separation_index = kwargs.get("constraint_separation_index", 0)
+        if _include_bias:
+            raise ValueError(
+                "Currently not able to include bias until PQ matrices are modified"
+            )
+        if hasattr(kwargs, "constraint_separation_index"):
+            constraint_separation_index = kwargs["constraint_separation_index"]
+        elif kwargs.get("inequality_constraints", False):
+            constraint_separation_index = kwargs["constraint_lhs"].shape[0]
+        else:
+            constraint_separation_index = 0
         constraint_rhs, constraint_lhs = _make_constraints(
             _n_tgts, include_bias=_include_bias
         )
-        constraint_order = kwargs.get("constraint_order", "feature")
+        constraint_order = kwargs.pop("constraint_order", "feature")
         if constraint_order == "target":
-            constraint_lhs = np.reshape(np.transpose(constraint_lhs, [1, 0, 2]))
+            constraint_lhs = np.transpose(constraint_lhs, [0, 2, 1])
         constraint_lhs = np.reshape(constraint_lhs, (constraint_lhs.shape[0], -1))
         try:
             constraint_lhs = np.concatenate(
@@ -460,22 +471,26 @@ def _reduce(self, x, y):
         self.PW_history_ = []
         self.PWeigs_history_ = []
         self.history_ = []
-        n_samples, n_features = x.shape
-        n_tgts = y.shape[1]
+        n_samples, n_tgts = y.shape
+        n_features = n_poly_features(
+            n_tgts,
+            2,
+            include_bias=self._include_bias,
+            interaction_only=self._interaction_only,
+        )
         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.
         self.PL_unsym_, self.PL_, self.PQ_ = self._set_Ptensors(n_tgts)
         # make sure dimensions/symmetries are correct
         self.PL_, self.PQ_ = self._check_P_matrix(
-            n_tgts, n_features, n_feat_expected, self.PL_, self.PQ_
+            n_tgts, n_features, n_features, self.PL_, self.PQ_
         )
 
         # Set initial coefficients
         if self.use_constraints and self.constraint_order.lower() == "target":
             self.constraint_lhs = reorder_constraints(
-                self.constraint_lhs, n_features, output_order="target"
+                self.constraint_lhs, n_features, output_order="feature"
             )
         coef_sparse = self.coef_.T
 
diff --git a/test/test_optimizers.py b/test/test_optimizers.py
index 839e05484..270e51d48 100644
--- a/test/test_optimizers.py
+++ b/test/test_optimizers.py
@@ -74,8 +74,8 @@ def _align_optimizer_and_1dfeatures(
 ) -> tuple[BaseOptimizer, NDArray]:
     # This is a hack until constraints are moved from init to fit
     if isinstance(opt, TrappingSR3):
-        opt = TrappingSR3(_n_tgts=1, _include_bias=True)
-        features = np.hstack([features, features, features])
+        opt = TrappingSR3(_n_tgts=1, _include_bias=False)
+        features = np.hstack([features, features])
     else:
         features = features
     return opt, features
@@ -435,6 +435,7 @@ def test_constrained_sr3_quadratic_library(params):
         dict(thresholder="l2", threshold=1, expected=1.5),
         dict(thresholder="weighted_l2", thresholds=np.ones((4, 1)), expected=2.5),
     ],
+    ids=lambda d: d["thresholder"],
 )
 def test_stable_linear_sr3_cost_function(params):
     expected = params.pop("expected")
@@ -773,21 +774,24 @@ def test_fit_warn(data_derivative_1d, optimizer):
     "optimizer",
     [
         (ConstrainedSR3, {"max_iter": 80}),
-        (TrappingSR3, {"_n_tgts": 5, "max_iter": 100}),
+        (TrappingSR3, {"_n_tgts": 3, "max_iter": 100, "eps_solver": 1e-5}),
         (MIOSR, {}),
     ],
+    ids=lambda param: param[0].__name__ + " " + ",".join([key for key in param[1]]),
 )
 @pytest.mark.parametrize("target_value", [0, -1, 3])
-def test_row_format_constraints(data_linear_combination, optimizer, target_value):
+def test_feature_format_constraints(data_linear_combination, optimizer, target_value):
     # Solution is x_dot = x.dot(np.array([[1, 1, 0], [0, 1, 1]]))
-    x, x_dot = data_linear_combination
+    x, y = data_linear_combination
 
     constraint_rhs = target_value * np.ones(2)
-    constraint_lhs = np.zeros((2, x.shape[1] * x_dot.shape[1]))
+    constraint_lhs = np.zeros((2, x.shape[1], y.shape[1]))
 
     # Should force corresponding entries of coef_ to be target_value
-    constraint_lhs[0, 0] = 1
-    constraint_lhs[1, 3] = 1
+    constraint_lhs[0, 1, 1] = 1
+    constraint_lhs[1, 2, 2] = 1
+    # reshape to "feature" order
+    constraint_lhs = np.reshape(constraint_lhs, (constraint_lhs.shape[0], -1))
 
     model = optimizer[0](
         constraint_lhs=constraint_lhs,
@@ -795,10 +799,10 @@ def test_row_format_constraints(data_linear_combination, optimizer, target_value
         constraint_order="feature",
         **optimizer[1],
     )
-    model.fit(x, x_dot)
+    model.fit(x, y)
 
     np.testing.assert_allclose(
-        np.array([model.coef_[0, 0], model.coef_[1, 1]]), target_value, atol=1e-8
+        np.array([model.coef_[1, 1], model.coef_[2, 2]]), target_value, atol=1e-7
     )
 
 
@@ -807,26 +811,37 @@ def test_row_format_constraints(data_linear_combination, optimizer, target_value
     [
         (ConstrainedSR3, {"max_iter": 80}),
         (StableLinearSR3, {}),
-        (TrappingSR3, {"max_iter": 100}),
+        (TrappingSR3, {"_n_tgts": 3, "max_iter": 200, "eps_solver": 1e-5}),
         (MIOSR, {}),
     ],
+    ids=lambda param: param[0].__name__ + " " + ",".join([key for key in param[1]]),
 )
 @pytest.mark.parametrize("target_value", [0, -1, 3])
 def test_target_format_constraints(data_linear_combination, optimizer, target_value):
-    x, x_dot = data_linear_combination
+    x, y = data_linear_combination
 
     constraint_rhs = target_value * np.ones(2)
-    constraint_lhs = np.zeros((2, x.shape[1] * x_dot.shape[1]))
+    constraint_lhs = np.zeros((2, x.shape[1], y.shape[1]))
 
     # Should force corresponding entries of coef_ to be target_value
-    constraint_lhs[0, 1] = 1
-    constraint_lhs[1, 4] = 1
+    constraint_lhs[0, 2, 1] = 1
+    constraint_lhs[1, 1, 2] = 1
+    # reshape to "target" order
+    constraint_lhs = np.reshape(
+        np.transpose(constraint_lhs, [0, 2, 1]), (constraint_lhs.shape[0], -1)
+    )
 
     model = optimizer[0](
-        constraint_lhs=constraint_lhs, constraint_rhs=constraint_rhs, **optimizer[1]
+        constraint_lhs=constraint_lhs,
+        constraint_rhs=constraint_rhs,
+        constraint_order="target",
+        **optimizer[1],
+    )
+    model.fit(x, y)
+
+    np.testing.assert_allclose(
+        np.array([model.coef_[1, 2], model.coef_[2, 1]]), target_value, atol=1e-7
     )
-    model.fit(x, x_dot)
-    np.testing.assert_allclose(model.coef_[:, 1], target_value, atol=1e-8)
 
 
 @pytest.mark.parametrize(
@@ -876,10 +891,11 @@ def test_constrained_inequality_constraints(data_lorenz, params):
             thresholder="weighted_l2", thresholds=0.5 * np.ones((1, 2)), expected=0.75
         ),
     ],
+    ids=lambda d: d["thresholder"],
 )
 def test_trapping_cost_function(params):
     expected = params.pop("expected")
-    opt = TrappingSR3(inequality_constraints=True, relax_optim=True, **params)
+    opt = TrappingSR3(relax_optim=True, **params)
     x = np.eye(2)
     y = np.ones(2)
     xi, cost = opt._create_var_and_part_cost(2, x, y)

From 90e0f6db844947330f9b0c61dcea691a399797fb Mon Sep 17 00:00:00 2001
From: Jake <37048747+Jacob-Stevens-Haas@users.noreply.github.com>
Date: Tue, 16 Apr 2024 13:08:24 -0700
Subject: [PATCH 15/16] CLN: Fix black/flake8 minor

---
 examples/16_noise_robustness/utils.py                | 1 -
 examples/17_parameterized_pattern_formation/utils.py | 8 ++++----
 pysindy/differentiation/finite_difference.py         | 1 -
 pysindy/feature_library/polynomial_library.py        | 3 ++-
 pysindy/feature_library/sindy_pi_library.py          | 1 -
 pysindy/feature_library/weak_pde_library.py          | 3 ---
 pysindy/optimizers/sbr.py                            | 1 -
 pysindy/optimizers/trapping_sr3.py                   | 1 -
 pysindy/utils/axes.py                                | 2 --
 pysindy/utils/odes.py                                | 1 +
 test/conftest.py                                     | 7 -------
 test/utils/test_axes.py                              | 7 -------
 12 files changed, 7 insertions(+), 29 deletions(-)

diff --git a/examples/16_noise_robustness/utils.py b/examples/16_noise_robustness/utils.py
index f81a89fbd..8b51b52c5 100644
--- a/examples/16_noise_robustness/utils.py
+++ b/examples/16_noise_robustness/utils.py
@@ -169,7 +169,6 @@ def make_test_trajectories(
     all_t_test = dict()
 
     for i, equation_name in enumerate(systems_list):
-
         dimension = all_properties[equation_name]["embedding_dimension"]
         all_sols_test[equation_name] = np.zeros((n, n_trajectories, dimension))
         all_t_test[equation_name] = np.zeros((n, n_trajectories))
diff --git a/examples/17_parameterized_pattern_formation/utils.py b/examples/17_parameterized_pattern_formation/utils.py
index 91edcae6c..c25778c14 100644
--- a/examples/17_parameterized_pattern_formation/utils.py
+++ b/examples/17_parameterized_pattern_formation/utils.py
@@ -38,7 +38,7 @@ def get_lorenz_trajectories(sigmas, rhos, betas, dt):
             x0_train,
             args=(sigmas[i], betas[i], rhos[i]),
             t_eval=t_train,
-            **integrator_keywords
+            **integrator_keywords,
         ).y.T
         x_trains = x_trains + [x_train]
         t_trains = t_trains + [t_train]
@@ -966,7 +966,7 @@ def oregonator(
                 args=(b,),
                 y0=ut,
                 first_step=dt,
-                **integrator_keywords
+                **integrator_keywords,
             )
             if not usol.success:
                 print(usol.message)
@@ -1164,7 +1164,7 @@ def oregonator_homogeneous(
             args=(b,),
             y0=ut,
             first_step=dt,
-            **integrator_keywords
+            **integrator_keywords,
         )
         dt = np.diff(usol.t)[-1]
         ut = usol.y[:, -1]
@@ -1267,7 +1267,7 @@ def cgle_homogeneous(t, u):
             (t[i], t[i + 1]),
             y0=ut,
             first_step=dt,
-            **integrator_keywords
+            **integrator_keywords,
         )
         #         print(usol.message,end='\r')
         dt = np.diff(usol.t)[-1]
diff --git a/pysindy/differentiation/finite_difference.py b/pysindy/differentiation/finite_difference.py
index 0c44917e9..25690ddcc 100644
--- a/pysindy/differentiation/finite_difference.py
+++ b/pysindy/differentiation/finite_difference.py
@@ -66,7 +66,6 @@ def __init__(
         drop_endpoints=False,
         periodic=False,
     ):
-
         if order <= 0 or not isinstance(order, int):
             raise ValueError("order must be a positive int")
         if d <= 0:
diff --git a/pysindy/feature_library/polynomial_library.py b/pysindy/feature_library/polynomial_library.py
index 817e6a963..21cc126a9 100644
--- a/pysindy/feature_library/polynomial_library.py
+++ b/pysindy/feature_library/polynomial_library.py
@@ -1,6 +1,7 @@
 from itertools import chain
 from math import comb
 from typing import Iterator
+from typing import Tuple
 
 import numpy as np
 from numpy.typing import NDArray
@@ -80,7 +81,7 @@ def _combinations(
         include_interaction: bool,
         interaction_only: bool,
         include_bias: bool,
-    ) -> Iterator[tuple[int, ...]]:
+    ) -> Iterator[Tuple[int, ...]]:
         """
         Create selection tuples of input indexes for each polynomail term
 
diff --git a/pysindy/feature_library/sindy_pi_library.py b/pysindy/feature_library/sindy_pi_library.py
index f45cf567f..407c67de1 100644
--- a/pysindy/feature_library/sindy_pi_library.py
+++ b/pysindy/feature_library/sindy_pi_library.py
@@ -393,7 +393,6 @@ def transform(self, x_full):
                         f_dot.__code__.co_argcount,
                         self.interaction_only,
                     ):
-
                         for i, f in enumerate(self.x_functions):
                             for f_combs in self._combinations(
                                 n_input_features,
diff --git a/pysindy/feature_library/weak_pde_library.py b/pysindy/feature_library/weak_pde_library.py
index 0f41ede84..edbdce40b 100644
--- a/pysindy/feature_library/weak_pde_library.py
+++ b/pysindy/feature_library/weak_pde_library.py
@@ -457,7 +457,6 @@ def _set_up_weights(self):
         deriv = np.zeros(self.grid_ndim)
         deriv[-1] = 1
         for k in range(self.K):
-
             ret = np.ones(shapes_k[k])
             for i in range(self.grid_ndim):
                 s = [0] * (self.grid_ndim + 1)
@@ -476,7 +475,6 @@ def _set_up_weights(self):
         # Product weights over the axes for pure derivative terms, shaped as inds_k
         self.fullweights0 = []
         for k in range(self.K):
-
             ret = np.ones(shapes_k[k])
             for i in range(self.grid_ndim):
                 s = [0] * (self.grid_ndim + 1)
@@ -922,7 +920,6 @@ def transform(self, x_full):
                         for deriv in derivs[
                             np.where(np.all(derivs <= derivs_mixed, axis=1))[0]
                         ]:
-
                             # indices for product rule terms
                             j0 = np.where(np.all(derivs == deriv, axis=1))[0][0]
                             j1 = np.where(
diff --git a/pysindy/optimizers/sbr.py b/pysindy/optimizers/sbr.py
index f06467b4c..2e304a34c 100644
--- a/pysindy/optimizers/sbr.py
+++ b/pysindy/optimizers/sbr.py
@@ -99,7 +99,6 @@ def __init__(
         unbias: bool = False,
         **kwargs,
     ):
-
         if unbias:
             raise ValueError("SBR is incompatible with unbiasing. Set unbias=False")
 
diff --git a/pysindy/optimizers/trapping_sr3.py b/pysindy/optimizers/trapping_sr3.py
index 415ce88c1..973267a26 100644
--- a/pysindy/optimizers/trapping_sr3.py
+++ b/pysindy/optimizers/trapping_sr3.py
@@ -539,7 +539,6 @@ def _reduce(self, x, y):
         # Begin optimization loop
         self.objective_history_ = []
         for k in range(self.max_iter):
-
             # update P tensor from the newest trap center
             mPQ = np.tensordot(trap_ctr, self.PQ_, axes=([0], [0]))
             p = self.PL_ - mPQ
diff --git a/pysindy/utils/axes.py b/pysindy/utils/axes.py
index 7de086dba..25c42db8f 100644
--- a/pysindy/utils/axes.py
+++ b/pysindy/utils/axes.py
@@ -430,7 +430,6 @@ def decorator(func):
     return decorator
 
 
-
 @_implements(np.ravel)
 def ravel(a, order="C"):
     out = np.ravel(np.asarray(a), order=order)
@@ -455,7 +454,6 @@ def concatenate(arrays, axis=0, out=None, dtype=None, casting="same_kind"):
     ax_list = [obj.axes for obj in arrays if isinstance(obj, AxesArray)]
     for ax1, ax2 in zip(ax_list[:-1], ax_list[1:]):
         if ax1 != ax2:
-
             raise ValueError("Concatenating >1 AxesArray with incompatible axes")
     result = np.concatenate(parents, axis, out=out, dtype=dtype, casting=casting)
     if isinstance(out, AxesArray):
diff --git a/pysindy/utils/odes.py b/pysindy/utils/odes.py
index eb2b93ea0..fd2be0bbf 100644
--- a/pysindy/utils/odes.py
+++ b/pysindy/utils/odes.py
@@ -334,6 +334,7 @@ def burgers_galerkin(sigma=0.1, nu=0.025, U=1.0):
 # Below this line are models only suitable for SINDy-PI, since they are implicit #
 # and therefore require a library Theta(X, Xdot) rather than just Theta(X) #
 
+
 # Michaelis–Menten model for enzyme kinetics
 def enzyme(t, x, jx=0.6, Vmax=1.5, Km=0.3):
     return jx - Vmax * x / (Km + x)
diff --git a/test/conftest.py b/test/conftest.py
index 2c4f08310..59fb61f74 100644
--- a/test/conftest.py
+++ b/test/conftest.py
@@ -60,7 +60,6 @@ def data_1d_bad_shape():
 
 @pytest.fixture(scope="session")
 def data_lorenz():
-
     t = np.linspace(0, 1, 12)
     x0 = [8, 27, -7]
     x = solve_ivp(lorenz, (t[0], t[-1]), x0, t_eval=t).y.T
@@ -70,7 +69,6 @@ def data_lorenz():
 
 @pytest.fixture
 def data_multiple_trajectories():
-
     n_points = [100, 200, 500]
     initial_conditions = [
         [8, 27, -7],
@@ -120,7 +118,6 @@ def diffuse(t, u, dx, nx):
 
 @pytest.fixture(scope="session")
 def data_discrete_time():
-
     n_steps = 100
     mu = 3.6
     x = np.zeros((n_steps))
@@ -133,7 +130,6 @@ def data_discrete_time():
 
 @pytest.fixture(scope="session")
 def data_discrete_time_multiple_trajectories():
-
     n_steps = 100
     mus = [1, 2.3, 3.6]
     x = [np.zeros((n_steps)) for mu in mus]
@@ -421,7 +417,6 @@ def u_fun(t):
 
 @pytest.fixture(scope="session")
 def data_discrete_time_c():
-
     n_steps = 100
     mu = 3.6
 
@@ -437,7 +432,6 @@ def data_discrete_time_c():
 
 @pytest.fixture(scope="session")
 def data_discrete_time_c_multivariable():
-
     n_steps = 100
     mu = 3.6
 
@@ -454,7 +448,6 @@ def data_discrete_time_c_multivariable():
 
 @pytest.fixture(scope="session")
 def data_discrete_time_multiple_trajectories_c():
-
     n_steps = 100
     mus = [1, 2.3, 3.6]
     u = [0.001 * np.random.randn(n_steps) for mu in mus]
diff --git a/test/utils/test_axes.py b/test/utils/test_axes.py
index f3302cb62..6e85cd38b 100644
--- a/test/utils/test_axes.py
+++ b/test/utils/test_axes.py
@@ -29,13 +29,6 @@ def test_bad_concat():
         np.concatenate((arr, arr2), axis=0)
 
 
-def test_concat_out():
-    arr = AxesArray(np.arange(3).reshape(1, 3), {"ax_a": 0, "ax_b": 1})
-    arr_out = np.empty((2, 3)).view(AxesArray)
-    result = np.concatenate((arr, arr), axis=0, out=arr_out)
-    assert_equal(result, arr_out)
-
-
 def test_reduce_mean_noinf_recursion():
     arr = AxesArray(np.array([[1]]), {"ax_a": [0, 1]})
     np.mean(arr, axis=0)

From 9de7b3be41920360c68278099867856f52919bc7 Mon Sep 17 00:00:00 2001
From: Jake <37048747+Jacob-Stevens-Haas@users.noreply.github.com>
Date: Tue, 16 Apr 2024 13:33:09 -0700
Subject: [PATCH 16/16] CI: bump CI versions

---
 .github/workflows/main.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 350f16890..68385b634 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -29,7 +29,7 @@ jobs:
     strategy:
       max-parallel: 4
       matrix:
-        python-version: ["3.8", "3.10"]
+        python-version: ["3.9", "3.11"]
 
     steps:
     - uses: actions/checkout@v3