diff --git a/pylops/optimization/basic.py b/pylops/optimization/basic.py index 9291b5e7..c4c3b0c0 100644 --- a/pylops/optimization/basic.py +++ b/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, diff --git a/pylops/optimization/cls_basic.py b/pylops/optimization/cls_basic.py index 4f47565f..dfc096f4 100644 --- a/pylops/optimization/cls_basic.py +++ b/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) diff --git a/pytests/test_solver.py b/pytests/test_solver.py index 13a2ca0c..d0b67d68 100644 --- a/pytests/test_solver.py +++ b/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)] )