Skip to content

qttools.wave_function_solver.superlu#

[docs] module qttools.wave_function_solver.superlu 

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
# Copyright (c) 2024-2026 ETH Zurich and the authors of the qttools package.

from qttools import NDArray, sparse
from qttools.wave_function_solver.solver import WFSolver

if "cupy" in sparse.__name__:
    from cupyx.scipy.sparse import linalg
else:
    from scipy.sparse import linalg


class SuperLU(WFSolver):
    """Wave function solver using LU decomposition for solving.

    This solver uses the SuperLU on the CPU for facorization. Depending
    on the chosen array module, the solution phase is computed on the
    CPU or GPU.

    """

    def solve(self, a: sparse.spmatrix, b: NDArray) -> NDArray:
        """Solves the sparse system a @ x = b using LU decomposition.

        Parameters
        ----------
        a : sparse.spmatrix
            The sparse system matrix.
        b : NDArray
            The right-hand side vector.

        Returns
        -------
        x : NDArray
            The solution vector.

        """
        lu = linalg.splu(a)
        return lu.solve(b)