Merge pull request #408 from dynamicslab/sparse-indices

Allow STLSQ to only sparsify/regularize a subset of indices

pysindy/optimizers/base.py

@@ -75,7 +75,7 @@ class BaseOptimizer(LinearRegression, ComplexityMixin):
         Weight vector(s).
 
     ind_ : array, shape (n_features,) or (n_targets, n_features)
-        Array of 0s and 1s indicating which coefficients of the
+        Array of bools indicating which coefficients of the
         weight vector have not been masked out.
 
     history_ : list
@@ -204,7 +204,7 @@ def fit(self, x_, y, sample_weight=None, **reduce_kws):
         self._set_intercept(X_offset, y_offset, X_scale)
         return self
 
-    def _unbias(self, x, y):
+    def _unbias(self, x: np.ndarray, y: np.ndarray):
         coef = np.zeros((y.shape[1], x.shape[1]))
         for i in range(self.ind_.shape[0]):
             if np.any(self.ind_[i]):

pysindy/optimizers/stlsq.py

@@ -1,8 +1,10 @@
 import warnings
+from typing import Union
 
 import numpy as np
 from scipy.linalg import LinAlgWarning
 from sklearn.exceptions import ConvergenceWarning
+from sklearn.linear_model import LinearRegression
 from sklearn.linear_model import ridge_regression
 from sklearn.utils.validation import check_is_fitted
 
@@ -58,6 +60,11 @@ class STLSQ(BaseOptimizer):
     verbose : bool, optional (default False)
         If True, prints out the different error terms every iteration.
 
+    sparse_ind : list, optional (default None)
+        Indices to threshold and perform ridge regression upon.
+        If None, sparse thresholding and ridge regression is applied to all
+        indices.
+
     Attributes
     ----------
     coef_ : array, shape (n_features,) or (n_targets, n_features)
