feat: added rtol to solvers

Author: mrava87

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(

pylops/optimization/callback.py

@@ -1,6 +1,7 @@
 __all__ = [
     "Callbacks",
     "MetricsCallback",
+    "ResidualNormCallback",
 ]
 
 from typing import TYPE_CHECKING, Dict, List, Optional, Sequence
@@ -28,6 +29,11 @@ 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 solvers.
+    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
@@ -181,3 +187,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

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:
@@ -592,22 +598,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 +632,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:
@@ -677,10 +692,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`

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,

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(
@@ -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: