diff --git a/pylops/optimization/cls_sparsity.py b/pylops/optimization/cls_sparsity.py index 59ce8a86e..f800a6138 100644 --- a/pylops/optimization/cls_sparsity.py +++ b/pylops/optimization/cls_sparsity.py @@ -23,7 +23,12 @@ from pylops.optimization.eigs import power_iteration from pylops.optimization.leastsquares import regularized_inversion from pylops.utils import deps -from pylops.utils.backend import get_array_module, get_module_name, inplace_set +from pylops.utils.backend import ( + get_array_module, + get_module_name, + get_real_dtype, + inplace_set, +) from pylops.utils.typing import InputDimsLike, NDArray, SamplingLike spgl1_message = deps.spgl1_import("the spgl1 solver") @@ -58,7 +63,7 @@ def _hardthreshold(x: NDArray, thresh: float) -> NDArray: """ x1 = x.copy() - x1[np.abs(x) <= sqrt(2 * thresh)] = 0 + x1[np.abs(x) <= sqrt(2 * thresh)] = 0.0 return x1 @@ -120,13 +125,17 @@ def _halfthreshold(x: NDArray, thresh: float) -> NDArray: Since version 1.17.0 does not produce ``np.nan`` on bad input. """ - arg = np.ones_like(x) - arg[x != 0] = (thresh / 8.0) * (np.abs(x[x != 0]) / 3.0) ** (-1.5) - arg = np.clip(arg, -1, 1) - phi = 2.0 / 3.0 * np.arccos(arg) - x1 = 2.0 / 3.0 * x * (1 + np.cos(2.0 * np.pi / 3.0 - phi)) - # x1[np.abs(x) <= 1.5 * thresh ** (2. / 3.)] = 0 - x1[np.abs(x) <= (54 ** (1.0 / 3.0) / 4.0) * thresh ** (2.0 / 3.0)] = 0 + ncp = get_array_module(x) + arg = ncp.ones_like(x) + arg[x != 0] = (thresh / 8.0) * (ncp.abs(x[x != 0.0]) / 3.0) ** (-1.5) + if ncp.iscomplexobj(arg): + arg.real = ncp.clip(arg.real, -1.0, 1.0) + arg.imag = ncp.clip(arg.imag, -1.0, 1.0) + else: + arg = ncp.clip(arg, -1.0, 1.0) + phi = 2.0 / 3.0 * ncp.arccos(arg) + x1 = 2.0 / 3.0 * x * (1.0 + ncp.cos(2.0 * np.pi / 3.0 - phi)) + x1[ncp.abs(x) <= (54.0 ** (1.0 / 3.0) / 4.0) * thresh ** (2.0 / 3.0)] = 0 return x1 @@ -1609,7 +1618,7 @@ def setup( # prepare decay (if not passed) if perc is None and decay is None: - self.decay = self.ncp.ones(niter, dtype=self.Op) + self.decay = self.ncp.ones(niter, dtype=get_real_dtype(self.Op.dtype)) # step size if alpha is not None: @@ -1751,6 +1760,7 @@ def step(self, x: NDArray, show: bool = False) -> Tuple[NDArray, float]: self.SOpx_unthesh[:] = self.SOprmatvec(self.x_unthesh) else: SOpx_unthesh = self.SOprmatvec(x_unthesh) + # threshold current solution or current solution projected onto SOp.H space if self.SOp is None: x_unthesh_or_SOpx_unthesh = ( @@ -1767,6 +1777,7 @@ def step(self, x: NDArray, show: bool = False) -> Tuple[NDArray, float]: ) else: x = self.threshf(x_unthesh_or_SOpx_unthesh, 100 - self.perc) + # apply SOp to thresholded x if self.SOp is not None: x = self.SOpmatvec(x) diff --git a/pytests/test_sparsity.py b/pytests/test_sparsity.py index ffa11871c..a996289ee 100644 --- a/pytests/test_sparsity.py +++ b/pytests/test_sparsity.py @@ -10,6 +10,7 @@ from numpy.testing import assert_array_almost_equal backend = "numpy" +import numpy as npp import pytest from pylops.basicoperators import FirstDerivative, Identity, MatrixMult @@ -18,7 +19,7 @@ from pylops.optimization.sparsity import fista, irls, ista, omp, spgl1, splitbregman # currently test spgl1 only if numpy<2.0.0 is installed... -np_version = np.__version__.split(".") +np_version = npp.__version__.split(".") par1 = { "ny": 11, @@ -60,35 +61,35 @@ "nx": 11, "imag": 1j, "x0": False, - "dtype": "complex64", + "dtype": "complex128", } # square complex, zero initial guess par2j = { "ny": 11, "nx": 11, "imag": 1j, "x0": True, - "dtype": "complex64", + "dtype": "complex128", } # square complex, non-zero initial guess par3j = { "ny": 31, "nx": 11, "imag": 1j, "x0": False, - "dtype": "complex64", + "dtype": "complex128", } # overdetermined complex, zero initial guess par4j = { "ny": 31, "nx": 11, "imag": 1j, "x0": True, - "dtype": "complex64", + "dtype": "complex128", } # overdetermined complex, non-zero initial guess par5j = { "ny": 21, "nx": 41, "imag": 1j, "x0": True, - "dtype": "complex64", + "dtype": "complex128", } # underdetermined complex, non-zero initial guess @@ -101,19 +102,22 @@ def test_IRLS_unknown_kind(): @pytest.mark.parametrize("par", [(par3), (par4), (par3j), (par4j)]) def test_IRLS_data(par): """Invert problem with outliers using data IRLS""" - np.random.seed(10) - G = np.random.normal(0, 10, (par["ny"], par["nx"])).astype("float32") + par[ + npp.random.seed(10) + A = npp.random.normal(0, 10, (par["ny"], par["nx"])) + par[ "imag" - ] * np.random.normal(0, 10, (par["ny"], par["nx"])).astype("float32") - Gop = MatrixMult(G, dtype=par["dtype"]) + ] * npp.random.normal(0, 10, (par["ny"], par["nx"])) + Aop = MatrixMult(np.asarray(A), dtype=par["dtype"]) + x = np.ones(par["nx"]) + par["imag"] * np.ones(par["nx"]) x0 = ( - np.random.normal(0, 1, par["nx"]) - + par["imag"] * np.random.normal(0, 1, par["nx"]) + np.asarray( + npp.random.normal(0, 1, par["nx"]) + + par["imag"] * npp.random.normal(0, 1, par["nx"]) + ) if par["x0"] else None ) - y = Gop * x + y = Aop * x # add outlier y[par["ny"] - 2] *= 5 @@ -121,7 +125,7 @@ def test_IRLS_data(par): # irls inversion for preallocate in [False, True]: xinv = irls( - Gop, + Aop, y, x0=x0, nouter=10, @@ -138,23 +142,26 @@ def test_IRLS_data(par): @pytest.mark.parametrize("par", [(par3), (par4), (par3j), (par4j)]) def test_IRLS_datamodel(par): """Invert problem with outliers using data-model IRLS""" - np.random.seed(10) - G = np.random.normal(0, 10, (par["ny"], par["nx"])).astype("float32") + par[ + npp.random.seed(10) + A = npp.random.normal(0, 10, (par["ny"], par["nx"])) + par[ "imag" - ] * np.random.normal(0, 10, (par["ny"], par["nx"])).astype("float32") - Gop = MatrixMult(G, dtype=par["dtype"]) + ] * npp.random.normal(0, 10, (par["ny"], par["nx"])) + Aop = MatrixMult(np.asarray(A), dtype=par["dtype"]) + + x = np.zeros(par["nx"]) + par["imag"] * np.zeros(par["nx"]) + x[par["nx"] // 2] = 1.0 + par["imag"] * 1.0 + x[3] = 1.0 + par["imag"] * 1.0 + x[par["nx"] - 4] = -1.0 - par["imag"] * 1.0 - x = np.zeros(par["nx"]) + par["imag"] * np.ones(par["nx"]) - x[par["nx"] // 2] = 1 - x[3] = 1 - x[par["nx"] - 4] = -1 x0 = ( - np.random.normal(0, 1, par["nx"]) - + par["imag"] * np.random.normal(0, 1, par["nx"]) + np.asarray( + npp.random.normal(0, 1, par["nx"]) + + par["imag"] * npp.random.normal(0, 1, par["nx"]) + ) if par["x0"] else None ) - y = Gop * x + y = Aop * x # add outlier y[par["ny"] - 2] *= 5 @@ -162,7 +169,7 @@ def test_IRLS_datamodel(par): # irls inversion for preallocate in [False, True]: xinv = irls( - Gop, + Aop, y, x0=x0, nouter=10, @@ -179,17 +186,19 @@ def test_IRLS_datamodel(par): @pytest.mark.parametrize("par", [(par1), (par3), (par5), (par1j), (par3j), (par5j)]) def test_IRLS_model(par): """Invert problem with model IRLS""" - np.random.seed(42) - Aop = MatrixMult(np.random.randn(par["ny"], par["nx"])) + npp.random.seed(42) + A = npp.random.randn(par["ny"], par["nx"]) + par["imag"] * npp.random.randn( + par["ny"], par["nx"] + ) + Aop = MatrixMult(np.asarray(A), dtype=par["dtype"]) - x = np.zeros(par["nx"]) - x[par["nx"] // 2] = 1 - x[3] = 1 - x[par["nx"] - 4] = -1 + x = np.zeros(par["nx"]) + par["imag"] * np.zeros(par["nx"]) + x[par["nx"] // 2] = 1.0 + par["imag"] * 1.0 + x[3] = 1.0 + par["imag"] * 1.0 + x[par["nx"] - 4] = -1.0 - par["imag"] * 1.0 y = Aop * x maxit = 100 - for preallocate in [False, True]: xinv = irls( Aop, y, nouter=maxit, tolIRLS=1e-3, kind="model", preallocate=preallocate @@ -201,17 +210,22 @@ def test_IRLS_model(par): def test_IRLS_model_stopping(par): """IRLS model testing stopping criterion rtol (here the class based solver is used as cost is not returned by the function based one)""" - np.random.seed(42) - Aop = MatrixMult(np.random.randn(par["ny"], par["nx"])) + npp.random.seed(42) + A = npp.random.randn(par["ny"], par["nx"]) + par["imag"] * npp.random.randn( + par["ny"], par["nx"] + ) + Aop = MatrixMult(np.asarray(A), dtype=par["dtype"]) - x = np.zeros(par["nx"]) - x[par["nx"] // 2] = 1 - x[3] = 1 - x[par["nx"] - 4] = -1 + x = np.zeros(par["nx"]) + par["imag"] * np.zeros(par["nx"]) + x[par["nx"] // 2] = 1.0 + par["imag"] * 1.0 + x[3] = 1.0 + par["imag"] * 1.0 + x[par["nx"] - 4] = -1.0 - par["imag"] * 1.0 y = Aop * x maxit = 100 rtol = 6e-1 + kwars_solver = dict(iter_lim=5) if backend == "numpy" else dict(niter=5) + rcallback = ResidualNormCallback(rtol) irlssolve = IRLS( Aop, @@ -220,7 +234,9 @@ def test_IRLS_model_stopping(par): ], ) irlssolve.setup(y=y, epsI=0.1, tolIRLS=0, kind="model") - _ = irlssolve.run(np.zeros(Aop.shape[1], dtype=Aop.dtype), nouter=maxit, iter_lim=2) + _ = irlssolve.run( + np.zeros(Aop.shape[1], dtype=Aop.dtype), nouter=maxit, **kwars_solver + ) assert irlssolve.cost[-2] / irlssolve.cost[0] >= rtol assert irlssolve.cost[-1] / irlssolve.cost[0] < rtol @@ -231,18 +247,20 @@ def test_IRLS_model_stopping(par): @pytest.mark.parametrize("par", [(par1), (par3), (par5), (par1j), (par3j), (par5j)]) def test_MP(par): """Invert problem with MP""" - np.random.seed(42) - Aop = MatrixMult(np.random.randn(par["ny"], par["nx"])) + npp.random.seed(42) + A = npp.random.randn(par["ny"], par["nx"]) + par["imag"] * npp.random.randn( + par["ny"], par["nx"] + ) + Aop = MatrixMult(np.asarray(A), dtype=par["dtype"]) - x = np.zeros(par["nx"]) - x[par["nx"] // 2] = 1 - x[3] = 1 - x[par["nx"] - 4] = -1 + x = np.zeros(par["nx"]) + par["imag"] * np.zeros(par["nx"]) + x[par["nx"] // 2] = 1.0 + par["imag"] * 1.0 + x[3] = 1.0 + par["imag"] * 1.0 + x[par["nx"] - 4] = -1.0 - par["imag"] * 1.0 y = Aop * x sigma = 1e-4 maxit = 100 - for preallocate in [False, True]: xinv, _, _ = omp( Aop, @@ -259,18 +277,20 @@ def test_MP(par): @pytest.mark.parametrize("par", [(par1), (par3), (par5), (par1j), (par3j), (par5j)]) def test_OMP(par): """Invert problem with OMP""" - np.random.seed(42) - Aop = MatrixMult(np.random.randn(par["ny"], par["nx"])) + npp.random.seed(42) + A = npp.random.randn(par["ny"], par["nx"]) + par["imag"] * npp.random.randn( + par["ny"], par["nx"] + ) + Aop = MatrixMult(np.asarray(A), dtype=par["dtype"]) - x = np.zeros(par["nx"]) - x[par["nx"] // 2] = 1 - x[3] = 1 - x[par["nx"] - 4] = -1 + x = np.zeros(par["nx"]) + par["imag"] * np.zeros(par["nx"]) + x[par["nx"] // 2] = 1.0 + par["imag"] * 1.0 + x[3] = 1.0 + par["imag"] * 1.0 + x[par["nx"] - 4] = -1.0 - par["imag"] * 1.0 y = Aop * x sigma = 1e-4 maxit = 100 - for preallocate in [False, True]: xinv, _, _ = omp(Aop, y, maxit, sigma=sigma, preallocate=preallocate) assert_array_almost_equal(x, xinv, decimal=1) @@ -279,17 +299,19 @@ def test_OMP(par): @pytest.mark.parametrize("par", [(par1), (par3), (par5), (par1j), (par3j), (par5j)]) def test_OMP_stopping(par): """OMP testing stopping criterion rtol""" - np.random.seed(42) - Aop = MatrixMult(np.random.randn(par["ny"], par["nx"])) + npp.random.seed(42) + A = npp.random.randn(par["ny"], par["nx"]) + par["imag"] * npp.random.randn( + par["ny"], par["nx"] + ) + Aop = MatrixMult(np.asarray(A), dtype=par["dtype"]) - x = np.zeros(par["nx"]) - x[par["nx"] // 2] = 1 - x[3] = 1 - x[par["nx"] - 4] = -1 + x = np.zeros(par["nx"]) + par["imag"] * np.zeros(par["nx"]) + x[par["nx"] // 2] = 1.0 + par["imag"] * 1.0 + x[3] = 1.0 + par["imag"] * 1.0 + x[par["nx"] - 4] = -1.0 - par["imag"] * 1.0 y = Aop * x maxit = 100 - for preallocate in [False, True]: rtol = 1e-2 _, _, cost = omp(Aop, y, maxit, sigma=0.0, rtol=rtol, preallocate=preallocate) @@ -316,17 +338,21 @@ def test_ISTA_FISTA_missing_perc(): @pytest.mark.parametrize("par", [(par1), (par3), (par5), (par1j), (par3j), (par5j)]) def test_ISTA_FISTA(par): """Invert problem with ISTA/FISTA""" - np.random.seed(42) - Aop = MatrixMult(np.random.randn(par["ny"], par["nx"])) + npp.random.seed(42) + A = npp.random.randn(par["ny"], par["nx"]) + par["imag"] * npp.random.randn( + par["ny"], par["nx"] + ) + Aop = MatrixMult(np.asarray(A), dtype=par["dtype"]) - x = np.zeros(par["nx"]) - x[par["nx"] // 2] = 1 - x[3] = 1 - x[par["nx"] - 4] = -1 + x = np.zeros(par["nx"]) + par["imag"] * np.zeros(par["nx"]) + x[par["nx"] // 2] = 1.0 + par["imag"] * 1.0 + x[3] = 1.0 + par["imag"] * 1.0 + x[par["nx"] - 4] = -1.0 - par["imag"] * 1.0 y = Aop * x - eps = 0.5 - perc = 30 + # some parameters need to be tuned differently for different problem sizes + eps = 1.0 if par["ny"] >= par["nx"] else 2.0 + perc = 50 if par["ny"] >= par["nx"] else 30 maxit = 500 # ISTA with too high alpha (check that exception is raised) @@ -401,18 +427,22 @@ def test_ISTA_FISTA(par): @pytest.mark.parametrize("par", [(par1), (par3), (par5), (par1j), (par3j), (par5j)]) def test_ISTA_FISTA_multiplerhs(par): """Invert problem with ISTA/FISTA with multiple RHS""" - np.random.seed(42) - Aop = MatrixMult(np.random.randn(par["ny"], par["nx"])) + npp.random.seed(42) + A = npp.random.randn(par["ny"], par["nx"]) + par["imag"] * npp.random.randn( + par["ny"], par["nx"] + ) + Aop = MatrixMult(np.asarray(A), dtype=par["dtype"]) - x = np.zeros(par["nx"]) - x[par["nx"] // 2] = 1 - x[3] = 1 - x[par["nx"] - 4] = -1 + x = np.zeros(par["nx"]) + par["imag"] * np.zeros(par["nx"]) + x[par["nx"] // 2] = 1.0 + par["imag"] * 1.0 + x[3] = 1.0 + par["imag"] * 1.0 + x[par["nx"] - 4] = -1.0 - par["imag"] * 1.0 x = np.outer(x, np.ones(3)) y = Aop * x - eps = 0.5 - perc = 30 + # some parameters need to be tuned differently for different problem sizes + eps = 1.0 if par["ny"] >= par["nx"] else 2.0 + perc = 50 if par["ny"] >= par["nx"] else 30 maxit = 500 # Regularization based ISTA and FISTA @@ -475,23 +505,27 @@ def test_ISTA_FISTA_multiplerhs(par): @pytest.mark.parametrize("par", [(par1), (par3), (par5), (par1j), (par3j), (par5j)]) def test_ISTA_FISTA_stopping(par): """ISTA/FISTA testing stopping criterion rtol""" - np.random.seed(42) - Aop = MatrixMult(np.random.randn(par["ny"], par["nx"])) + npp.random.seed(42) + A = npp.random.randn(par["ny"], par["nx"]) + par["imag"] * npp.random.randn( + par["ny"], par["nx"] + ) + Aop = MatrixMult(np.asarray(A), dtype=par["dtype"]) - x = np.zeros(par["nx"]) - x[par["nx"] // 2] = 1 - x[3] = 1 - x[par["nx"] - 4] = -1 + x = np.zeros(par["nx"]) + par["imag"] * np.zeros(par["nx"]) + x[par["nx"] // 2] = 1.0 + par["imag"] * 1.0 + x[3] = 1.0 + par["imag"] * 1.0 + x[par["nx"] - 4] = -1.0 - par["imag"] * 1.0 y = Aop * x eps = 0.5 maxit = 500 + rtol = 5e-1 # Regularization based ISTA and FISTA threshkinds = ["hard", "soft", "half"] if backend == "numpy" else ["soft", "half"] for threshkind in threshkinds: for preallocate in [False, True]: - rtol = 5e-1 + # ISTA _, _, cost = ista( Aop, @@ -529,10 +563,13 @@ def test_ISTA_FISTA_stopping(par): ) def test_SPGL1(par): """Invert problem with SPGL1""" - np.random.seed(42) - Aop = MatrixMult(np.random.randn(par["ny"], par["nx"])) + npp.random.seed(42) + A = npp.random.normal(0, 10, (par["ny"], par["nx"])) + par[ + "imag" + ] * npp.random.normal(0, 10, (par["ny"], par["nx"])) + Aop = MatrixMult(np.asarray(A), dtype=par["dtype"]) - x = np.zeros(par["nx"]) + x = np.zeros(par["nx"]) + par["imag"] * np.zeros(par["nx"]) x[par["nx"] // 2] = 1 x[3] = 1 x[par["nx"] - 4] = -1