Source code for femorph_solver.solvers.linear._pyamg

"""Algebraic multigrid (AMG) preconditioned CG via ``pyamg``.

AMG is the standard industrial preconditioner for elliptic problems like
structural stiffness — often 10-100× fewer iterations than Jacobi-CG on
large unstructured SPD systems.

References
----------
- Ruge, J. W. & Stüben, K.  "Algebraic multigrid."  In *Multigrid
  Methods* (S. F. McCormick, ed.), SIAM Frontiers in Applied
  Mathematics, Chapter 4, pp. 73-130 (1987).  [the foundational
  classical-AMG algorithm]
- Vaněk, P., Mandel, J., & Brezina, M.  "Algebraic multigrid by
  smoothed aggregation for second and fourth order elliptic
  problems."  Computing 56 (3), 179-196 (1996).
  [smoothed-aggregation variant, pyamg's default ``smoothed_aggregation_solver``]
- Bell, N., Olson, L. N., & Schroder, J.  "PyAMG: Algebraic
  Multigrid Solvers in Python v5.0."  Journal of Open Source
  Software 8 (81), 5495 (2023).  [the shipped implementation]
"""

from __future__ import annotations

from typing import ClassVar

import numpy as np

from ._base import LinearSolverBase

try:
    import pyamg

    _HAVE = True
except ImportError:
    _HAVE = False




class PyAMGSolver(LinearSolverBase):
    """Smoothed-aggregation AMG preconditioned CG (SPD)."""

    name: ClassVar[str] = "pyamg"
    kind: ClassVar[str] = "iterative"
    spd_only: ClassVar[bool] = True
    install_hint: ClassVar[str] = "`pip install pyamg`"

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



    @staticmethod
    def available() -> bool:
        return _HAVE


    def _factor(self, A):
        # pyamg's C++ kernels are bound to int32 indices only; coerce
        # here so callers can pass the int64-indexed CSR/CSC that SciPy
        # produces on 64-bit systems.
        A_csr = A.tocsr().astype(np.float64)
        A_csr = A_csr.__class__(
            (A_csr.data, A_csr.indices.astype(np.int32), A_csr.indptr.astype(np.int32)),
            shape=A_csr.shape,
        )
        self._A = A_csr
        self._ml = pyamg.smoothed_aggregation_solver(self._A)



    def solve(self, b: np.ndarray) -> np.ndarray:
        return self._ml.solve(
            np.ascontiguousarray(b, dtype=np.float64),
            tol=self._rtol,
            maxiter=self._maxiter,
            accel="cg",
        )