From f39187c577c6569e605c326006c35dbe98af86a4 Mon Sep 17 00:00:00 2001 From: Vincent Gao Date: Fri, 29 May 2026 12:19:27 +0200 Subject: [PATCH 1/2] Add damp support to CG solver (closes #406) The CG solver only solved Op x = y. Re-introduce the damp coefficient (removed in v2) so it can solve the damped system (Op + damp*I) x = y, consistent with cgls/lsqr. Damping is applied by adding damp*c to every operator application Op*c in the iterations (and to the initial residual when x0 is given); damp=0 reproduces the previous behaviour exactly. Thread damp through cg(), CG.solve and CG.setup, and document it. Add test_cg_damp verifying that internal damping matches both adding damp*I to the operator (with damp=0) and the dense (Op + damp*I)^-1 y, and that damp=0 still recovers the undamped solution. --- pylops/optimization/basic.py | 7 +++++++ pylops/optimization/cls_basic.py | 33 +++++++++++++++++++++++++++++++- pytests/test_solver.py | 32 +++++++++++++++++++++++++++++++ 3 files changed, 71 insertions(+), 1 deletion(-) diff --git a/pylops/optimization/basic.py b/pylops/optimization/basic.py index 9291b5e79..6e064e858 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,11 @@ def cg( Initial guess 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. tol : :obj:`float`, optional Absolute tolerance on residual norm. Stops the solver when the residual norm is below this value. @@ -101,6 +107,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 4f47565ff..e102fef65 100644 --- a/pylops/optimization/cls_basic.py +++ b/pylops/optimization/cls_basic.py @@ -68,6 +68,13 @@ class CG(Solver): 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). + .. [1] Hestenes, M R., Stiefel, E., “Methods of Conjugate Gradients for Solving Linear Systems”, Journal of Research of the National Bureau of Standards. vol. 49. 1952. @@ -134,6 +141,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 +158,10 @@ 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. 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}`. tol : :obj:`float`, optional Absolute tolerance on residual norm. Stops the solver when the residual norm is below this value. @@ -170,6 +182,7 @@ def setup( """ self.y = y self.niter = niter + self.damp = damp self.tol = tol self.ncp = get_array_module(y) @@ -187,6 +200,9 @@ 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: + self.r -= self.damp * x self.c = self.r.copy() self.kold = self.ncp.abs(self.r.dot(self.r.conj())) @@ -221,6 +237,10 @@ 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 + 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 +337,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 +354,10 @@ def solve( internally as zero vector 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}`. tol : :obj:`float`, optional Absolute tolerance on residual norm. Stops the solver when the residual norm is below this value. @@ -360,7 +385,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 13a2ca0cd..1a3dde154 100644 --- a/pytests/test_solver.py +++ b/pytests/test_solver.py @@ -102,6 +102,38 @@ 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 (issue #406). + + 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"]) + y = np.random.normal(0, 1, n) + par["imag"] * np.random.normal(0, 1, n) + 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] + 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)] ) From 6828c57cf055040c6fccf01236732c2c0bb68a76 Mon Sep 17 00:00:00 2001 From: Vincent Gao Date: Fri, 29 May 2026 23:12:54 +0200 Subject: [PATCH 2/2] 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 --- pylops/optimization/basic.py | 5 +---- pylops/optimization/cls_basic.py | 28 ++++++++++------------------ pytests/test_solver.py | 11 ++++++----- 3 files changed, 17 insertions(+), 27 deletions(-) diff --git a/pylops/optimization/basic.py b/pylops/optimization/basic.py index 6e064e858..c4c3b0c09 100644 --- a/pylops/optimization/basic.py +++ b/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. diff --git a/pylops/optimization/cls_basic.py b/pylops/optimization/cls_basic.py index e102fef65..dfc096f4f 100644 --- a/pylops/optimization/cls_basic.py +++ b/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. diff --git a/pytests/test_solver.py b/pytests/test_solver.py index 1a3dde154..d0b67d68f 100644 --- a/pytests/test_solver.py +++ b/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)