Source code for femorph_solver.solvers.eigen._primme

"""PRIMME via ``primme.eigsh`` — a state-of-the-art block-Davidson solver.

Often beats ARPACK on large SPD generalized eigenproblems because of its
restart + soft-locking scheme.  Pip-installable and self-contained (no
extra system libs).

References
----------
- Stathopoulos, A. & McCombs, J. R.  "PRIMME: PReconditioned
  Iterative MultiMethod Eigensolver — Methods and software
  description."  ACM Trans. Math. Softw. 37 (2), 21:1-21:30 (2010).
  [the shipped PRIMME library: block-JD + GD+k restart + the
  soft-locking convergence machinery]
- Stathopoulos, A.  "Nearly Optimal Preconditioned Methods for
  Hermitian Eigenproblems Under Limited Memory.  Part I: Seeking
  One Eigenvalue."  SIAM J. Sci. Comput. 29 (2), 481-514 (2007).
  [the GD+k restart scheme inside PRIMME]
- Davidson, E. R.  "The Iterative Calculation of a Few of the
  Lowest Eigenvalues and Corresponding Eigenvectors of Large
  Real-Symmetric Matrices."  J. Comp. Phys. 17 (1), 87-94 (1975).
  [the original Davidson method PRIMME descends from]
"""

from __future__ import annotations

from typing import ClassVar

import numpy as np
import scipy.sparse as sp

from ._base import EigenSolverBase

try:
    import primme

    _HAVE = True
except ImportError:
    _HAVE = False




class PrimmeSolver(EigenSolverBase):
    """primme.eigsh for generalized SPD eigenproblems."""

    name: ClassVar[str] = "primme"
    kind: ClassVar[str] = "sparse-davidson"
    spd_only: ClassVar[bool] = True
    install_hint: ClassVar[str] = "`pip install primme`"

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



    @staticmethod
    def solve(
        K: sp.spmatrix | sp.sparray,
        M: sp.spmatrix | sp.sparray,
        n_modes: int,
        *,
        sigma: float = 0.0,
        linear_solver: str = "auto",
        tol: float = 1e-10,
    ) -> tuple[np.ndarray, np.ndarray]:
        # Map ``sigma=0`` (the modal-default) to PRIMME's pure
        # Davidson path (``sigma=None``, ``which="SA"``).  Passing
        # ``sigma=0.0`` triggers PRIMME's internal shift-invert: it
        # then needs to apply ``(K - 0·M)^{-1} = K^{-1}`` on every
        # restart and, with no ``OPinv`` callback supplied, falls
        # back to its own iterative inner solver (GMRES/QMR).  On a
        # 500k-DOF SOLID186 rotor with K nnz ≈ 75 M and the bore-
        # clamp making K ill-conditioned (rcond ≈ 3e-18), that inner
        # solver returns NaN and PRIMME aborts with error -41
        # ("matvec failed").  Pure Davidson on (K, M) avoids the
        # inner solve entirely — it converges via M-orthogonal
        # projection alone, which is cheap, robust, and the right
        # algorithm for the SPD generalized eigenproblem we have.
        Kcsr = K.tocsr()
        Mcsr = M.tocsr()
        if sigma == 0.0 or sigma is None:
            w, v = primme.eigsh(
                Kcsr,
                M=Mcsr,
                k=n_modes,
                which="SA",
                tol=tol,
            )
        else:
            # Non-zero shift — caller wants modes near σ.  PRIMME's
            # built-in shift-invert needs ``OPinv`` to be useful;
            # without it we re-route through ARPACK (the autotune
            # already does this for σ≠0, so reaching this branch
            # is a bypass-the-autotune call).  Pass through and let
            # PRIMME try its iterative inner solve — caller chose
            # this path explicitly.
            w, v = primme.eigsh(
                Kcsr,
                M=Mcsr,
                k=n_modes,
                which="SM",
                sigma=sigma,
                tol=tol,
            )
        order = np.argsort(w)
        return w[order], v[:, order]