Gram-Schmidt orthogonalization using NumPy einsum

gram_schmidt.py

import numpy as np

def gram_schmidt(X):
    '''
    Perform Gram-Schmidt orthogonalization on the rows of X.
    Return a new array Y with the same shape as Y of which
    the rows form an orthonormal basis.
    
    The input vectors must be linearly independent.
    The number of input vectors (rows) must not exceed the
    dimensionality (number of columns).
    '''
    
    X = np.asarray(X, dtype=float)
    assert len(X.shape) == 2, 'expected matrix of row vectors'
    assert X.shape[0] <= X.shape[1], 'more vectors than dimensions'
    
    # the first basis vector is parallel to the first input vector
    Y = np.empty(shape=X.shape)
    Y[0, :] = X[0, :] / np.linalg.norm(X[0, :])
    
    # construct the remaining N-1 basis vectors
    for i in range(1, len(X)):
        # project the next input vector (X[i]) to every
        # basis vector so far, then sum all projections
        current = X[i, :]
        basis = Y[:i, :]
        proj = np.einsum('ij, j, ik -> k', basis, current, basis)
        
        # subtract the combination of projections to obtain
        # a vector that is orthogonal to everything so far
        b = current - proj
        
        # normalize it and set it as the next basis vector
        Y[i, :] = b / np.linalg.norm(b)
    
    return Y

# Y = gram_schmidt([[1, 2, 3], [-1, 0, 1], [2, -2, 1]])
Y = gram_schmidt([[1, 2, 3, 4], [-1, 0, 1, 0], [2, -2, 1, 1]])

Z = Y @ Y.T
I = np.eye(Y.shape[0])
assert np.allclose(Z, I), 'returned basis is not orthonormal'

print(Y)
print()
print(Z.round(6))