@@ -102,6 +109,7 @@ def __init__(
         copy_X=True,
         initial_guess=None,
         verbose=False,
+        sparse_ind=None,
         unbias=True,
     ):
         super().__init__(
@@ -121,28 +129,55 @@ def __init__(
         self.ridge_kw = ridge_kw
         self.initial_guess = initial_guess
         self.verbose = verbose
+        self.sparse_ind = sparse_ind
 
-    def _sparse_coefficients(self, dim, ind, coef, threshold):
-        """Perform thresholding of the weight vector(s)"""
+    def _sparse_coefficients(
+        self, dim: int, ind_nonzero: np.ndarray, coef: np.ndarray, threshold: float
+    ) -> (np.ndarray, np.ndarray):
+        """Perform thresholding of the weight vector(s) (on specific indices
+        if ``self.sparse_ind`` is not None)"""
         c = np.zeros(dim)
-        c[ind] = coef
+        c[ind_nonzero] = coef
         big_ind = np.abs(c) >= threshold
+        if self.sparse_ind is not None:
+            nonsparse_ind_mask = np.ones_like(ind_nonzero)
+            nonsparse_ind_mask[self.sparse_ind] = False
+            big_ind = big_ind | nonsparse_ind_mask
         c[~big_ind] = 0
         return c, big_ind
 
-    def _regress(self, x, y):
-        """Perform the ridge regression"""
+    def _regress(self, x: np.ndarray, y: np.ndarray, dim: int, sparse_sub: np.ndarray):
+        """Perform the ridge regression (on specific indices if
+        ``self.sparse_ind`` is not None)"""
         kw = self.ridge_kw or {}
-
-        with warnings.catch_warnings():
-            warnings.filterwarnings("ignore", category=LinAlgWarning)
-            try:
-                coef = ridge_regression(x, y, self.alpha, **kw)
-            except LinAlgWarning:
-                # increase alpha until warning stops
-                self.alpha = 2 * self.alpha
-        self.iters += 1
-        return coef
+        if self.sparse_ind is None:
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", category=LinAlgWarning)
+                try:
+                    coef = ridge_regression(x, y, self.alpha, **kw)
+                except LinAlgWarning:
+                    # increase alpha until warning stops
+                    self.alpha = 2 * self.alpha
+            self.iters += 1
+            return coef
+        if self.sparse_ind is not None:
+            alpha_array = np.zeros((dim, dim))
+            alpha_array[sparse_sub, sparse_sub] = np.sqrt(self.alpha)
+            x_lin = np.concatenate((x, alpha_array), axis=0)
+            y_lin = np.concatenate((y, np.zeros((dim,))))
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", category=LinAlgWarning)
+                try:
+                    coef = (
+                        LinearRegression(fit_intercept=False, **kw)
+                        .fit(x_lin, y_lin)
+                        .coef_
+                    )
+                except LinAlgWarning:
+                    # increase alpha until warning stops
+                    self.alpha = 2 * self.alpha
+            self.iters += 1
+            return coef
 
     def _no_change(self):
         """Check if the coefficient mask has changed after thresholding"""
@@ -167,6 +202,7 @@ def _reduce(self, x, y):
         n_samples, n_features = x.shape
         n_targets = y.shape[1]
         n_features_selected = np.sum(ind)
+        sparse_sub = [np.array(self.sparse_ind)] * y.shape[1]
 
         # Print initial values for each term in the optimization
         if self.verbose:
@@ -188,29 +224,38 @@ def _reduce(self, x, y):
                     "Sparsity parameter is too big ({}) and eliminated all "
                     "coefficients".format(self.threshold)
                 )
-                coef = np.zeros((n_targets, n_features))
+                optvar = np.zeros((n_targets, n_features))
                 break
 
-            coef = np.zeros((n_targets, n_features))
+            optvar = np.zeros((n_targets, n_features))
             for i in range(n_targets):
                 if np.count_nonzero(ind[i]) == 0:
                     warnings.warn(
                         "Sparsity parameter is too big ({}) and eliminated all "
                         "coefficients".format(self.threshold)
                     )
                     continue
-                coef_i = self._regress(x[:, ind[i]], y[:, i])
+                coef_i = self._regress(
+                    x[:, ind[i]], y[:, i], np.count_nonzero(ind[i]), sparse_sub[i]
+                )
                 coef_i, ind_i = self._sparse_coefficients(
                     n_features, ind[i], coef_i, self.threshold
                 )
-                coef[i] = coef_i
+                if self.sparse_ind is not None:
+                    vals_to_remove = np.intersect1d(
+                        self.sparse_ind, np.where(coef_i == 0)
+                    )
+                    sparse_sub[i] = _remove_and_decrement(
+                        self.sparse_ind, vals_to_remove
+                    )
+                optvar[i] = coef_i
                 ind[i] = ind_i
 
-            self.history_.append(coef)
+            self.history_.append(optvar)
             if self.verbose:
-                R2 = np.sum((y - np.dot(x, coef.T)) ** 2)
-                L2 = self.alpha * np.sum(coef**2)
-                L0 = np.count_nonzero(coef)
+                R2 = np.sum((y - np.dot(x, optvar.T)) ** 2)
+                L2 = self.alpha * np.sum(optvar**2)
+                L0 = np.count_nonzero(optvar)
                 row = [k, R2, L2, L0, R2 + L2]
                 print(
                     "{0:10d} ... {1:10.4e} ... {2:10.4e} ... {3:10d}"
@@ -226,16 +271,15 @@ def _reduce(self, x, y):
                 ),
                 ConvergenceWarning,
             )
-            try:
-                coef
-            except NameError:
-                coef = self.coef_
-                warnings.warn(
-                    "STLSQ._reduce has no iterations left to determine coef",
-                    ConvergenceWarning,
-                )
-        self.coef_ = coef
-        self.ind_ = ind
+        if self.sparse_ind is None:
+            self.coef_ = optvar
+            self.ind_ = ind
+        else:
+            non_sparse_ind = np.setxor1d(self.sparse_ind, range(n_features))
+            self.coef_ = optvar[:, self.sparse_ind]
+            self.ind_ = ind[:, self.sparse_ind]
+            self.optvar_non_sparse_ = optvar[:, non_sparse_ind]
+            self.ind_non_sparse_ = ind[:, non_sparse_ind]
 
     @property
     def complexity(self):
@@ -244,3 +288,32 @@ def complexity(self):
         return np.count_nonzero(self.coef_) + np.count_nonzero(
             [abs(self.intercept_) >= self.threshold]
         )
