Preserve sparse-group axes and finite zero-group limits

examples/plot_norms.py

@@ -8,6 +8,7 @@
 - L2 norm (:class:`pyproximal.L2`)
 - L1 norm (:class:`pyproximal.L1`)
 - L21 norm (:class:`pyproximal.L21`)
+- Sparse-group norm (:class:`pyproximal.L21_plus_L1`)
 
 """
 
@@ -189,3 +190,10 @@
 plt.title(r"$||x||_1$")
 plt.legend()
 plt.tight_layout()
+
+###############################################################################
+# Sparse-group shrinkage can group either rows or columns. Its pure L1 limit
+# also handles zero groups without producing NaNs.
+sparse_group = pyproximal.L21_plus_L1(sigma=1.0, rho=1.0)
+grouped_x = np.array([[0.0, 1.0, 3.0], [0.0, 0.5, -4.0]])
+print("Pure L1 sparse-group shrinkage: ", sparse_group.prox(grouped_x, 0.5, axis=1))

pyproximal/proximal/L21_plus_L1.py

@@ -23,6 +23,10 @@ class L21_plus_L1(ProxOperator):
     The proximal operator of the :math:`L_{2,1} + L_1` mixed-norm is simply the
     product of each individual proximal operator [1]_.
 
+    Grouping is performed along ``axis`` in :meth:`prox` (columns by default).
+    When ``rho=1``, only elementwise L1 shrinkage remains. With ``sigma=0``,
+    the proximal is the identity. Zero groups remain zero in both cases.
+
     .. [1] Gramfort, Alexandre, Daniel Strohmeier, Jens Haueisen, Matti Hamalainen,
         and Matthieu Kowalski. "Functional brain imaging with M/EEG using structured
         sparsity in time-frequency dictionaries." In Biennial International Conference
@@ -46,11 +50,13 @@ def __call__(self, x: NDArray) -> float:
     @_check_tau
     def prox(self, x: NDArray, tau: float, axis: int = 0) -> NDArray:
         thresh = self.sigma * tau
-        l1 = np.maximum(np.abs(x) - thresh * self.rho, 0)
+        absx = np.abs(x)
+        l1 = np.maximum(absx - thresh * self.rho, 0)
         # Axis defines what dimension to perform grouping over
-        aux_l21 = np.sqrt(
-            np.sum(np.maximum(np.abs(x) - thresh * self.rho, 0) ** 2, axis=axis)
-        )
-        l21 = np.maximum(1 - thresh * (1 - self.rho) / aux_l21, 0)
-        x = np.nan_to_num(x / np.abs(x)) * l1 * l21
+        aux_l21 = np.sqrt(np.sum(l1**2, axis=axis, keepdims=True))
+        denominator = np.maximum(aux_l21, thresh * (1 - self.rho))
+        denominator = np.where(denominator == 0, 1, denominator)
+        l21 = np.maximum(1 - thresh * (1 - self.rho) / denominator, 0)
+        phase = np.where(absx == 0, 0, x) / np.where(absx == 0, 1, absx)
+        x = phase * l1 * l21
         return x

pytests/test_sparse_group_limits.py

@@ -0,0 +1,35 @@
+import numpy as np
+import pytest
+
+from pyproximal import L21_plus_L1
+
+
+@pytest.mark.parametrize("dtype", [np.float64, np.complex128])
+@pytest.mark.parametrize("axis", [0, 1, -1])
+@pytest.mark.parametrize("sigma,rho", [(1.0, 1.0), (0.0, 0.8)])
+def test_sparse_group_limits(dtype, axis, sigma, rho):
+    """Pure L1 and zero regularization are finite, including zero groups."""
+    x = np.array([[0, 1, 3], [0, 0.5, -4]], dtype=dtype)
+    if np.issubdtype(dtype, np.complexfloating):
+        x *= (1 + 1j) / np.sqrt(2)
+    tau = 0.5
+    expected = np.array(
+        [
+            0 if abs(value) <= tau * sigma else value * (1 - tau * sigma / abs(value))
+            for value in x.ravel()
+        ]
+    ).reshape(x.shape)
+    with np.errstate(divide="raise", invalid="raise"):
+        result = L21_plus_L1(sigma=sigma, rho=rho).prox(x, tau, axis=axis)
+    np.testing.assert_allclose(result, expected, atol=1e-14)
+
+
+@pytest.mark.parametrize("dtype", [np.float64, np.complex128])
+def test_sparse_group_row_axis(dtype):
+    """Grouping rows is equivalent to grouping columns of the transpose."""
+    x = np.array([[0, 1, 3], [0, 0.5, -4]], dtype=dtype)
+    if np.issubdtype(dtype, np.complexfloating):
+        x *= (1 + 1j) / np.sqrt(2)
+    operator = L21_plus_L1(sigma=0.8, rho=0.4)
+    expected = operator.prox(x.T, 0.5).T
+    np.testing.assert_allclose(operator.prox(x, 0.5, axis=1), expected)