BUG #394 fix prox and regularization

examples/1_feature_overview/example.py

@@ -406,7 +406,7 @@ def f(x):
 model.print()
 
 # With thresholds matrix
-thresholds = 2 * np.ones((10, 3))
+thresholds = 2 * np.ones((3, 10))
 thresholds[4:, :] = 0.1
 sr3_optimizer = ps.SR3(thresholder="weighted_l0", thresholds=thresholds)
 model = ps.SINDy(optimizer=sr3_optimizer).fit(x_train, t=dt)

pysindy/optimizers/constrained_sr3.py

@@ -279,8 +279,8 @@ def _calculate_penalty(
         regularization: 'l0' | 'weighted_l0' | 'l1' | 'weighted_l1' |
                         'l2' | 'weighted_l2'
         regularization_weight: float | np.array, can be a scalar
-                               or an array of shape (n_targets, n_features)
-        xi: cp.Variable
+                               or an array of shape (m, n)
+        xi: cp.Variable of length m*n
 
         Returns:
         --------

pysindy/utils/base.py

@@ -153,7 +153,7 @@ def reorder_constraints(arr, n_features, output_order="feature"):
     return arr.reshape(starting_shape).transpose([0, 2, 1]).reshape((n_constraints, -1))
 
 
-def validate_prox_and_reg_inputs(func, regularization):
+def _validate_prox_and_reg_inputs(func, regularization):
     def wrapper(x, regularization_weight):
         if regularization[:8] == "weighted":
             if not isinstance(regularization_weight, np.ndarray):
@@ -163,15 +163,11 @@ def wrapper(x, regularization_weight):
             weight_shape = regularization_weight.shape
             if weight_shape != x.shape:
                 raise ValueError(
-                    f"Invalid shape for 'regularization_weight': \
-                        {weight_shape}. Must be the same shape as x: {x.shape}."
+                    f"Invalid shape for 'regularization_weight':"
+                    f"{weight_shape}. Must be the same shape as x: {x.shape}."
                 )
-        else:
-            if not isinstance(regularization_weight, (int, float)) and (
-                isinstance(regularization_weight, np.ndarray)
-                and regularization_weight.shape not in [(1, 1), (1,)]
-            ):
-                raise ValueError("'regularization_weight' must be a scalar")
+        elif not isinstance(regularization_weight, (int, float)):
+            raise ValueError("'regularization_weight' must be a scalar")
         return func(x, regularization_weight)
 
     return wrapper
@@ -180,7 +176,7 @@ def wrapper(x, regularization_weight):
 def get_prox(
     regularization: str,
 ) -> Callable[
-    [NDArray[np.float64], Union[np.float64, NDArray[np.float64]]], NDArray[np.float64]
+    [NDArray[np.float64], Union[float, NDArray[np.float64]]], NDArray[np.float64]
 ]:
     """
     Args:
@@ -190,59 +186,47 @@ def get_prox(
     Returns:
     --------
     proximal_operator: (x: np.array, reg_weight: float | np.array) -> np.array
-        A function that takes an input x of shape (n_targets, n_features)
+        A function that takes an input x of shape (m, n)
         and regularization weight factor which can be a scalar or
-        an array of shape (n_targets, n_features),
-        and returns an array of shape (n_targets, n_features)
+        an array of shape (m, n),
+        and returns an array of shape (m, n)
     """
 
-    def prox_l0(x: NDArray[np.float64], regularization_weight: np.float64):
-        """Proximal operator for L0 regularization."""
-        threshold = np.sqrt(2 * regularization_weight)
-        return x * (np.abs(x) > threshold)
-
-    def prox_weighted_l0(
-        x: NDArray[np.float64], regularization_weight: NDArray[np.float64]
+    def prox_l0(
+        x: NDArray[np.float64],
+        regularization_weight: Union[float, NDArray[np.float64]],
     ):
-        """Proximal operator for weighted l0 regularization."""
         threshold = np.sqrt(2 * regularization_weight)
         return x * (np.abs(x) > threshold)
 
-    def prox_l1(x: NDArray[np.float64], regularization_weight: np.float64):
-        """Proximal operator for L1 regularization."""
-        return np.sign(x) * np.maximum(np.abs(x) - regularization_weight, 0)
-
-    def prox_weighted_l1(
-        x: NDArray[np.float64], regularization_weight: NDArray[np.float64]
+    def prox_l1(
+        x: NDArray[np.float64],
+        regularization_weight: Union[float, NDArray[np.float64]],
     ):
-        """Proximal operator for weighted l1 regularization."""
-        return np.sign(x) * np.maximum(np.abs(x) - regularization_weight, 0)
 
-    def prox_l2(x: NDArray[np.float64], regularization_weight: np.float64):
-        """Proximal operator for ridge regularization."""
-        return x / (1 + 2 * regularization_weight)
+        return np.sign(x) * np.maximum(np.abs(x) - regularization_weight, 0)
 
-    def prox_weighted_l2(
-        x: NDArray[np.float64], regularization_weight: NDArray[np.float64]
+    def prox_l2(
+        x: NDArray[np.float64],
+        regularization_weight: Union[float, NDArray[np.float64]],
     ):
-        """Proximal operator for ridge regularization."""
         return x / (1 + 2 * regularization_weight)
 
     prox = {
         "l0": prox_l0,
-        "weighted_l0": prox_weighted_l0,
+        "weighted_l0": prox_l0,
         "l1": prox_l1,
-        "weighted_l1": prox_weighted_l1,
+        "weighted_l1": prox_l1,
         "l2": prox_l2,
-        "weighted_l2": prox_weighted_l2,
+        "weighted_l2": prox_l2,
     }
     regularization = regularization.lower()
-    return validate_prox_and_reg_inputs(prox[regularization], regularization)
+    return _validate_prox_and_reg_inputs(prox[regularization], regularization)
 
 
 def get_regularization(
     regularization: str,
-) -> Callable[[NDArray[np.float64], Union[np.float64, NDArray[np.float64]]], float]:
+) -> Callable[[NDArray[np.float64], Union[float, NDArray[np.float64]]], float]:
     """
     Args:
     -----
@@ -251,46 +235,43 @@ def get_regularization(
     Returns:
     --------
     regularization_function: (x: np.array, reg_weight: float | np.array) -> np.array
-        A function that takes an input x of shape (n_targets, n_features)
+        A function that takes an input x of shape (m, n)
         and regularization weight factor which can be a scalar or
-        an array of shape (n_targets, n_features),
+        an array of shape (m, n),
         and returns a float
     """
 
-    def regularization_l0(x: NDArray[np.float64], regularization_weight: np.float64):
-        return regularization_weight * np.count_nonzero(x)
-
-    def regualization_weighted_l0(
-        x: NDArray[np.float64], regularization_weight: NDArray[np.float64]
+    def regularization_l0(
+        x: NDArray[np.float64],
+        regularization_weight: Union[float, NDArray[np.float64]],
     ):
-        return np.sum(regularization_weight[np.nonzero(x)])
 
-    def regularization_l1(x: NDArray[np.float64], regularization_weight: np.float64):
-        return np.sum(regularization_weight * np.abs(x))
+        return np.sum(regularization_weight * (x != 0))
 
-    def regualization_weighted_l1(
-        x: NDArray[np.float64], regularization_weight: NDArray[np.float64]
+    def regularization_l1(
+        x: NDArray[np.float64],
+        regularization_weight: Union[float, NDArray[np.float64]],
     ):
-        return np.sum(regularization_weight * np.abs(x))
 
-    def regularization_l2(x: NDArray[np.float64], regularization_weight: np.float64):
-        return np.sum(regularization_weight * x**2)
+        return np.sum(regularization_weight * np.abs(x))
 
-    def regualization_weighted_l2(
-        x: NDArray[np.float64], regularization_weight: NDArray[np.float64]
+    def regularization_l2(
+        x: NDArray[np.float64],
+        regularization_weight: Union[float, NDArray[np.float64]],
     ):
+
         return np.sum(regularization_weight * x**2)
 
     regularization_fn = {
         "l0": regularization_l0,
-        "weighted_l0": regualization_weighted_l0,
+        "weighted_l0": regularization_l0,
         "l1": regularization_l1,
-        "weighted_l1": regualization_weighted_l1,
+        "weighted_l1": regularization_l1,
         "l2": regularization_l2,
-        "weighted_l2": regualization_weighted_l2,
+        "weighted_l2": regularization_l2,
     }
     regularization = regularization.lower()
-    return validate_prox_and_reg_inputs(
+    return _validate_prox_and_reg_inputs(
         regularization_fn[regularization], regularization
     )
 

test/utils/test_utils.py

@@ -83,35 +83,16 @@ def test_get_regularization(regularization, lam, expected):
     assert result == expected
 
 
-@pytest.mark.parametrize("regularization", ["l0", "l1", "l2"])
-@pytest.mark.parametrize("lam", [1, np.array([1]), np.array([[1]])])
-def test_get_prox_and_regularization_shape(regularization, lam):
-    data = np.array([[-2, 5]]).T
-    reg = get_regularization(regularization)
-    reg_result = reg(data, lam)
-    prox = get_prox(regularization)
-    prox_result = prox(data, lam)
-    assert reg_result is not None
-    assert prox_result is not None
-
-
-@pytest.mark.parametrize(
-    "regularization", ["weighted_l0", "weighted_l1", "weighted_l2"]
-)
-@pytest.mark.parametrize("lam", [np.array([[1, 2]]).T])
-def test_get_weighted_prox_and_regularization_shape(regularization, lam):
-    data = np.array([[-2, 5]]).T
-    reg = get_regularization(regularization)
-    reg_result = reg(data, lam)
-    prox = get_prox(regularization)
-    prox_result = prox(data, lam)
-    assert reg_result is not None
-    assert prox_result is not None
-
-
 @pytest.mark.parametrize("regularization", ["l0", "l1", "l2"])
 @pytest.mark.parametrize(
-    "lam", [np.array([[1, 2]]), np.array([1, 2]), np.array([[1, 2]]).T]
+    "lam",
+    [
+        np.array([[1, 2]]),
+        np.array([1, 2]),
+        np.array([[1, 2]]).T,
+        np.array([1]),
+        np.array([[1]]),
+    ],
 )
 def test_get_prox_and_regularization_bad_shape(regularization, lam):
     data = np.array([[-2, 5]]).T
@@ -127,7 +108,15 @@ def test_get_prox_and_regularization_bad_shape(regularization, lam):
     "regularization", ["weighted_l0", "weighted_l1", "weighted_l2"]
 )
 @pytest.mark.parametrize(
-    "lam", [np.array([[1, 2]]), np.array([1, 2, 3]), np.array([[1, 2, 3]]).T, 1]
+    "lam",
+    [
+        np.array([[1, 2]]),
+        np.array([1, 2, 3]),
+        np.array([[1, 2, 3]]).T,
+        1,
+        np.array([1]),
+        np.array([[1]]),
+    ],
 )
 def test_get_weighted_prox_and_regularization_bad_shape(regularization, lam):
     data = np.array([[-2, 5]]).T