Whether you’re solving differential equations, training neural networks, inferring parameters in a statistical model or solving optimization problems, you’ll find linear algebra at the core of almost every routine. Still, mistakes in the application of linear algebra are all too common. Even though it’s a novice numerical linear algebra error, I’ve seen professors solve \(A \boldsymbol{x} = b\) as b = np.linalg.inv(A) @ x and call it a day.

Here’s a list of beginner linear algebra mistakes, followed by some tips for writing better and faster numerical code. Remember to start with a correct implementation before working on speed.

Don’t compare floats

Floats are an approximation of real numbers, and comparing them bit-by-bit is too stringent:

np.power(np.sqrt(13), 2) == 13.0  # False
np.isclose(np.power(np.sqrt(13), 2), 13.0)  # True

A = np.random.randn(10, 10)
Q, R = np.linalg.qr(A)
np.all(Q @ R == A)  # False
np.allclose(Q @ R, A)  # True

Use numpy.isclose and numpy.allclose in general, and consider numpy.testing.assert_allclose for testing.

Don’t multiply diagonal matrices

The multiplication \(A \operatorname{diag}(\boldsymbol{v})\) is just a column-wise scaling of \(A\) by \(\boldsymbol{v}\). Similarly, \(\operatorname{diag}(\boldsymbol{v}) A\) is a row-wise scaling. Implement the scaling directly instead of multiplying with the diagonal matrix, since using a diagonal matrix mostly does unnecessary work by multiplying the entries of \(A\) with zeros.

import numpy as np
import scipy as sp

A = np.random.randn(1000, 500)
v = np.random.randn(500)

A @ np.diag(v)  # 6.3 ms ± 1.01 ms per loop
A * v  # 369 µs ± 35.9 µs per loop

assert np.allclose(A @ np.diag(v), A * v)

Don’t add diagonal matrices

Don’t explicitly form \(I\) when computing \(A + I\). If \(A\) has size \(n \times n\), then you’re adding \(n\) ones and \(n^2 - n\) unnecessary zeroes. Instead use numpy.fill_diagonal and do it like this:

A = np.random.randn(1000, 1000)
I = np.eye(1000)

A + I  # 1.07 ms ± 68.2 µs per loop
np.fill_diagonal(A, A.diagonal() + 1)  # 8.66 μs ± 141 ns per loop

This trick can be used to speed up adding arbitrary diagonal matrixes too.

Don’t multiply matrices in arbitrary order

Matrix multiplication is associative, so the expressions \((AB)\boldsymbol{x}\) and \(A(B\boldsymbol{x})\) evaluate to the same value—but the computational cost is not the same. Your can either work out the optimal ordering by hand, or let numpy.linalg.multi_dot do the job for you.

A = np.random.randn(1000, 500)
B = np.random.randn(500, 2000)
x = np.random.randn(2000)

(A @ B) @ x  # 18.8 ms ± 1.11 ms per loop
A @ (B @ x)  # 378 µs ± 114 µs per loop
A @ B @ x  # 19 ms ± 766 µs per loop
np.linalg.multi_dot([A, B, x])  # 378 µs ± 110 µs per loop

Consider another example: computing \(A \boldsymbol{v} \boldsymbol{v}^T A\). Here numpy.linalg.multi_dot won’t be of much help, since \(\boldsymbol{v}\) and \(\boldsymbol{v}^T\) are the same 1 dimensional array in NumPy. Should we implement the computation as \(A (\boldsymbol{v} \boldsymbol{v}^T) A\) or \((A \boldsymbol{v}) (\boldsymbol{v}^T A)\)?

A = np.random.randn(3000, 3000)
v = np.random.randn(3000)

A @ np.outer(v, v) @ A  # 910 ms ± 61.6 ms per loop
np.outer(A @ v, A.T @ v)  # 41.5 ms ± 3.95 ms per loop

assert np.allclose(A @ np.outer(v, v) @ A, np.outer(A @ v, A.T @ v))

• In the first approach we perform one outer product and two matrix multiplications.
• In the second approach we perform one outer product and two matrix-vector multiplications.

If \(A \in \mathbb{R}^{m \times n}\) and \(\boldsymbol{v} \in \mathbb{R}^{n}\), then computing the matrix-vector multiplication \(A \boldsymbol{v}\) uses \(\mathcal{O}(mn)\) operations. If \(A \in \mathbb{R}^{m \times n}\) and \(B \in \mathbb{R}^{n \times k}\), then computing the matrix-matrix multiplication \(A B\) uses \(\mathcal{O}(mnk)\) operations. So a matrix-vector multiplication is \(\mathcal{O}(k)\) faster than a matrix-matrix multiplication.

