doc: added tests and documentation about rtol in solvers

Author: mrava87

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,
@@ -134,7 +134,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 +188,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(

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)