Skip to content

cagpjax.solvers.cholesky

Linear solvers based on Cholesky decomposition.

Cholesky

Bases: AbstractLinearSolverMethod

Solve a linear system using the Cholesky decomposition.

Due to numerical imprecision, Cholesky factorization may fail even for positive-definite \(A\). Optionally, a small amount of jitter (\(\epsilon\)) can be added to \(A\) to ensure positive-definiteness. Note that the resulting system solved is slightly different from the original system.

Attributes:

Name Type Description
jitter ScalarFloat | None

Small amount of jitter to add to \(A\) to ensure positive-definiteness.
Source code in src/cagpjax/solvers/cholesky.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
class Cholesky(AbstractLinearSolverMethod):
    """
    Solve a linear system using the Cholesky decomposition.

    Due to numerical imprecision, Cholesky factorization may fail even for
    positive-definite $A$. Optionally, a small amount of `jitter` ($\\epsilon$) can
    be added to $A$ to ensure positive-definiteness. Note that the resulting system
    solved is slightly different from the original system.

    Attributes:
        jitter: Small amount of jitter to add to $A$ to ensure positive-definiteness.
    """

    jitter: ScalarFloat | None

    def __init__(self, jitter: ScalarFloat | None = None):
        self.jitter = jitter

    @override
    def __call__(self, A: LinearOperator) -> AbstractLinearSolver:
        return CholeskySolver(A, jitter=self.jitter)

CholeskySolver

Bases: AbstractLinearSolver

Solve a linear system by computing the Cholesky decomposition.

Source code in src/cagpjax/solvers/cholesky.py 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
class CholeskySolver(AbstractLinearSolver):
    """
    Solve a linear system by computing the Cholesky decomposition.
    """

    lchol: LinearOperator

    def __init__(self, A: LinearOperator, jitter: ScalarFloat | None = None):
        self.lchol = lower_cholesky(A, jitter)

    @override
    def solve(self, b: Float[Array, "N #K"]) -> Float[Array, "N #K"]:
        Linv = cola.linalg.inv(self.lchol)
        return Linv.T @ (Linv @ b)

    @override
    def logdet(self) -> ScalarFloat:
        return 2 * jnp.sum(jnp.log(cola.linalg.diag(self.lchol)))

    @override
    def inv_congruence_transform(
        self, B: LinearOperator | Float[Array, "K N"]
    ) -> LinearOperator | Float[Array, "K K"]:
        Linv = cola.linalg.inv(self.lchol)
        right_term = Linv @ B
        return right_term.T @ right_term

    @override
    def trace_solve(self, B: Self) -> ScalarFloat:
        L = cola.linalg.inv(self.lchol) @ B.lchol.to_dense()
        return jnp.sum(jnp.square(L))