Don’t compute a product if you only need the diagonal

Consider computing \(\operatorname{diag}(A B)\). Forming the full product \(A B\) only to extract the diagonal is wasteful. Instead, compute the diagonal of the product as (B.T * A).sum(axis=1) in NumPy:

A = np.random.randn(5000, 10000)
B = np.random.randn(10000, 5000)

np.diag(A @ B)  # 3.55 s ± 151 ms per loop
(B.T * A).sum(axis=1)  # 188 ms ± 5.44 ms per loop

assert np.allclose(np.diag(A @ B), (B.T * A).sum(axis=1))

Don’t form the inverse matrix

There are several reasons why you don’t invert that matrix: it’s slow and it’s not numerically stable. Instead of explicitly solving \(A\boldsymbol{x} = b\) by forming and multiplying by \(A^{-1}\), use a solver.

A = np.random.randn(2000, 2000)
x = np.random.randn(2000)
b = A @ x

np.linalg.inv(A) @ x  # 355 ms ± 52.5 ms per loop
np.linalg.solve(A, b)  # 86.7 ms ± 8.33 ms per loop

Now let’s look at a slightly more advanced computation, which combines the idea of not forming inverses with the idea of not computing products in arbitrary order. We want to compute the product \(A B^{-1} C\), where \(A \in \mathbb{R}^{m \times n}\), \(B \in \mathbb{R}^{n \times n}\) and \(C \in \mathbb{R}^{n \times k}\).

• If \(m > k\), we should compute \(A (B^{-1} C)\), which uses \(\mathcal{O}( kn(m+n) )\) floating point operations.
• If \(k > m\), we should compute \((A B^{-1}) C\), which uses \(\mathcal{O}( mn(n +k) )\) floating point operations.

Below we solve the case when \(k \gg m\) without forming the inverse.

m, n, k = 10, 1000, 100_000

A = np.random.randn(m, n)
B = np.random.randn(n, n)
C = np.random.randn(n, k)

(A @ np.linalg.inv(B)) @ C  # 233 ms ± 56.4 ms per loop
A @ (np.linalg.inv(B) @ C)  # 2.32 s ± 145 ms per loop

np.linalg.solve(B.T, A.T).T @ C  # 157 ms ± 7.29 ms per loop
A @ np.linalg.solve(B, C)  # 2.71 s ± 174 ms per loop

Don’t use memory you don’t need

If you can, use inplace operations to avoid filling up the memory with unnecessary arrays:

v = np.random.randn(100)
v -= 1  # Mutate v without creating any new arrays
v = v - 1  # Create a new array (v - 1) in memory, assign it to v

Inplace operations are also typically a bit faster. Win-win.

Don’t use pure Python loops

Avoid pure Python loops in numerical code. If the obvious approach is to loop, first try to vectorize the computation. If that fails, look into Numba or Cython.

Here’s an example. The code below computes the empirical cross-covariance between two data sets. Each variable is a column, each row is a set of observations.

def covariance_loop(A, B):
    A = A - np.mean(A, axis=0)
    B = B - np.mean(B, axis=0)
    cov = np.zeros((A.shape[1], B.shape[1]))

    for j in range(A.shape[0]):
        cov += np.outer(A[j, :], B[j, :])

    return cov / (A.shape[0] - 1.0)


A = np.random.randn(1000, 1000)
B = np.random.randn(1000, 2000)

covariance_loop(A, B)  # 3.62 s ± 241 ms per loop

The code above computes an outer product plus a matrix-matrix addition. This is done for each row.

Although the asymptotic complexity is the same, the code below is much faster because it eliminates the Python loop and expresses the computation as a single matrix-matrix product instead.

def covariance(A, B):
    A = A - np.mean(A, axis=0)
    B = B - np.mean(B, axis=0)

    return A.T @ B / (A.shape[0] - 1.0)


covariance(A, B)  # 48.3 ms ± 8.75 ms per loop

Do use the Woodbury identity

Learn to use the Woodbury matrix identity and related formulas. The Woodbury equation is:

If \(P\) and \(R\) are positive definite, then the following holds:

A typical use case is when we have inverted \(A\), and want to compute the inverse of a rank 1 update of \(A\). In other words we wish to compute \(\left(A + \boldsymbol{v}\boldsymbol{v}^T \right)^{-1}\), and we already have \(A^{-1}\). This is a special case of the Woodbury matrix identity, called the Sherman–Morrison formula:

A = np.random.randn(1000, 1000)
v = np.random.randn(1000)
A_inv = np.linalg.inv(A)

np.linalg.inv(A + np.outer(v, v))  # 50.2 ms ± 6.32 ms per loop


def rank1update(A_inv, v):
    return A_inv - np.outer(A_inv @ v, v.T @ A_inv) / (1 + v.T @ A_inv @ v)

rank1update(A_inv, v)  # 8.61 ms ± 2.9 ms per loop

Do think about what the meaning is

Just like \(A \operatorname{diag}(\boldsymbol{v})\) simply scales the columns of \(A\) with the entries of \(\boldsymbol{v}\), more tangled expressions occasionally have simple interpretations. Consider the expression \(A(I - \frac{1}{N}\boldsymbol{1}\boldsymbol{1}^T)\). Looks fancy, right? It simply subtracts the mean from each row. Instead of forming \(I\), then forming \(\boldsymbol{1}\boldsymbol{1}^T\), performing the subtraction, multiplying, etc—simply implement the row scaling directly.

If you’re working in a thousand dimensions, its easy to get lost. Work out simple examples in three or four dimensions before generalizing.

A = np.random.randn(1000, 1000)
factor = np.eye(1000) - np.outer(np.ones(1000), np.ones(1000)) / 1000

A @ factor  # 25 ms ± 9.58 ms per loop
A - np.mean(A, axis=1, keepdims=True)  # 1.61 ms ± 60.4 µs per loop

assert np.allclose(A @ factor, A - np.mean(A, axis=1, keepdims=True))

Here’s another example from the wild. I once saw someone implement

v = np.exp(np.random.randn(5_000))
sp.linalg.cholesky(np.diag(v**2))  # 724 ms ± 99.9 ms per loop

because that’s what a paper said. But the Cholesky decomposition of a diagonal matrix is just the square of the diagonal, so the entire factorization is a waste of time—it does nothing at all.

Do time everything

We’ve been using the ubiquitous %timeit Ipython magic command in almost every example so far. If you’re in doubt about whether one approach or the other is better, then implement both and time them.

Consider solving the Ridge regression optimization problem

where \(\lVert \boldsymbol{x} \rVert^2_W = \boldsymbol{x}^T W \boldsymbol{x}\) denotes the weighted squared norm. We assume that \(W\) is a diagonal matrix with entries given by a vector \(\boldsymbol{w}\).

We will examine three approaches to solving this problem.

The first approach is to differentiate the objective function and set it equal to zero. We obtain the normal equations \((X^T W X + I) \boldsymbol{\beta} = X^T W \boldsymbol{y}\). Forming the normal equations and solving could look like this:

X = np.random.randn(10_000, 1000)
y = np.random.randn(10_000)
w = np.exp(np.random.randn(10_000))

def solve_normal_eqns(X, w, y):
    """Solve |X @ beta - y|_w^2 + |beta|^2 for beta."""
    XT_W = X.T * w  # Form X.T @ diag(w)
    lhs = XT_W @ X
    lhs.flat[:: lhs.shape[0] + 1] += 1  # Add identity
    rhs = XT_W @ y
    return sp.linalg.solve(
        lhs,
        rhs,
        assume_a="pos",
        overwrite_a=True,
        overwrite_b=True,
    )

solve_normal_eqns(X, w, y)  # 321 ms ± 97.2 ms per loop

The second approach is to realize that

and make a call to a least squares solver with these matrices directly.

def solve_lstsq(X, w, y):
    """Solve |X @ beta - y|_w^2 + |beta|^2 for beta."""
    squared_w = w**0.5
    lhs = np.vstack(((X.T * squared_w).T, np.eye(X.shape[1])))
    rhs = np.hstack((squared_w * y, np.zeros(X.shape[1])))
    beta, *_ = sp.linalg.lstsq(
        lhs,
        rhs,
        overwrite_a=True,
        overwrite_b=True,
    )
    return beta

solve_lstlsq(X, w, y)  # 1.34 s ± 429 ms per loop

The third approach uses the singular value decomposition (SVD). Let \(W^{1/2}X = U D V^T\), then the factor \((X^T W X + I)\) in the normal equations may be written as

and \((X^T W X + I)^{-1} = V^T (D^2 + I)^{-1} V\). Plugging this back into the normal equations, we obtain the solution as

In the implementation, we exploit the fact that \((D^2 + I)^{-1}\), \(D\) and \(W^{1/2}\) are diagonal matrices:

def solve_svd(X, w, y):
    """Solve |X @ beta - y|_w^2 + |beta|^2 for beta."""
    w_sq = w**0.5
    W_sq_X = (X.T * w_sq).T  # Form W^1/2 @ X
    U, d, VT = sp.linalg.svd(W_sq_X, full_matrices=False, overwrite_a=True)
    inv_diag = d / (1 + d**2)  # Form (D^2 + I)^-1 @ D
    return np.linalg.multi_dot([VT.T * inv_diag, U.T, (w_sq * y)])

solve_svd(X, w, y)  # 2.47 s ± 329 ms per loop

The approaches return the same answer:

assert np.allclose(solve_normal_eqns(X, w, y), solve_lstsq(X, w, y))
assert np.allclose(solve_normal_eqns(X, w, y), solve_svd(X, w, y))

The third approach using the SVD is what you might find in a paper. However, sklearn implements approach one, which is significantly faster. Under the hood, scipy.linalg.solve calls the LAPACK routine sposv, which uses a Cholesky decomposition. It’s typically instructive to implement a few different approaches, then test them and time them. We have not discussed the numerical stability of the least-squares problem here, but have a look at the references at the end of this article for more information about that issue.

Do read the documentation

Consider again the function solve_normal_eqns(X, w, y) that we used to solve the normal equations \((X^T W X + I) \boldsymbol{\beta} = X^T W \boldsymbol{y}\) with. It turns out that scipy.linalg.solve(a, b) has a parameter lower=False, and the documentation reads:

If False (default), the calculation uses only the data in the upper triangle of a; entries below the diagonal are ignored.

It also turns out that there’s a BLAS routine with the cryptic name dsyrk that only computes the upper-triangular part of \(\alpha A A^T\).

We can improve the speed of our solver by only computing the upper-triangular part of \(X^T W X + I\) as follows:

def solve_normal_eqns2(X, w, y):
    """Solve |X @ beta - y|_w^2 + |beta|^2 for beta."""

    # Compute upper triangular entries in X.T @ diag(w) @ X
    XT_W_X = sp.linalg.blas.dsyrk(alpha=1, a=X.T * w**0.5)

    # Add identity matrix
    XT_W_X.flat[:: XT_W_X.shape[0] + 1] += 1

    return sp.linalg.solve(
        XT_W_X,
        X.T @ (w*y),
        assume_a="pos",
        lower=False, # Solver only uses upper triangular elements
        overwrite_a=True,
        overwrite_b=True,
    )

X = np.random.randn(10_000, 5000)
y = np.random.randn(10_000)
w = np.exp(np.random.randn(10_000))

solve_normal_eqns(X, w, y)  # 4.78 s ± 393 ms per loop
solve_normal_eqns2(X, w, y)  # 3.33 s ± 436 ms per loop

If you’re working with linear algebra, familiarize yourself with all the functions in scipy.linalg. When using a function, read carefully through the arguments. If speed and accuracy truly matters, look into how functions are implemented and research the specific problem you’re working on carefully.

Do consider approximations

Sometimes we don’t need exact answers. For instance, we can approximately solve the Ridge regression problem by running five Conjugate Gradient iterations on the normal equations.

X = np.random.randn(10_000, 1000)
y = np.random.randn(10_000)
w = np.exp(np.random.randn(10_000))
XT_W = X.T * w


def matvec(v):
    # Implements (X.T @ diag(v) @ X + I) @ v
    return np.linalg.multi_dot([XT_W, X, v]) + v


A = sp.sparse.linalg.LinearOperator((X.shape[1], X.shape[1]), matvec=matvec)

# 56.9 ms ± 13 ms per loop
approx, _ = sp.sparse.linalg.cg(A, XT_W @ y, tol=0.0, maxiter=5)

# 295 ms ± 55.1 ms per loop
exact = solve_normal_eqns(X, w, y)

The Pearson correlation between the exact and approximate answer is approximately \(0.998\). What we gain in speed might make up for the loss of accuracy, but this is highly problem dependent. An approximate singular value decomposition can be obtained using sklearn.utils.extmath.randomized_svd.

Resources and further reading

Three excellent books:

• Numerical Linear Algebra by Trefethen et al (1997).
• Matrix Computations by Golub et al (2013).
• Accuracy and Stability of Numerical Algorithms by Nicholas J. Higham (2002).

Two reference manuals:

• The Matrix Reference Manual
• The Matrix Cookbook

An open-source book on vectorization:

• From Python to Numpy