+
+    def _unbias(self, x: np.ndarray, y: np.ndarray):
+        if not self.sparse_ind:
+            return super()._unbias(x, y)
+        regression_col_mask = np.zeros((y.shape[1], x.shape[1]), dtype=bool)
+        regression_col_mask[:, self.sparse_ind] = self.ind_
+        non_sparse_ind = np.setxor1d(self.sparse_ind, range(x.shape[1]))
+        regression_col_mask[:, non_sparse_ind] = self.ind_non_sparse_
+
+        for i in range(self.ind_.shape[0]):
+            if np.any(self.ind_[i]):
+                optvar = (
+                    LinearRegression(fit_intercept=False)
+                    .fit(x[:, regression_col_mask[i]], y[:, i])
+                    .coef_
+                )
+                self.coef_[i] = optvar[self.sparse_ind]
+                self.optvar_non_sparse_[i] = optvar[non_sparse_ind]
+
+
+def _remove_and_decrement(
+    existing_vals: Union[np.ndarray, list], vals_to_remove: Union[np.ndarray, list]
+) -> np.ndarray:
+    """Remove elements from existing values and decrement the elements
+    that are greater than the removed elements"""
+    for s in reversed(vals_to_remove):
+        existing_vals = np.delete(existing_vals, np.where(s == existing_vals))
+        existing_vals = np.where(existing_vals > s, existing_vals - 1, existing_vals)
+    return existing_vals

test/test_optimizers.py

@@ -29,6 +29,7 @@
 from pysindy.optimizers import STLSQ
 from pysindy.optimizers import TrappingSR3
 from pysindy.optimizers import WrappedOptimizer
+from pysindy.optimizers.stlsq import _remove_and_decrement
 from pysindy.utils import supports_multiple_targets
 from pysindy.utils.odes import enzyme
 
@@ -1082,3 +1083,49 @@ def test_frols_error_linear_dependence():
     y = np.array([[1.0, 1.0]])
     with pytest.raises(ValueError):
         opt.fit(x, y)
+
+
+def test_sparse_subset_multitarget():
+    A = np.eye(4)
+    b = np.array([[1, 1, 0.5, 1], [1, 1, 1, 0.5]]).T
+    opt = STLSQ(unbias=False, threshold=0.5, alpha=0.1, sparse_ind=[2, 3])
+    opt.fit(A, b)
+    X = opt.coef_
+    Y = opt.optvar_non_sparse_
+    assert X[0, 0] == 0.0
+    assert 0.0 < X[0, 1] < 1.0
+    np.testing.assert_equal(Y[:, :2], np.ones((2, 2)))
+    assert X[1, 1] == 0.0
+    assert 0.0 < X[1, 0] < 1.0
+
+
+def test_sparse_subset_off_diagonal():
+    A = np.array([[1, 1], [0, 1]])
+    b = np.array([1, 1])
+    opt = STLSQ(unbias=False, threshold=0.1, alpha=0.1, sparse_ind=[1])
+    opt.fit(A, b)
+    X = opt.coef_
+    Y = opt.optvar_non_sparse_
+    assert Y[0, 0] > 0.0 and Y[0, 0] < 0.5
+    assert X[0, 0] > 0.5 and X[0, 0] < 1.0
+
+
+def test_sparse_subset_unbias():
+    A = np.array([[1, 1], [0, 1]])
+    b = np.array([1, 1])
+    opt = STLSQ(unbias=True, threshold=0.1, alpha=0.1, sparse_ind=[1])
+    opt.fit(A, b)
+    X = opt.coef_
+    Y = opt.optvar_non_sparse_
+    assert np.abs(Y[0, 0]) < 2e-16
+    assert np.abs(X[0, 0] - 1.0) < 2e-16
+
+
+def test_remove_and_decrement():
+    existing_vals = np.array([2, 3, 4, 5])
+    vals_to_remove = np.array([3, 5])
+    expected = np.array([2, 3])
+    result = _remove_and_decrement(
+        existing_vals=existing_vals, vals_to_remove=vals_to_remove
+    )
+    np.testing.assert_array_equal(expected, result)