From 4ccebadf0f8b3dd4eec7b598f58e7ace19b3411d Mon Sep 17 00:00:00 2001 From: Anton Soloviev Date: Sat, 17 Oct 2026 00:24:56 -0700 Subject: [PATCH] fix: preserve sparse-group axes and zero-group limits --- examples/plot_norms.py | 8 +++++++ pyproximal/proximal/L21_plus_L1.py | 18 ++++++++++----- pytests/test_sparse_group_limits.py | 35 +++++++++++++++++++++++++++++ 3 files changed, 55 insertions(+), 6 deletions(-) create mode 100644 pytests/test_sparse_group_limits.py diff --git a/examples/plot_norms.py b/examples/plot_norms.py index bd80d371..b1f720a0 100755 --- a/examples/plot_norms.py +++ b/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)) diff --git a/pyproximal/proximal/L21_plus_L1.py b/pyproximal/proximal/L21_plus_L1.py index fe6c90dc..f69611e4 100644 --- a/pyproximal/proximal/L21_plus_L1.py +++ b/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 diff --git a/pytests/test_sparse_group_limits.py b/pytests/test_sparse_group_limits.py new file mode 100644 index 00000000..2056f450 --- /dev/null +++ b/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)