Source code for femorph_solver.solvers.linear._cholmod

"""CHOLMOD sparse Cholesky via ``scikit-sparse``.

For SPD ``A`` (stiffness matrices after BCs, for example), CHOLMOD's
supernodal Cholesky is ~2-3× faster than SuperLU's LU and uses about
half the memory.

References
----------
- Chen, Y., Davis, T. A., Hager, W. W., & Rajamanickam, S.
  "Algorithm 887: CHOLMOD, Supernodal Sparse Cholesky Factorization
  and Update/Downdate."  ACM Trans. Math. Softw. 35 (3), 22:1-22:14
  (2008).  [the shipped supernodal algorithm + its update paths]
- Davis, T. A. & Hager, W. W.  "Dynamic supernodes in sparse
  Cholesky update/downdate and triangular solves."  ACM Trans.
  Math. Softw. 35 (4), 27:1-27:23 (2009).
  [the row-and-column modification path used when updating
  factorisations across parametric sweeps]
"""

from __future__ import annotations

from typing import ClassVar

import numpy as np
import scipy.sparse as sp

from ._base import LinearSolverBase

_HAVE = True
try:
    # scikit-sparse ≥0.5 renamed the factorization entry point to
    # ``cho_factor`` (returning a CholeskyFactor); older builds still
    # export ``cholesky`` (returning a Factor).  Prefer the new name
    # and fall back for 0.4.x.
    from sksparse.cholmod import cho_factor as _cho_factor
except ImportError:
    try:
        from sksparse.cholmod import cholesky as _cho_factor
    except ImportError:
        _HAVE = False
        _cho_factor = None  # type: ignore[assignment]

#: kwarg that picks the supernodal-vs-simplicial strategy — 0.5.x renamed
#: ``mode`` → ``supernodal_mode``.  Resolved once at import so the hot
#: factor path doesn't pay for exception-based dispatch per call.
_SUPER_KW: str | None = None
if _HAVE:
    import inspect as _inspect

    _params = _inspect.signature(_cho_factor).parameters
    if "supernodal_mode" in _params:
        _SUPER_KW = "supernodal_mode"
    elif "mode" in _params:
        _SUPER_KW = "mode"




class CholmodSolver(LinearSolverBase):
    """CHOLMOD supernodal sparse Cholesky (SPD only)."""

    name: ClassVar[str] = "cholmod"
    kind: ClassVar[str] = "direct"
    spd_only: ClassVar[bool] = True
    install_hint: ClassVar[str] = "install SuiteSparse system library + `pip install scikit-sparse`"



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


    def _factor(self, A):
        # scikit-sparse predates scipy's new sparse-array API; coerce to
        # the legacy csc_matrix type it insists on.  Skip the astype copy
        # when the data is already float64 — at ~60k DOFs that saves
        # ~80ms per call.  ``mode="auto"`` lets CHOLMOD pick between
        # simplicial and supernodal factorization; its default is plain
        # supernodal, which is ~2× slower than the simplicial path on
        # 3D hex meshes with many small supernodes.
        if A.dtype != np.float64:
            A = A.astype(np.float64)
        if isinstance(A, sp.csc_matrix):
            A_csc = A
        else:
            A_csc = sp.csc_matrix(A)
        # ``simplicial`` beats CHOLMOD's ``"auto"`` heuristic on FE
        # stiffness matrices coming out of 3D hex meshes: the factor is
        # sparse enough that the extra bookkeeping of the supernodal
        # path isn't recouped.  Measured ~1.7× faster on 60k-DOF hex
        # problems than the auto-selected supernodal path.  The kwarg
        # was renamed in 0.5.x (``supernodal_mode``) — ``_SUPER_KW`` is
        # resolved once at import.
        if _SUPER_KW is not None:
            self._factor_obj = _cho_factor(A_csc, **{_SUPER_KW: "simplicial"})
        else:
            self._factor_obj = _cho_factor(A_csc)



    def solve(self, b: np.ndarray) -> np.ndarray:
        b = np.ascontiguousarray(b, dtype=np.float64)
        # 0.5.x exposes ``solve(b, system="A")``; 0.4.x exposes
        # ``solve_A(b)``.  Prefer the newer one.
        if hasattr(self._factor_obj, "solve_A"):
            return self._factor_obj.solve_A(b)
        return self._factor_obj.solve(b, system="A")