linalg.py

import numpy as np

def plu_inplace(A):
    P = np.arange(len(A))
    for i in range(len(A)):
        j = i + np.argmax(np.abs(A[i:, i]))
        P[[i, j]], A[[i, j]] = P[[j, i]], A[[j, i]]
        A[i+1:, i] /= A[i, i]
        A[i+1:, i+1:] -= np.outer(A[i+1:, i], A[i, i+1:])
    return P, A, A

def l1_solve(L, y):
    x = y.copy()
    for i in range(len(x)):
        x[i] -= L[i, :i] @ x[:i]
    return x

def l_solve(L, y):
    x = y.copy()
    for i in range(len(x)):
        x[i] = (x[i] - L[i, :i] @ x[:i]) / L[i, i]
    return x

def u_solve(U, y):
    return l_solve(U[::-1, ::-1], y[::-1])[::-1]

def lu_solve(L, U, y):
    return u_solve(U, l1_solve(L, y))

def plu_solve(P, L, U, y):
    return lu_solve(L, U, y[P])

def cholesky_inplace(A):
    for i in range(len(A)):
        A[i, i] = np.sqrt(A[i, i])
        A[i+1:, i] /= A[i, i]
        A[i+1:, i+1:] -= np.outer(A[i+1:, i], A[i+1:, i])
    return A

def cholesky_solve(L, y):
    return u_solve(L.T, l_solve(L, y))

def qr_inplace(A):
    Q = []
    for i in range(len(A)):
        v = A[i:, i].copy()
        v[0] += np.sign(v[0]) * np.linalg.norm(v)
        v /= np.linalg.norm(v)
        Q.append(v)
        A[i:, i:] -= np.outer(v * 2, v @ A[i:, i:])
    return Q, A

def q_solve(Q, y):
    x = y.copy()
    for i, v in enumerate(Q):
        x[i:] -= v * 2 * (v @ x[i:])
    return x

def qr_solve(Q, R, y):
    return u_solve(R, q_solve(Q, y))

# Make a solver for A + u v^T given a solver for A, using the Sherman-Morrison formula.
def rank1_update(A_solve, u, v):
    w = A_solve(u)
    w /= 1 + np.dot(v, w)
    def solve(y):
        x = A_solve(y)
        return x - np.dot(v, x) * w
    return solve