Skip to content

cagpjax.linalg.eigh

Hermitian eigenvalue decomposition.

Eigh

Bases: Algorithm

Eigh algorithm for eigenvalue decomposition.

Source code in src/cagpjax/linalg/eigh.py 
27
28
29
30
class Eigh(cola.linalg.Algorithm):
    """
    Eigh algorithm for eigenvalue decomposition.
    """

EighResult

Bases: NamedTuple

Result of Hermitian eigenvalue decomposition.

Attributes:

Name Type Description
eigenvalues Float[Array, N]

Eigenvalues of the operator.
eigenvectors LinearOperator

Eigenvectors of the operator.
Source code in src/cagpjax/linalg/eigh.py 
15
16
17
18
19
20
21
22
23
24
class EighResult(NamedTuple):
    """Result of Hermitian eigenvalue decomposition.

    Attributes:
        eigenvalues: Eigenvalues of the operator.
        eigenvectors: Eigenvectors of the operator.
    """

    eigenvalues: Float[Array, "N"]
    eigenvectors: LinearOperator

Lanczos

Bases: Algorithm

Lanczos algorithm for approximate partial eigenvalue decomposition.

Parameters:

Name Type Description Default
max_iters int | None

Maximum number of iterations (number of eigenvalues/vectors to compute). If None, all eigenvalues/eigenvectors are computed.

 None
v0 Float[Array, N] | None

Initial vector. If None, a random vector is generated using key.

 None
key PRNGKeyArray | None

Random key for generating a random initial vector if v0 is not provided.

 None
Source code in src/cagpjax/linalg/eigh.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
class Lanczos(cola.linalg.Algorithm):
    """Lanczos algorithm for approximate partial eigenvalue decomposition.

    Args:
        max_iters: Maximum number of iterations (number of eigenvalues/vectors to compute).
            If `None`, all eigenvalues/eigenvectors are computed.
        v0: Initial vector. If `None`, a random vector is generated using `key`.
        key: Random key for generating a random initial vector if `v0` is
            not provided.
    """

    max_iters: int | None
    v0: Float[Array, "N"] | None
    key: PRNGKeyArray | None

    def __init__(
        self,
        max_iters: int | None = None,
        /,
        *,
        v0: Float[Array, "N"] | None = None,
        key: PRNGKeyArray | None = None,
    ):
        self.max_iters = max_iters
        self.v0 = v0
        self.key = key

eigh(A, alg=Eigh(), grad_rtol=None)

Compute the Hermitian eigenvalue decomposition of a linear operator.

For some algorithms, the decomposition may be approximate or partial.

Parameters:

Name Type Description Default
A LinearOperator

Hermitian linear operator.

 required
alg Algorithm

Algorithm for eigenvalue decomposition.

 Eigh()
grad_rtol float | None

Specifies the cutoff for similar eigenvalues, used to improve gradient computation for (almost-)degenerate matrices. If not provided, the default is 0.0. If None or negative, all eigenvalues are treated as distinct.

 None

Returns:

Type Description
EighResult

A named tuple of (eigenvalues, eigenvectors) where eigenvectors is a (semi-)orthogonal LinearOperator.
Note

Degenerate matrices have repeated eigenvalues. The set of eigenvectors that correspond to the same eigenvalue is not unique but instead forms a subspace. grad_rtol only improves stability of gradient-computation if the function being differentiated depends only depends on these subspaces and not the specific eigenvectors themselves.

Source code in src/cagpjax/linalg/eigh.py 
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
def eigh(
    A: LinearOperator,
    alg: cola.linalg.Algorithm = Eigh(),
    grad_rtol: float | None = None,
) -> EighResult:
    """Compute the Hermitian eigenvalue decomposition of a linear operator.

    For some algorithms, the decomposition may be approximate or partial.

    Args:
        A: Hermitian linear operator.
        alg: Algorithm for eigenvalue decomposition.
        grad_rtol: Specifies the cutoff for similar eigenvalues, used to improve
            gradient computation for (almost-)degenerate matrices.
            If not provided, the default is 0.0.
            If None or negative, all eigenvalues are treated as distinct.

    Returns:
        A named tuple of `(eigenvalues, eigenvectors)` where `eigenvectors` is a
            (semi-)orthogonal `LinearOperator`.

    Note:
        Degenerate matrices have repeated eigenvalues.
        The set of eigenvectors that correspond to the same eigenvalue is not unique
        but instead forms a subspace.
        `grad_rtol` only improves stability of gradient-computation if the function
        being differentiated depends only depends on these subspaces and not the
        specific eigenvectors themselves.
    """
    if grad_rtol is None:
        grad_rtol = -1.0
    vals, vecs = _eigh(A, alg, grad_rtol)  # pyright: ignore[reportArgumentType]
    if vecs.shape[-1] == A.shape[-1]:
        vecs = cola.Unitary(vecs)
    else:
        vecs = cola.Stiefel(vecs)
    return EighResult(vals, vecs)