Merge pull request #493 from dynamicslab/trapping-plumes

Trapping part 2/3

.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

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))

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]

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:

pysindy/feature_library/polynomial_library.py

@@ -1,7 +1,10 @@
 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
 from scipy import sparse
 from sklearn.preprocessing import PolynomialFeatures
 from sklearn.utils.validation import check_is_fitted
@@ -73,8 +76,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 +110,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 +230,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 +249,36 @@ def transform(self, x_full):
                     ),
                     x.axes,
                 )
-                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

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,

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(

pysindy/optimizers/constrained_sr3.py

@@ -1,4 +1,6 @@
 import warnings
+from copy import deepcopy
+from typing import Optional
 from typing import Tuple
 
 try:
@@ -72,8 +74,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.
@@ -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(

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")