Source code for femorph_solver.solvers.linear._gmres

"""GMRES iterative solver (general).

Nonsymmetric Krylov solver — use when ``A`` is not SPD.  Jacobi
preconditioner by default.

References
----------
- Saad, Y. & Schultz, M. H.  "GMRES: A Generalized Minimal Residual
  Algorithm for Solving Nonsymmetric Linear Systems."  SIAM J.
  Sci. Stat. Comput. 7 (3), 856-869 (1986).  [the original GMRES
  algorithm implemented by ``scipy.sparse.linalg.gmres``]
- Saad, Y.  *Iterative Methods for Sparse Linear Systems*, 2nd ed.
  SIAM (2003), §6.5.  [modern treatment including restart + the
  restart-length trade-off scipy exposes via its ``restart`` kwarg]
"""

from __future__ import annotations

from typing import ClassVar

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

from ._base import LinearSolverBase




class GMRESSolver(LinearSolverBase):
    """Preconditioned GMRES (general sparse)."""

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

    def __init__(
        self,
        A,
        *,
        assume_spd: bool = False,
        rtol: float = 1e-10,
        atol: float = 0.0,
        restart: int = 50,
        maxiter: int | None = None,
    ) -> None:
        self._rtol = rtol
        self._atol = atol
        self._restart = restart
        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.gmres(
            self._A,
            b,
            M=self._M,
            rtol=self._rtol,
            atol=self._atol,
            restart=self._restart,
            maxiter=self._maxiter,
        )
        if info != 0:
            raise RuntimeError(f"GMRES failed to converge (info={info})")
        return x