Source code for femorph_solver.solvers.linear._cg

"""Conjugate Gradient iterative solver (SPD).

Uses :func:`scipy.sparse.linalg.cg` with a diagonal (Jacobi)
preconditioner.  Always available; best for very large sparse SPD
systems where factorization memory would be prohibitive.

References
----------
- Hestenes, M. R. & Stiefel, E.  "Methods of Conjugate Gradients
  for Solving Linear Systems."  J. Res. Natl. Bur. Stand. 49 (6),
  409-436 (1952).  [the original CG algorithm]
- Saad, Y.  *Iterative Methods for Sparse Linear Systems*, 2nd ed.
  SIAM (2003), §6.  [modern treatment + convergence theory]
"""

from __future__ import annotations

from typing import ClassVar

import numpy as np
import scipy.sparse.linalg as spla

from ._base import LinearSolverBase




class CGSolver(LinearSolverBase):
    """Preconditioned Conjugate Gradient for SPD ``A``."""

    name: ClassVar[str] = "cg"
    kind: ClassVar[str] = "iterative"
    spd_only: ClassVar[bool] = True
    install_hint: ClassVar[str] = "bundled with SciPy — no install needed"

    def __init__(
        self,
        A,
        *,
        assume_spd: bool = True,
        rtol: float = 1e-10,
        atol: float = 0.0,
        maxiter: int | None = None,
    ) -> None:
        self._rtol = rtol
        self._atol = atol
        self._maxiter = maxiter
        super().__init__(A, assume_spd=assume_spd)

    def _factor(self, A):
        self._A = A.tocsr()
        d = self._A.diagonal()
        inv = np.where(np.abs(d) > 0, 1.0 / d, 1.0)
        self._M = spla.LinearOperator(self._A.shape, matvec=lambda x: inv * x)



    def solve(self, b: np.ndarray) -> np.ndarray:
        x, info = spla.cg(
            self._A, b, M=self._M, rtol=self._rtol, atol=self._atol, maxiter=self._maxiter
        )
        if info != 0:
            raise RuntimeError(f"CG failed to converge (info={info})")
        return x