Skip to content

matfree.low_rank

matfree.low_rank

Low-rank approximations (like partial Cholesky decompositions) of matrices.

matfree.low_rank.cholesky_partial(mat_el: Callable, /, *, nrows: int, rank: int) -> Callable

Compute a partial Cholesky factorisation.

Source code in matfree/low_rank.py 
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
def cholesky_partial(mat_el: Callable, /, *, nrows: int, rank: int) -> Callable:
    """Compute a partial Cholesky factorisation."""

    def cholesky(*params):
        if rank > nrows:
            msg = f"Rank exceeds n: {rank} >= {nrows}."
            raise ValueError(msg)
        if rank < 1:
            msg = f"Rank must be positive, but {rank} < {1}."
            raise ValueError(msg)

        step = _cholesky_partial_body(mat_el, nrows, *params)
        chol = np.zeros((nrows, rank))
        return control_flow.fori_loop(0, rank, step, chol), {}

    return cholesky

matfree.low_rank.cholesky_partial_pivot(mat_el: Callable, /, *, nrows: int, rank: int) -> Callable

Compute a partial Cholesky factorisation with pivoting.

Source code in matfree/low_rank.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
def cholesky_partial_pivot(mat_el: Callable, /, *, nrows: int, rank: int) -> Callable:
    """Compute a partial Cholesky factorisation with pivoting."""

    def cholesky(*params):
        if rank > nrows:
            msg = f"Rank exceeds nrows: {rank} >= {nrows}."
            raise ValueError(msg)
        if rank < 1:
            msg = f"Rank must be positive, but {rank} < {1}."
            raise ValueError(msg)

        body = _cholesky_partial_pivot_body(mat_el, nrows, *params)

        L = np.zeros((nrows, rank))
        P = np.arange(nrows)

        init = (L, P, P, True)
        (L, P, _matrix, success) = control_flow.fori_loop(0, rank, body, init)
        return _pivot_invert(L, P), {"success": success}

    return cholesky

matfree.low_rank.preconditioner(cholesky: Callable) -> Callable

Turn a low-rank approximation into a preconditioner.

Parameters:

Name Type Description Default
cholesky Callable

(Partial) Cholesky decomposition. Usually, the result of either cholesky_partial or cholesky_partial_pivot.

 required

Returns:

Type Description
solve

A function that computes

\[ (v, s, *p) \mapsto (sI + L(*p) L(*p)^\top)^{-1} v, \]

where \(K = [k(i,j,*p)]_{ij} \approx L(*p) L(*p)^\top\) and \(L\) comes from the low-rank approximation.
Source code in matfree/low_rank.py 
 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
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
def preconditioner(cholesky: Callable, /) -> Callable:
    r"""Turn a low-rank approximation into a preconditioner.

    Parameters
    ----------
    cholesky
        (Partial) Cholesky decomposition.
        Usually, the result of either
        [cholesky_partial][matfree.low_rank.cholesky_partial]
        or
        [cholesky_partial_pivot][matfree.low_rank.cholesky_partial_pivot].


    Returns
    -------
    solve
        A function that computes

        $$
        (v, s, *p) \mapsto (sI + L(*p) L(*p)^\top)^{-1} v,
        $$

        where $K = [k(i,j,*p)]_{ij} \approx L(*p) L(*p)^\top$
        and $L$ comes from the low-rank approximation.
    """

    def solve(v: Array, s: float, *cholesky_params):
        chol, info = cholesky(*cholesky_params)

        # Assert that the low-rank matrix is tall,
        # not wide (every sign has a story...)
        N, n = np.shape(chol)
        assert n <= N, (N, n)

        # Scale
        U = chol / np.sqrt(s)
        V = chol.T / np.sqrt(s)
        v /= s

        # Cholesky decompose the capacitance matrix
        # and solve the system
        eye_n = np.eye(n)
        chol_cap = linalg.cho_factor(eye_n + V @ U)
        sol = linalg.cho_solve(chol_cap, V @ v)
        return v - U @ sol, info

    return solve