Address review: align CG damp with other solvers

- shorten damp docstrings to 'Damping coefficient' (cls_basic CG.setup/solve,
  basic.cg) and simplify CG Notes to a single damped-system sentence
- follow preallocate vs non-preallocate residual-update pattern for the
  initial damping correction in setup
- shorten the damping comment in step to '# add damping contribution'
- test_cg_damp: build y from a known x (y = Aop * x), rename loop var to
  'preallocate', drop issue ref from docstring

Author: gaoflow

pylops/optimization/basic.py

@@ -47,10 +47,7 @@ def cg(
     niter : :obj:`int`, optional
         Number of iterations
     damp : :obj:`float`, optional
-        Damping coefficient. When non-zero, the damped system
-        :math:`(\mathbf{Op} + \epsilon\mathbf{I})\,\mathbf{x} = \mathbf{y}` is
-        solved instead of :math:`\mathbf{Op}\,\mathbf{x} = \mathbf{y}`. ``Op``
-        must be square and symmetric positive-definite.
+        Damping coefficient
     tol : :obj:`float`, optional
         Absolute tolerance on residual norm. Stops the solver when the
         residual norm is below this value.

pylops/optimization/cls_basic.py

@@ -65,15 +65,9 @@ class CG(Solver):
 
     Notes
     -----
-    Solve the :math:`\mathbf{y} = \mathbf{Op}\,\mathbf{x}` problem using conjugate gradient
-    iterations [1]_.
-
-    When a non-zero damping coefficient :math:`\epsilon` is provided to ``setup``/``solve``,
-    the damped system :math:`(\mathbf{Op} + \epsilon\mathbf{I})\,\mathbf{x} = \mathbf{y}` is
-    solved instead; this is achieved by adding :math:`\epsilon\mathbf{c}` to every operator
-    application :math:`\mathbf{Op}\,\mathbf{c}` in the iterations. ``Op`` is still required to
-    be square and symmetric positive-definite (with :math:`\epsilon \geq 0` the damped
-    operator remains so).
+    Solve the :math:`\mathbf{y} = (\mathbf{Op} + \epsilon\mathbf{I})\,\mathbf{x}` problem
+    using conjugate gradient iterations [1]_, where :math:`\epsilon` is the damping
+    coefficient.
 
     .. [1] Hestenes, M R., Stiefel, E., “Methods of Conjugate Gradients for Solving
        Linear Systems”, Journal of Research of the National Bureau of Standards.
@@ -159,9 +153,7 @@ def setup(
             Number of iterations (default to ``None`` in case a user wants to
             manually step over the solver)
         damp : :obj:`float`, optional
-            Damping coefficient. When non-zero, the damped system
-            :math:`(\mathbf{Op} + \epsilon\mathbf{I})\,\mathbf{x} = \mathbf{y}`
-            is solved instead of :math:`\mathbf{Op}\,\mathbf{x} = \mathbf{y}`.
+            Damping coefficient
         tol : :obj:`float`, optional
             Absolute tolerance on residual norm. Stops the solver when the
             residual norm is below this value.
@@ -202,7 +194,10 @@ def setup(
                 self.ncp.subtract(self.y, self.Op.matvec(x), out=self.r)
             # account for the damping term in the initial residual
             if self.damp != 0.0:
-                self.r -= self.damp * x
+                if not self.preallocate:
+                    self.r = self.r - self.damp * x
+                else:
+                    self.ncp.subtract(self.r, self.damp * x, out=self.r)
         self.c = self.r.copy()
         self.kold = self.ncp.abs(self.r.dot(self.r.conj()))
 
@@ -237,8 +232,7 @@ def step(self, x: NDArray, show: bool = False) -> NDArray:
 
         """
         Opc = self.Op.matvec(self.c)
-        # solve the damped system (Op + damp*I) x = y by adding the damping
-        # contribution to every application of the operator
+        # add damping contribution
         if self.damp != 0.0:
             Opc = Opc + self.damp * self.c
         cOpc = self.ncp.abs(self.c.dot(Opc.conj()))
@@ -355,9 +349,7 @@ def solve(
         niter : :obj:`int`, optional
             Number of iterations
         damp : :obj:`float`, optional
-            Damping coefficient. When non-zero, the damped system
-            :math:`(\mathbf{Op} + \epsilon\mathbf{I})\,\mathbf{x} = \mathbf{y}`
-            is solved instead of :math:`\mathbf{Op}\,\mathbf{x} = \mathbf{y}`.
+            Damping coefficient
         tol : :obj:`float`, optional
             Absolute tolerance on residual norm. Stops the solver when the
             residual norm is below this value.

pytests/test_solver.py

@@ -104,7 +104,7 @@ def test_cg(par):
 
 @pytest.mark.parametrize("par", [(par1), (par2)])
 def test_cg_damp(par):
-    """CG with damping solves the damped system (Op + damp*I) x = y (issue #406).
+    """CG with damping solves the damped system (Op + damp*I) x = y.
 
     Verify equivalence between internal damping, adding ``damp*I`` to the
     operator (with ``damp=0`` in the solver), and the dense solution; and that
@@ -116,16 +116,17 @@ def test_cg_damp(par):
     A = np.random.normal(0, 1, (n, n)) + par["imag"] * np.random.normal(0, 1, (n, n))
     A = np.conj(A).T @ A + np.eye(n)  # symmetric positive-definite
     Aop = MatrixMult(A, dtype=par["dtype"])
-    y = np.random.normal(0, 1, n) + par["imag"] * np.random.normal(0, 1, n)
+    x = np.ones(n) + par["imag"] * np.ones(n)
+    y = Aop * x
     damp = 0.8
 
     x_dense = np.linalg.solve(A + damp * np.eye(n), y)
     # adding damp*I to the operator and using damp=0 must match internal damping
     Aop_damped = MatrixMult(A + damp * np.eye(n), dtype=par["dtype"])
 
-    for prealloc in [False, True]:
-        x_int = cg(Aop, y, niter=4 * n, tol=1e-10, damp=damp, preallocate=prealloc)[0]
-        x_ext = cg(Aop_damped, y, niter=4 * n, tol=1e-10, preallocate=prealloc)[0]
+    for preallocate in [False, True]:
+        x_int = cg(Aop, y, niter=4 * n, tol=1e-10, damp=damp, preallocate=preallocate)[0]
+        x_ext = cg(Aop_damped, y, niter=4 * n, tol=1e-10, preallocate=preallocate)[0]
         assert_array_almost_equal(x_int, x_dense, decimal=5)
         assert_array_almost_equal(x_int, x_ext, decimal=5)