Merge pull request #752 from gaoflow/feature/cg-damp

Add damp support to CG solver (closes #406)

Author: mrava87

pylops/optimization/basic.py

@@ -22,6 +22,7 @@ def cg(
     y: NDArray,
     x0: NDArray | None = None,
     niter: int = 10,
+    damp: float = 0.0,
     tol: float = 1e-4,
     rtol: float = 0.0,
     rtol1: float = 0.0,
@@ -45,6 +46,8 @@ def cg(
         Initial guess
     niter : :obj:`int`, optional
         Number of iterations
+    damp : :obj:`float`, optional
+        Damping coefficient
     tol : :obj:`float`, optional
         Absolute tolerance on residual norm. Stops the solver when the
         residual norm is below this value.
@@ -101,6 +104,7 @@ def cg(
         x0=x0,
         tol=tol,
         niter=niter,
+        damp=damp,
         show=show,
         itershow=itershow,
         preallocate=preallocate,

pylops/optimization/cls_basic.py

@@ -65,8 +65,9 @@ class CG(Solver):
 
     Notes
     -----
-    Solve the :math:`\mathbf{y} = \mathbf{Op}\,\mathbf{x}` problem using conjugate gradient
-    iterations [1]_.
+    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.
@@ -134,6 +135,7 @@ def setup(
         y: NDArray,
         x0: NDArray | None = None,
         niter: int | None = None,
+        damp: float = 0.0,
         tol: float = 1e-4,
         preallocate: bool = False,
         show: bool = False,
@@ -150,6 +152,8 @@ def setup(
         niter : :obj:`int`, optional
             Number of iterations (default to ``None`` in case a user wants to
             manually step over the solver)
+        damp : :obj:`float`, optional
+            Damping coefficient
         tol : :obj:`float`, optional
             Absolute tolerance on residual norm. Stops the solver when the
             residual norm is below this value.
@@ -170,6 +174,7 @@ def setup(
         """
         self.y = y
         self.niter = niter
+        self.damp = damp
         self.tol = tol
 
         self.ncp = get_array_module(y)
@@ -187,6 +192,12 @@ def setup(
             else:
                 self.r = self.ncp.empty_like(self.y)
                 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:
+                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()))
 
@@ -221,6 +232,9 @@ def step(self, x: NDArray, show: bool = False) -> NDArray:
 
         """
         Opc = self.Op.matvec(self.c)
+        # add damping contribution
+        if self.damp != 0.0:
+            Opc = Opc + self.damp * self.c
         cOpc = self.ncp.abs(self.c.dot(Opc.conj()))
         a = self.kold / cOpc
         if not self.preallocate:
@@ -317,6 +331,7 @@ def solve(
         y: NDArray,
         x0: NDArray | None = None,
         niter: int = 10,
+        damp: float = 0.0,
         tol: float = 1e-4,
         preallocate: bool = False,
         show: bool = False,
@@ -333,6 +348,8 @@ def solve(
             internally as zero vector
         niter : :obj:`int`, optional
             Number of iterations
+        damp : :obj:`float`, optional
+            Damping coefficient
         tol : :obj:`float`, optional
             Absolute tolerance on residual norm. Stops the solver when the
             residual norm is below this value.
@@ -360,7 +377,13 @@ def solve(
 
         """
         x = self.setup(
-            y=y, x0=x0, niter=niter, tol=tol, preallocate=preallocate, show=show
+            y=y,
+            x0=x0,
+            niter=niter,
+            damp=damp,
+            tol=tol,
+            preallocate=preallocate,
+            show=show,
         )
         x = self.run(x, niter, show=show, itershow=itershow)
         self.finalize(show)

pytests/test_solver.py

@@ -102,6 +102,39 @@ def test_cg(par):
         assert_array_almost_equal(x, xinv, decimal=4)
 
 
+@pytest.mark.parametrize("par", [(par1), (par2)])
+def test_cg_damp(par):
+    """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
+    ``damp=0`` reproduces the undamped result.
+    """
+    np.random.seed(10)
+
+    n = par["nx"]
+    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"])
+    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 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)
+
+    # damp=0 reproduces the undamped solution
+    x_undamped = cg(Aop, y, niter=4 * n, tol=1e-10, damp=0.0)[0]
+    assert_array_almost_equal(x_undamped, np.linalg.solve(A, y), decimal=5)
+
+
 @pytest.mark.parametrize(
     "par", [(par1), (par2), (par3), (par4), (par1j), (par2j), (par3j), (par3j)]
 )