Function _cholesky

dask/array/linalg.py:1321–1399

Private function to perform Cholesky decomposition, which returns both lower and upper triangulars.

(a)

def _cholesky(a):
    """
    Private function to perform Cholesky decomposition, which returns both
    lower and upper triangulars.
    """

    if a.ndim != 2:
        raise ValueError("Dimension must be 2 to perform cholesky decomposition")

    xdim, ydim = a.shape
    if xdim != ydim:
        raise ValueError(
            "Input must be a square matrix to perform cholesky decomposition"
        )
    if len(set(a.chunks[0] + a.chunks[1])) != 1:
        msg = (
            "All chunks must be a square matrix to perform cholesky decomposition. "
            "Use .rechunk method to change the size of chunks."
        )
        raise ValueError(msg)

    vdim = len(a.chunks[0])
    hdim = len(a.chunks[1])

    token = tokenize(a)
    name = f"cholesky-{token}"

    # (name_lt_dot, i, j, k, l) corresponds to l_ij.dot(l_kl.T)
    name_lt_dot = f"cholesky-lt-dot-{token}"
    # because transposed results are needed for calculation,
    # we can build graph for upper triangular simultaneously
    name_upper = f"cholesky-upper-{token}"

    # calculates lower triangulars because subscriptions get simpler
    dsk = {}
    for i in range(vdim):
        for j in range(hdim):
            if i < j:
                dsk[name, i, j] = (
                    partial(np.zeros_like, shape=(a.chunks[0][i], a.chunks[1][j])),
                    meta_from_array(a),
                )
                dsk[name_upper, j, i] = (name, i, j)
            elif i == j:
                target = (a.name, i, j)
                if i > 0:
                    prevs = []
                    for p in range(i):
                        prev = name_lt_dot, i, p, i, p
                        dsk[prev] = (np.dot, (name, i, p), (name_upper, p, i))
                        prevs.append(prev)
                    target = (operator.sub, target, (sum, prevs))
                dsk[name, i, i] = (_cholesky_lower, target)
                dsk[name_upper, i, i] = (_conj_transpose, (name, i, i))
            else:
                # solving x.dot(L11.T) = (A21 - L20.dot(L10.T)) is equal to
                # L11.dot(x.T) = A21.T - L10.dot(L20.T)
                # L11.dot(x.T) = A12 - L10.dot(L02)