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]:
        w, v = primme.eigsh(
            K.tocsr(),
            M=M.tocsr(),
            k=n_modes,
            which="SM",
            sigma=sigma,
            tol=tol,
        )
        order = np.argsort(w)
        return w[order], v[:, order]