Gram-Schmidt row-wise vs column-wise

gram-schmidt.py

import numpy as np
import perfplot

rng = np.random.default_rng(1)


def setup(n):
    U = rng.random((n, 10))
    return U, np.ascontiguousarray(U.T)


def gram_schmidt_cols(data):
    V, _ = data
    U = np.copy(V)
    n = U.shape[1]
    for i in range(n):
        # for j in range(i):
        #     U[:, i] -= (U[:, i] @ U[:, j]) * U[:, j]
        U[:, i] -= U[:, :i] @ (U[:, i].T @ U[:, :i])
        U[:, i] /= np.sqrt((U[:, i] @ U[:, i]))
    return U


def gram_schmidt_rows(data):
    _, V = data
    U = np.copy(V)
    n = U.shape[0]
    for i in range(n):
        # for j in range(i):
        #     U[i] -= (U[i] @ U[j]) * U[j]
        U[i] -= (U[:i] @ U[i]) @ U[:i]
        U[i] /= np.sqrt(U[i] @ U[i])
    return U.T


def modified_gram_schmidt_cols(data):
    V, _ = data
    U = np.copy(V)
    n = U.shape[1]
    for i in range(n):
        U[:, i] /= np.sqrt(U[:, i] @ U[:, i])
        for j in range(i + 1, n):
            U[:, j] -= U[:, i] * (U[:, i] @ U[:, j])
    return U


def modified_gram_schmidt_cols2(data):
    V, _ = data
    U = np.copy(V)
    n = U.shape[1]
    for i in range(n):
        U[:, i] /= np.sqrt(U[:, i] @ U[:, i])
        U[:, i + 1 :] -= np.outer(U[:, i], U[:, i].T @ U[:, i + 1 :])
    return U


def modified_gram_schmidt_rows(data):
    _, V = data
    U = np.copy(V)
    n = U.shape[0]
    for i in range(n):
        U[i] /= np.sqrt(U[i] @ U[i])
        for j in range(i + 1, n):
            U[j] -= U[i] * (U[i] @ U[j])
    return U.T


def modified_gram_schmidt_rows2(data):
    _, V = data
    U = np.copy(V)
    n = U.shape[0]
    for i in range(n):
        U[i] /= np.sqrt(U[i] @ U[i])
        U[i + 1 :] -= np.outer(U[i + 1 :] @ U[i], U[i])
    return U.T


b = perfplot.bench(
    setup=setup,
    kernels=[
        gram_schmidt_cols,
        gram_schmidt_rows,
        modified_gram_schmidt_cols,
        modified_gram_schmidt_cols2,
        modified_gram_schmidt_rows,
        modified_gram_schmidt_rows2,
    ],
    n_range=[2 ** k for k in range(4, 24)],
    title="10 vectors of length n",
    xlabel="n",
)
b.save("out0.png")
b.show()

Row-wise is faster. Reason: U[j] is a contiguous block in memory, U[:, j] isn't.