diff --git a/docs/source/addingsolver.rst b/docs/source/addingsolver.rst index d64cc996d..2002c6876 100755 --- a/docs/source/addingsolver.rst +++ b/docs/source/addingsolver.rst @@ -4,8 +4,8 @@ Implementing new solvers ======================== Users are welcome to create new solvers and add them to the PyLops library. -In this tutorial, we will go through the key steps in the definition of a solver, using the -:py:class:`pylops.optimization.basic.CG` as an example. +In this tutorial, we will go through the key steps in the definition of a solver, using a +sligthly simplified version of :py:class:`pylops.optimization.basic.CG` as an example. .. note:: In case the solver that you are planning to create falls within the category of proximal solvers, @@ -83,7 +83,6 @@ note that the ``setup`` method returns the created starting guess ``x`` (does no .. code-block:: python def setup(self, y, x0=None, niter=None, tol=1e-4, show=False): - self.y = y self.tol = tol self.niter = niter @@ -134,7 +133,26 @@ can add additional input parameters. For CG, the step is: Similarly, we also implement a ``run`` method that is in charge of running a number of iterations by repeatedly -calling the ``step`` method. It is also usually convenient to implement a finalize method; this method can do any required post-processing that should +calling the ``step`` method. + +.. code-block:: python + + def run(self, x, niter, show, itershow): + while self.iiter < niter and self.kold > self.tol: + x = self.step(x, showstep) + self.callback(x) + # check if any callback has raised a stop flag + stop = _callback_stop(self.callbacks) + if stop: + break + return x + +It is worth noting that any number of callbacks can be attached to the solver; some of these +callbacks can implement a stopping criterion and set the ``stop`` member to True when a given +condition is met. The ``_callback_stop`` method is in change of checking if any of the callbacks +has set ``stop`` to True and in the case break the iterations. + +Finally, it is also usually convenient to implement a ``finalize`` method; this method can do any required post-processing that should not be applied at the end of each step, rather at the end of the entire optimization process. For CG, this is as simple as converting the ``cost`` variable from a list to a ``numpy`` array. For more details, see our implementations for CG. @@ -169,8 +187,10 @@ input and returns some of the most valuable properties of the class-based solver .. code-block:: python - def cg(Op, y, x0, niter=10, tol=1e-4, show=False, itershow=(10, 10, 10), callback=None): - cgsolve = CG(Op) + def cg(Op, y, x0, niter=10, tol=1e-4, rtol=0.0, + show=False, itershow=(10, 10, 10), callback=None): + rcallback = ResidualNormCallback(rtol) + cgsolve = CG(Op, callbacks=[rcallback, ]) if callback is not None: cgsolve.callback = callback x, iiter, cost = cgsolve.solve( diff --git a/pylops/optimization/basic.py b/pylops/optimization/basic.py index f8b060d2d..32f3658c0 100644 --- a/pylops/optimization/basic.py +++ b/pylops/optimization/basic.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Callable, Optional, Tuple +from pylops.optimization.callback import ResidualNormCallback from pylops.optimization.cls_basic import CG, CGLS, LSQR from pylops.utils.decorators import add_ndarray_support_to_solver from pylops.utils.typing import NDArray @@ -21,6 +22,7 @@ def cg( x0: Optional[NDArray] = None, niter: int = 10, tol: float = 1e-4, + rtol: bool = 0.0, show: bool = False, itershow: Tuple[int, int, int] = (10, 10, 10), callback: Optional[Callable] = None, @@ -42,7 +44,12 @@ def cg( niter : :obj:`int`, optional Number of iterations tol : :obj:`float`, optional - Tolerance on residual norm + Absolute tolerance on residual norm. Stops the solver when the + residual norm is below this value. + rtol : :obj:`float`, optional + Relative tolerance on residual norm. Stops the solver when the + ratio of the current residual norm to the initial residual norm is + below this value. show : :obj:`bool`, optional Display iterations log itershow : :obj:`tuple`, optional @@ -71,7 +78,13 @@ def cg( See :class:`pylops.optimization.cls_basic.CG` """ - cgsolve = CG(Op) + rcallback = ResidualNormCallback(rtol) + cgsolve = CG( + Op, + callbacks=[ + rcallback, + ], + ) if callback is not None: cgsolve.callback = callback x, iiter, cost = cgsolve.solve( @@ -94,6 +107,7 @@ def cgls( niter: int = 10, damp: float = 0.0, tol: float = 1e-4, + rtol: float = 0.0, show: bool = False, itershow: Tuple[int, int, int] = (10, 10, 10), callback: Optional[Callable] = None, @@ -117,7 +131,12 @@ def cgls( damp : :obj:`float`, optional Damping coefficient tol : :obj:`float`, optional - Tolerance on residual norm + Absolute tolerance on residual norm. Stops the solver when the + residual norm is below this value. + rtol : :obj:`float`, optional + Relative tolerance on residual norm. Stops the solver when the + ratio of the current residual norm to the initial residual norm is + below this value. show : :obj:`bool`, optional Display iterations log itershow : :obj:`tuple`, optional @@ -161,7 +180,13 @@ def cgls( See :class:`pylops.optimization.cls_basic.CGLS` """ - cgsolve = CGLS(Op) + rcallback = ResidualNormCallback(rtol) + cgsolve = CGLS( + Op, + callbacks=[ + rcallback, + ], + ) if callback is not None: cgsolve.callback = callback x, istop, iiter, r1norm, r2norm, cost = cgsolve.solve( diff --git a/pylops/optimization/callback.py b/pylops/optimization/callback.py index 946e82696..095dd2e1e 100644 --- a/pylops/optimization/callback.py +++ b/pylops/optimization/callback.py @@ -1,6 +1,7 @@ __all__ = [ "Callbacks", "MetricsCallback", + "ResidualNormCallback", ] from typing import TYPE_CHECKING, Dict, List, Optional, Sequence @@ -28,25 +29,32 @@ class Callbacks: All methods take two input parameters: the solver itself, and the vector ``x``. + Moreover, some callback may be used to implement custom stopping criteria for the solver. + This can be done by adding a boolean attribute ``stop`` to the callback object, which will + be initially set to ``False``. As soon as the callback sets this attribute to ``True``, the + ``run`` method of the solver will stop iterating and return the current model vector. + Examples -------- >>> import numpy as np >>> from pylops.basicoperators import MatrixMult - >>> from pylops.optimization.solver import CG + >>> from pylops.optimization.basic import CG >>> from pylops.optimization.callback import Callbacks + >>> >>> class StoreIterCallback(Callbacks): ... def __init__(self): ... self.stored = [] ... def on_step_end(self, solver, x): ... self.stored.append(solver.iiter) - >>> cb_sto = StoreIterCallback() + >>> >>> Aop = MatrixMult(np.random.normal(0., 1., 36).reshape(6, 6)) >>> Aop = Aop.H @ Aop >>> y = Aop @ np.ones(6) + >>> cb_sto = StoreIterCallback() >>> cgsolve = CG(Aop, callbacks=[cb_sto, ]) >>> xest = cgsolve.solve(y=y, x0=np.zeros(6), tol=0, niter=6, show=False)[0] - >>> xest - array([1., 1., 1., 1., 1., 1.]) + >>> xest, cb_sto.stored + (array([1., 1., 1., 1., 1., 1.]), [1, 2, 3, 4, 5, 6]) """ @@ -181,3 +189,53 @@ def on_step_end(self, solver: "Solver", x: NDArray) -> None: self.metrics["snr"].append(snr(self.xtrue, x)) if "psnr" in self.which: self.metrics["psnr"].append(psnr(self.xtrue, x)) + + +class ResidualNormCallback(Callbacks): + """Residual norm callback + + This callback can be used to stop the solver when the residual norm + is below a certain threshold defined as a percentage of the + initial residual norm. + + Parameters + ---------- + rtol : :obj:`float` + Percentage of the initial residual norm below which the solver + will stop iterating. For example, if `rtol` is 0.1, the solver + will stop when the residual norm is below 10% of the initial + residual norm. + + """ + + def __init__(self, rtol: float) -> None: + self.rtol = rtol + self.stop = False + + def on_step_end(self, solver: "Solver", x: NDArray) -> None: + if solver.cost[-1] < self.rtol * solver.cost[0]: + self.stop = True + + +def _callback_stop(callbacks: Sequence[Callbacks]) -> bool: + """Check if any callback has raised a stop flag + + Parameters + ---------- + callbacks : :obj:`pylops.optimization.callback.Callbacks` + List of callbacks to evaluate + + Returns + ------- + stop : :obj:`bool` + Whether to stop the solver or not + + """ + if callbacks is not None: + stop = [ + False if not hasattr(callback, "stop") else callback.stop + for callback in callbacks + ] + if any(stop): + return True + return False diff --git a/pylops/optimization/cls_basic.py b/pylops/optimization/cls_basic.py index 0713773ea..843b9596b 100644 --- a/pylops/optimization/cls_basic.py +++ b/pylops/optimization/cls_basic.py @@ -10,6 +10,7 @@ import numpy as np from pylops.optimization.basesolver import Solver, _units +from pylops.optimization.callback import _callback_stop from pylops.utils.backend import ( get_array_module, get_module_name, @@ -121,7 +122,8 @@ def setup( Number of iterations (default to ``None`` in case a user wants to manually step over the solver) tol : :obj:`float`, optional - Tolerance on residual norm + Absolute tolerance on residual norm. Stops the solver when the + residual norm is below this value. preallocate : :obj:`bool`, optional .. versionadded:: 2.6.0 @@ -260,6 +262,10 @@ def run( ) x = self.step(x, showstep) self.callback(x) + # check if any callback has raised a stop flag + stop = _callback_stop(self.callbacks) + if stop: + break return x def finalize(self, show: bool = False) -> None: @@ -299,7 +305,8 @@ def solve( niter : :obj:`int`, optional Number of iterations tol : :obj:`float`, optional - Tolerance on residual norm + Absolute tolerance on residual norm. Stops the solver when the + residual norm is below this value. preallocate : :obj:`bool`, optional .. versionadded:: 2.6.0 @@ -440,7 +447,8 @@ def setup( damp : :obj:`float`, optional Damping coefficient tol : :obj:`float`, optional - Tolerance on residual norm + Absolute tolerance on residual norm. Stops the solver when the + residual norm is below this value. preallocate : :obj:`bool`, optional .. versionadded:: 2.6.0 @@ -592,22 +600,26 @@ def run( Estimated model of size :math:`[M \times 1]` """ - niter = self.niter if niter is None else niter - if niter is None: + self.niter = self.niter if niter is None else niter + if self.niter is None: raise ValueError("niter must not be None") - while self.iiter < niter and self.kold > self.tol: + while self.iiter < self.niter and self.kold > self.tol: showstep = ( True if show and ( self.iiter < itershow[0] - or niter - self.iiter < itershow[1] + or self.niter - self.iiter < itershow[1] or self.iiter % itershow[2] == 0 ) else False ) x = self.step(x, showstep) self.callback(x) + # check if any callback has raised a stop flag + stop = _callback_stop(self.callbacks) + if stop: + break return x def finalize(self, show: bool = False) -> None: @@ -622,7 +634,12 @@ def finalize(self, show: bool = False) -> None: self.tend = time.time() self.telapsed = self.tend - self.tstart # reason for termination - self.istop = 1 if self.kold < self.tol else 2 + if self.kold < self.tol: + self.istop = 1 + elif self.iiter >= self.niter: + self.istop = 2 + else: + self.istop = 3 self.r1norm = self.kold self.r2norm = self.cost1[self.iiter] if show: @@ -655,7 +672,8 @@ def solve( damp : :obj:`float`, optional Damping coefficient tol : :obj:`float`, optional - Tolerance on residual norm + Absolute tolerance on residual norm. Stops the solver when the + residual norm is below this value. preallocate : :obj:`bool`, optional .. versionadded:: 2.6.0 @@ -677,10 +695,14 @@ def solve( Gives the reason for termination ``1`` means :math:`\mathbf{x}` is an approximate solution to - :math:`\mathbf{y} = \mathbf{Op}\,\mathbf{x}` + :math:`\mathbf{y} = \mathbf{Op}\,\mathbf{x}` with the provided + tolerance ``tol`` ``2`` means :math:`\mathbf{x}` approximately solves the least-squares - problem + problem (reached the maximum number of iterations ``niter``) + + ``3`` means another stopping criterion implemented via a callback + was reached iit : :obj:`int` Iteration number upon termination r1norm : :obj:`float` diff --git a/pylops/optimization/cls_leastsquares.py b/pylops/optimization/cls_leastsquares.py index 557997a10..4cd51ee61 100644 --- a/pylops/optimization/cls_leastsquares.py +++ b/pylops/optimization/cls_leastsquares.py @@ -9,7 +9,7 @@ import numpy as np from scipy.sparse.linalg import cg as sp_cg -from scipy.sparse.linalg import lsqr +from scipy.sparse.linalg import lsqr as sp_lsqr from pylops.basicoperators import Diagonal, VStack from pylops.optimization.basesolver import Solver, _units @@ -676,7 +676,7 @@ def run( if engine == "scipy" and self.ncp == np: if show: kwargs_solver["show"] = 1 - xinv, istop, itn, r1norm, r2norm = lsqr( + xinv, istop, itn, r1norm, r2norm = sp_lsqr( self.RegOp, self.datatot, x0=x, **kwargs_solver )[0:5] elif engine == "pylops" or self.ncp != np: @@ -938,7 +938,7 @@ def run( if engine == "scipy" and self.ncp == np: if show: kwargs_solver["show"] = 1 - pinv, istop, itn, r1norm, r2norm = lsqr( + pinv, istop, itn, r1norm, r2norm = sp_lsqr( self.POp, self.y, x0=x, diff --git a/pylops/optimization/cls_sparsity.py b/pylops/optimization/cls_sparsity.py index 093efc856..59ce8a86e 100644 --- a/pylops/optimization/cls_sparsity.py +++ b/pylops/optimization/cls_sparsity.py @@ -19,6 +19,7 @@ from pylops.basicoperators import Diagonal, Identity, VStack from pylops.optimization.basesolver import Solver, _units from pylops.optimization.basic import cgls +from pylops.optimization.callback import _callback_stop from pylops.optimization.eigs import power_iteration from pylops.optimization.leastsquares import regularized_inversion from pylops.utils import deps @@ -399,14 +400,14 @@ def setup( epsR : :obj:`float`, optional Damping to be applied to residuals for weighting term epsI : :obj:`float`, optional - Tikhonov damping (for ``kind="data"``) or L1 model damping - (for ``kind="datamodel"``) + Tikhonov damping tolIRLS : :obj:`float`, optional Tolerance. Stop outer iterations if difference between inverted model at subsequent iterations is smaller than ``tolIRLS`` warm : :obj:`bool`, optional - Warm start each inversion inner step with previous estimate (``True``) or not (``False``). - This only applies to ``kind="data"`` and ``kind="datamodel"`` + Warm start each inversion inner step with previous estimate (``True``) + or not (``False``). This only applies to ``kind="data"`` and + ``kind="datamodel"`` kind : :obj:`str`, optional Kind of solver (``model``, ``data`` or ``datamodel``) preallocate : :obj:`bool`, optional @@ -432,9 +433,6 @@ def setup( self.isjax = get_module_name(self.ncp) == "jax" self._setpreallocate(preallocate) - # initiate outer iteration counter - self.iiter = 0 - # choose step to use if self.kind == "data": self._step = self._step_data @@ -456,6 +454,13 @@ def setup( self.rw = self.ncp.empty_like(self.y) else: self.rw = self.ncp.empty(self.Op.shape[1], dtype=self.Op.dtype) + + # create variables to track the residual norm and iterations + self.cost = [ + float(np.linalg.norm(self.y)), + ] + self.iiter = 0 + # print setup if show: self._print_setup() @@ -619,6 +624,7 @@ def step( self.rnorm = self.ncp.linalg.norm(self.r) self.iiter += 1 + self.cost.append(float(self.rnorm)) if show: self._print_step(x) return x @@ -687,6 +693,10 @@ def run( xold = x.copy() x = self.step(x, engine, showstep, **kwargs_solver) self.callback(x) + # check if any callback has raised a stop flag + stop = _callback_stop(self.callbacks) + if stop: + break # adding initial guess if hasattr(self, "x0"): @@ -1134,7 +1144,7 @@ def step( self.ncp.subtract(self.res, self.y, out=self.res) self.iiter += 1 - self.cost.append(float(np.linalg.norm(self.res))) + self.cost.append(float(self.ncp.linalg.norm(self.res))) if show: self._print_step(x) return x, cols @@ -1187,6 +1197,10 @@ def run( ) x, cols = self.step(x, cols, engine, showstep) self.callback(x, cols) + # check if any callback has raised a stop flag + stop = _callback_stop(self.callbacks) + if stop: + break return x, cols def finalize( @@ -1757,7 +1771,7 @@ def step(self, x: NDArray, show: bool = False) -> Tuple[NDArray, float]: if self.SOp is not None: x = self.SOpmatvec(x) - # check model update + # compute model update norm if not self.preallocate: xupdate = np.linalg.norm(x - xold) else: @@ -1768,7 +1782,7 @@ def step(self, x: NDArray, show: bool = False) -> Tuple[NDArray, float]: ) xupdate = np.linalg.norm(self.xold) - # cost functions + # compute cost functions costdata = 0.5 * np.linalg.norm(self.res if self.preallocate else res) ** 2 costreg = self.eps * np.linalg.norm(x, ord=1) self.cost.append(float(costdata + costreg)) @@ -1824,6 +1838,10 @@ def run( ) x, xupdate = self.step(x, showstep) self.callback(x) + # check if any callback has raised a stop flag + stop = _callback_stop(self.callbacks) + if stop: + break if xupdate <= self.tol: logger.info("Update smaller that tolerance for iteration %d", self.iiter) return x @@ -2205,6 +2223,10 @@ def run( ) x, z, xupdate = self.step(x, z, showstep) self.callback(x) + # check if any callback has raised a stop flag + stop = _callback_stop(self.callbacks) + if stop: + break if xupdate <= self.tol: logger.warning( "Update smaller that tolerance for " "iteration %d", self.iiter @@ -2943,7 +2965,10 @@ def run( ) x = self.step(x, engine, showstep, show_inner, **kwargs_lsqr) self.callback(x) - + # check if any callback has raised a stop flag + stop = _callback_stop(self.callbacks) + if stop: + break return x def finalize(self, show: bool = False) -> NDArray: diff --git a/pylops/optimization/sparsity.py b/pylops/optimization/sparsity.py index c7c720efb..54e155d72 100644 --- a/pylops/optimization/sparsity.py +++ b/pylops/optimization/sparsity.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple +from pylops.optimization.callback import ResidualNormCallback from pylops.optimization.cls_sparsity import FISTA, IRLS, ISTA, OMP, SPGL1, SplitBregman from pylops.utils.decorators import add_ndarray_support_to_solver from pylops.utils.typing import NDArray, SamplingLike @@ -140,6 +141,7 @@ def omp( niter_outer: int = 10, niter_inner: int = 40, sigma: float = 1e-4, + rtol: float = 0.0, normalizecols: bool = False, Opbasis: Optional["LinearOperator"] = None, optimal_coeff: bool = False, @@ -167,8 +169,12 @@ def omp( niter_inner : :obj:`int`, optional Number of iterations of inner loop. By choosing ``niter_inner=0``, the Matching Pursuit (MP) algorithm is implemented. - sigma : :obj:`list` + sigma : :obj:`float`, optional Maximum :math:`L_2` norm of residual. When smaller stop iterations. + rtol : :obj:`float`, optional + Relative tolerance on residual. Stops the solver when the + ratio of the current residual norm to the initial residual norm + is below this value. normalizecols : :obj:`list`, optional Normalize columns (``True``) or not (``False``). Note that this can be expensive as it requires applying the forward operator @@ -201,6 +207,7 @@ def omp( Pre-allocate all variables used by the solver. Note that if ``y`` is a JAX array, this option is ignored and variables are not pre-allocated since JAX does not support in-place operations. + Returns ------- xinv : :obj:`numpy.ndarray` @@ -222,7 +229,13 @@ def omp( See :class:`pylops.optimization.cls_sparsity.OMP` """ - ompsolve = OMP(Op) + rcallback = ResidualNormCallback(rtol) + ompsolve = OMP( + Op, + callbacks=[ + rcallback, + ], + ) if callback is not None: ompsolve.callback = callback x, niter_outer, cost = ompsolve.solve( @@ -251,6 +264,7 @@ def ista( alpha: Optional[float] = None, eigsdict: Optional[Dict[str, Any]] = None, tol: float = 1e-10, + rtol: bool = 0.0, threshkind: str = "soft", perc: Optional[float] = None, decay: Optional[NDArray] = None, @@ -292,8 +306,12 @@ def ista( Dictionary of parameters to be passed to :func:`pylops.LinearOperator.eigs` method when computing the maximum eigenvalue tol : :obj:`float`, optional - Tolerance. Stop iterations if difference between inverted model + Absolute tolerance on model update. Stop iterations if difference between inverted model at subsequent iterations is smaller than ``tol`` + rtol : :obj:`float`, optional + Relative tolerance on total cost function. Stops the solver when the + ratio of the current cost function to the initial cost function + is below this value. threshkind : :obj:`str`, optional Kind of thresholding ('hard', 'soft', 'half', 'hard-percentile', 'soft-percentile', or 'half-percentile' - 'soft' used as default) @@ -327,7 +345,7 @@ def ista( niter : :obj:`int` Number of effective iterations cost : :obj:`numpy.ndarray` - History of cost function + History of cost (including regularization term) Raises ------ @@ -352,7 +370,13 @@ def ista( See :class:`pylops.optimization.cls_sparsity.ISTA` """ - istasolve = ISTA(Op) + rcallback = ResidualNormCallback(rtol) + istasolve = ISTA( + Op, + callbacks=[ + rcallback, + ], + ) if callback is not None: istasolve.callback = callback x, iiter, cost = istasolve.solve( @@ -385,6 +409,7 @@ def fista( alpha: Optional[float] = None, eigsdict: Optional[Dict[str, Any]] = None, tol: float = 1e-10, + rtol: float = 0.0, threshkind: str = "soft", perc: Optional[float] = None, decay: Optional[NDArray] = None, @@ -426,8 +451,12 @@ def fista( Dictionary of parameters to be passed to :func:`pylops.LinearOperator.eigs` method when computing the maximum eigenvalue tol : :obj:`float`, optional - Tolerance. Stop iterations if difference between inverted model + Absolute tolerance on model update. Stop iterations if difference between inverted model at subsequent iterations is smaller than ``tol`` + rtol : :obj:`float`, optional + Relative tolerance on total cost function. Stops the solver when the + ratio of the current cost function to the initial cost function + is below this value. threshkind : :obj:`str`, optional Kind of thresholding ('hard', 'soft', 'half', 'soft-percentile', or 'half-percentile' - 'soft' used as default) @@ -461,7 +490,7 @@ def fista( niter : :obj:`int` Number of effective iterations cost : :obj:`numpy.ndarray`, optional - History of cost function + History of cost (including regularization term) Raises ------ @@ -484,7 +513,13 @@ def fista( See :class:`pylops.optimization.cls_sparsity.FISTA` """ - fistasolve = FISTA(Op) + rcallback = ResidualNormCallback(rtol) + fistasolve = FISTA( + Op, + callbacks=[ + rcallback, + ], + ) if callback is not None: fistasolve.callback = callback x, iiter, cost = fistasolve.solve( @@ -637,6 +672,7 @@ def splitbregman( epsRL1s: Optional[SamplingLike] = None, epsRL2s: Optional[SamplingLike] = None, tol: float = 1e-10, + rtol: float = 0.0, tau: float = 1.0, restart: bool = False, engine: str = "scipy", @@ -687,8 +723,12 @@ def splitbregman( :math:`L_2` Regularization dampings (must have the same number of elements as ``RegsL2``) tol : :obj:`float`, optional - Tolerance. Stop outer iterations if difference between inverted model + Tolerance. Stop the solver if difference between inverted model at subsequent iterations is smaller than ``tol`` + rtol : :obj:`float`, optional + Relative tolerance on total cost function. Stops the solver when the + ratio of the current cost function to the initial cost function + is below this value. tau : :obj:`float`, optional Scaling factor in the Bregman update (must be close to 1) restart : :obj:`bool`, optional @@ -732,7 +772,13 @@ def splitbregman( See :class:`pylops.optimization.cls_sparsity.SplitBregman` """ - sbsolve = SplitBregman(Op) + rcallback = ResidualNormCallback(rtol) + sbsolve = SplitBregman( + Op, + callbacks=[ + rcallback, + ], + ) if callback is not None: sbsolve.callback = callback xinv, itn_out, cost = sbsolve.solve( diff --git a/pytests/test_memoizeoperator.py b/pytests/test_memoizeoperator.py index de3190a35..d249b45a5 100644 --- a/pytests/test_memoizeoperator.py +++ b/pytests/test_memoizeoperator.py @@ -10,8 +10,6 @@ from numpy.testing import assert_array_almost_equal backend = "numpy" -import itertools - import pytest from pylops import MemoizeOperator @@ -73,7 +71,7 @@ def test_memoize_evals_2(par): # Approach 1 Aop1 = Aop.toreal(forw=False, adj=True) - xinv1 = Aop1.div(y) + xinv1 = Aop1.div(y).real assert_array_almost_equal(x, xinv1) # Approach 2 diff --git a/pytests/test_solver.py b/pytests/test_solver.py index 79a7ba23d..a59c08df1 100644 --- a/pytests/test_solver.py +++ b/pytests/test_solver.py @@ -164,6 +164,38 @@ def test_cg_forceflat(par): assert_array_almost_equal(x.ravel(), xinv, decimal=4) +@pytest.mark.parametrize( + "par", [(par1), (par2), (par3), (par4), (par1j), (par2j), (par3j), (par3j)] +) +def test_cg_stopping(par): + """CG testing stopping criterion rtol""" + np.random.seed(10) + + A = np.random.normal(0, 10, (par["ny"], par["nx"])) + par[ + "imag" + ] * np.random.normal(0, 10, (par["ny"], par["nx"])) + A = np.conj(A).T @ A # to ensure definite positive matrix + Aop = MatrixMult(A, dtype=par["dtype"]) + + x = np.ones(par["nx"]) + par["imag"] * np.ones(par["nx"]) + if par["x0"]: + x0 = np.random.normal(0, 10, par["nx"]) + par["imag"] * np.random.normal( + 0, 10, par["nx"] + ) + else: + x0 = None + + y = Aop * x + + for preallocate in [False, True]: + rtol = 1e-2 + _, _, cost = cg( + Aop, y, x0=x0, niter=par["nx"], tol=0, rtol=rtol, preallocate=preallocate + ) + assert cost[-2] / cost[0] >= rtol + assert cost[-1] / cost[0] < rtol + + @pytest.mark.parametrize( "par", [(par1), (par2), (par3), (par4), (par1j), (par2j), (par3j), (par3j)] ) @@ -193,14 +225,105 @@ def test_cgls(par): assert_array_almost_equal(x, xinv, decimal=4) +@pytest.mark.parametrize( + "par", [(par1), (par2), (par3), (par4), (par1j), (par2j), (par3j), (par3j)] +) +def test_cgls_ndarray(par): + """CGLS with linear operator (and ndarray as input and output)""" + np.random.seed(10) + + dims = dimsd = (par["nx"], par["ny"]) + x = np.ones(dims) + par["imag"] * np.ones(dims) + + A = np.random.normal(0, 10, (x.size, x.size)) + par["imag"] * np.random.normal( + 0, 10, (x.size, x.size) + ) + Aop = MatrixMult(A, dtype=par["dtype"]) + Aop.dims = dims + Aop.dimsd = dimsd + + if par["x0"]: + x0 = np.random.normal(0, 10, dims) + par["imag"] * np.random.normal(0, 10, dims) + else: + x0 = None + + y = Aop * x + + for preallocate in [False, True]: + xinv = cgls(Aop, y, x0=x0, niter=2 * x.size, tol=0, preallocate=preallocate)[0] + assert xinv.shape == x.shape + assert_array_almost_equal(x, xinv, decimal=4) + + +@pytest.mark.parametrize( + "par", [(par1), (par2), (par3), (par4), (par1j), (par2j), (par3j), (par3j)] +) +def test_cgls_forceflat(par): + """CGLS with linear operator (and forced 1darray as input and output)""" + np.random.seed(10) + + dims = dimsd = (par["nx"], par["ny"]) + x = np.ones(dims) + par["imag"] * np.ones(dims) + + A = np.random.normal(0, 10, (x.size, x.size)) + par["imag"] * np.random.normal( + 0, 10, (x.size, x.size) + ) + Aop = MatrixMult(A, dtype=par["dtype"], forceflat=True) + Aop.dims = dims + Aop.dimsd = dimsd + + if par["x0"]: + x0 = np.random.normal(0, 10, dims) + par["imag"] * np.random.normal(0, 10, dims) + else: + x0 = None + + y = Aop * x + + for preallocate in [False, True]: + xinv = cgls(Aop, y, x0=x0, niter=2 * x.size, tol=0, preallocate=preallocate)[0] + assert xinv.shape == x.ravel().shape + assert_array_almost_equal(x.ravel(), xinv, decimal=4) + + +@pytest.mark.parametrize( + "par", [(par1), (par2), (par3), (par4), (par1j), (par2j), (par3j), (par3j)] +) +def test_cgls_stopping(par): + """CGLS testing stopping criterion rtol""" + np.random.seed(10) + + A = np.random.normal(0, 10, (par["ny"], par["nx"])) + par[ + "imag" + ] * np.random.normal(0, 10, (par["ny"], par["nx"])) + Aop = MatrixMult(A, dtype=par["dtype"]) + + x = np.ones(par["nx"]) + par["imag"] * np.ones(par["nx"]) + if par["x0"]: + x0 = np.random.normal(0, 10, par["nx"]) + par["imag"] * np.random.normal( + 0, 10, par["nx"] + ) + else: + x0 = None + + y = Aop * x + + for preallocate in [False, True]: + rtol = 1e-2 + cost = cgls( + Aop, y, x0=x0, niter=par["nx"], tol=0, rtol=rtol, preallocate=preallocate + )[-1] + assert cost[-2] / cost[0] >= rtol + assert cost[-1] / cost[0] < rtol + + @pytest.mark.skipif( int(os.environ.get("TEST_CUPY_PYLOPS", 0)) == 1, reason="Not CuPy enabled" ) @pytest.mark.parametrize( "par", [(par1), (par2), (par3), (par4), (par1j), (par2j), (par3j), (par3j)] ) -def test_lsqr(par): - """Compare local Pylops and scipy LSQR""" +def test_lsqr_pylops_scipy(par): + """Compare Pylops and scipy LSQR""" np.random.seed(10) A = np.random.normal(0, 10, (par["ny"], par["nx"])) + par[ @@ -230,3 +353,92 @@ def test_lsqr(par): assert_array_almost_equal(xinv, x, decimal=4) assert_array_almost_equal(xinv_sp, x, decimal=4) + + +@pytest.mark.parametrize( + "par", [(par1), (par2), (par3), (par4), (par1j), (par2j), (par3j), (par3j)] +) +def test_lsqr(par): + """LSQR with linear operator""" + np.random.seed(10) + + A = np.random.normal(0, 10, (par["ny"], par["nx"])) + par[ + "imag" + ] * np.random.normal(0, 10, (par["ny"], par["nx"])) + Aop = MatrixMult(A, dtype=par["dtype"]) + + x = np.ones(par["nx"]) + par["imag"] * np.ones(par["nx"]) + if par["x0"]: + x0 = np.random.normal(0, 10, par["nx"]) + par["imag"] * np.random.normal( + 0, 10, par["nx"] + ) + else: + x0 = None + + y = Aop * x + + for preallocate in [False, True]: + xinv = lsqr(Aop, y, x0=x0, niter=par["nx"], atol=1e-5, preallocate=preallocate)[ + 0 + ] + assert_array_almost_equal(x, xinv, decimal=4) + + +@pytest.mark.parametrize( + "par", [(par1), (par2), (par3), (par4), (par1j), (par2j), (par3j), (par3j)] +) +def test_lsqr_ndarray(par): + """LSQR with linear operator (and ndarray as input and output)""" + np.random.seed(10) + + dims = dimsd = (par["nx"], par["ny"]) + x = np.ones(dims) + par["imag"] * np.ones(dims) + + A = np.random.normal(0, 10, (x.size, x.size)) + par["imag"] * np.random.normal( + 0, 10, (x.size, x.size) + ) + Aop = MatrixMult(A, dtype=par["dtype"]) + Aop.dims = dims + Aop.dimsd = dimsd + + if par["x0"]: + x0 = np.random.normal(0, 10, dims) + par["imag"] * np.random.normal(0, 10, dims) + else: + x0 = None + + y = Aop * x + + for preallocate in [False, True]: + xinv = lsqr(Aop, y, x0=x0, niter=2 * x.size, atol=0, preallocate=preallocate)[0] + assert xinv.shape == x.shape + assert_array_almost_equal(x, xinv, decimal=4) + + +@pytest.mark.parametrize( + "par", [(par1), (par2), (par3), (par4), (par1j), (par2j), (par3j), (par3j)] +) +def test_lsqr_forceflat(par): + """LSQR with linear operator (and forced 1darray as input and output)""" + np.random.seed(10) + + dims = dimsd = (par["nx"], par["ny"]) + x = np.ones(dims) + par["imag"] * np.ones(dims) + + A = np.random.normal(0, 10, (x.size, x.size)) + par["imag"] * np.random.normal( + 0, 10, (x.size, x.size) + ) + Aop = MatrixMult(A, dtype=par["dtype"], forceflat=True) + Aop.dims = dims + Aop.dimsd = dimsd + + if par["x0"]: + x0 = np.random.normal(0, 10, dims) + par["imag"] * np.random.normal(0, 10, dims) + else: + x0 = None + + y = Aop * x + + for preallocate in [False, True]: + xinv = lsqr(Aop, y, x0=x0, niter=2 * x.size, atol=0, preallocate=preallocate)[0] + assert xinv.shape == x.ravel().shape + assert_array_almost_equal(x.ravel(), xinv, decimal=4) diff --git a/pytests/test_sparsity.py b/pytests/test_sparsity.py index a7092ec7f..ffa11871c 100644 --- a/pytests/test_sparsity.py +++ b/pytests/test_sparsity.py @@ -13,6 +13,8 @@ import pytest from pylops.basicoperators import FirstDerivative, Identity, MatrixMult +from pylops.optimization.callback import ResidualNormCallback +from pylops.optimization.cls_sparsity import IRLS from pylops.optimization.sparsity import fista, irls, ista, omp, spgl1, splitbregman # currently test spgl1 only if numpy<2.0.0 is installed... @@ -90,6 +92,12 @@ } # underdetermined complex, non-zero initial guess +def test_IRLS_unknown_kind(): + """Check error is raised if unknown kind is passed""" + with pytest.raises(NotImplementedError): + _ = irls(Identity(5), np.ones(5), 10, kind="foo") + + @pytest.mark.parametrize("par", [(par3), (par4), (par3j), (par4j)]) def test_IRLS_data(par): """Invert problem with outliers using data IRLS""" @@ -189,6 +197,34 @@ def test_IRLS_model(par): assert_array_almost_equal(x, xinv, decimal=1) +@pytest.mark.parametrize("par", [(par1), (par3), (par5), (par1j), (par3j), (par5j)]) +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"])) + + x = np.zeros(par["nx"]) + x[par["nx"] // 2] = 1 + x[3] = 1 + x[par["nx"] - 4] = -1 + y = Aop * x + + maxit = 100 + rtol = 6e-1 + rcallback = ResidualNormCallback(rtol) + irlssolve = IRLS( + Aop, + callbacks=[ + rcallback, + ], + ) + 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) + assert irlssolve.cost[-2] / irlssolve.cost[0] >= rtol + assert irlssolve.cost[-1] / irlssolve.cost[0] < rtol + + @pytest.mark.skipif( int(os.environ.get("TEST_CUPY_PYLOPS", 0)) == 1, reason="Not CuPy enabled" ) @@ -240,6 +276,27 @@ def test_OMP(par): assert_array_almost_equal(x, xinv, decimal=1) +@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"])) + + x = np.zeros(par["nx"]) + x[par["nx"] // 2] = 1 + x[3] = 1 + x[par["nx"] - 4] = -1 + 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) + assert cost[-2] / cost[0] >= rtol + assert cost[-1] / cost[0] < rtol + + def test_ISTA_FISTA_unknown_threshkind(): """Check error is raised if unknown threshkind is passed""" with pytest.raises(NotImplementedError): @@ -270,7 +327,7 @@ def test_ISTA_FISTA(par): eps = 0.5 perc = 30 - maxit = 2000 + maxit = 500 # ISTA with too high alpha (check that exception is raised) with pytest.raises(ValueError): @@ -356,7 +413,7 @@ def test_ISTA_FISTA_multiplerhs(par): eps = 0.5 perc = 30 - maxit = 2000 + maxit = 500 # Regularization based ISTA and FISTA threshkinds = ["hard", "soft", "half"] if backend == "numpy" else ["soft", "half"] @@ -415,6 +472,55 @@ def test_ISTA_FISTA_multiplerhs(par): assert_array_almost_equal(x, xinv, decimal=1) +@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"])) + + x = np.zeros(par["nx"]) + x[par["nx"] // 2] = 1 + x[3] = 1 + x[par["nx"] - 4] = -1 + y = Aop * x + + eps = 0.5 + maxit = 500 + + # 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, + y, + niter=maxit, + eps=eps, + threshkind=threshkind, + tol=0.0, + rtol=rtol, + preallocate=preallocate, + ) + assert cost[-2] / cost[0] >= rtol + assert cost[-1] / cost[0] < rtol + + # FISTA + _, _, cost = fista( + Aop, + y, + niter=maxit, + eps=eps, + threshkind=threshkind, + tol=0.0, + rtol=rtol, + preallocate=preallocate, + ) + assert cost[-2] / cost[0] >= rtol + assert cost[-1] / cost[0] < rtol + + @pytest.mark.skipif( int(os.environ.get("TEST_CUPY_PYLOPS", 0)) == 1, reason="Not CuPy enabled" )