Joe Despres

5 min read

Linear algebra from scratch in Rust

Four months, 58 commits, and a matrix inverse built from cofactors. The type system flagged the float sorting problem. It didn't catch a single shape error.

I wanted to see how much of the math in linear algebra a strict type system could check for me, and I wanted to know how the operations actually work. So when I started working through Mike X Cohen’s Practical Linear Algebra for Data Science in April 2023, I did the exercises in Rust, though the book’s are in Python, and wrote each operation as loops before calling the library version.

The repo has 58 commits across 28 days, from April 11 to August 15, 2023, most of them written in the evenings or before work.

Dates What got built
Apr 11–18 Vector operations, lengths, transpose, dot products, Pearson correlation and cosine similarity
Apr 30–May 10 A 1D convolution (“kernel method”), matrix helpers, matrix multiplication, a symmetry check
May 30–Jun 5 Frobenius norm, trace, random-walk matrices, covariance and correlation matrices
Jun 10–13 Got ndarray-linalg linking, submatrices, minors, the sign grid, a full matrix inverse
Jul 27–30 Bug fixes, a hand-written matrix multiply checked against the library, left and right inverses
Aug 15 The Hilbert matrix

Everything is ndarray

I didn’t write my own array type. Everything is ndarray::Array1<f64> or Array2<f64>, and I put my own methods on those types with extension traits. The trait bounds are the part I liked best, because they say which kinds of numbers an operation makes sense for:

pub trait MyVecMethods<T>
where
    T: ndarray::LinalgScalar,
{
    fn kernel_mul(&self, kernel: &Array1<T>) -> Array1<T>;
}

pub trait FloatMethods<T>
where
    T: ndarray::NdFloat,
{
    fn l2_norm(&self) -> T;
    fn unit_vec(&self) -> Array1<T>;
    fn set_magnitude(&self, l: T) -> Array1<T>;
}

A convolution only needs + and *, so LinalgScalar is enough and it works on integers. A vector length needs a square root, so l2_norm requires NdFloat. If you call .l2_norm() on an Array1<i32>, the compiler won’t let you. In R or NumPy you’d only find that out when a result came back looking wrong.

The other habit I picked up was returning a named struct when a function computes more than one thing:

pub struct OrthogonalDecomp {
    parallel_component: Array1<f64>,
    orthogonal_component: Array1<f64>,
}

A tuple (Array1<f64>, Array1<f64>) would let a caller swap the two components by accident without noticing. The struct doesn’t.

Stuck at sorting floats

I wanted max, min, and abs_max on a matrix. The obvious move is to flatten it and sort. In Rust, f64 doesn’t implement Ord, so .sort() won’t compile. That’s because of NaN: NaN < 1.0 and NaN > 1.0 are both false, so floats don’t have a total order.

The usual workaround is this:

vec.sort_by(|a, b| a.partial_cmp(b).unwrap());

It compiles, but it panics the first time it sees a NaN. The compiler is pointing at a real edge case, and .unwrap() only makes it go away. f64::total_cmp, already in stable Rust at the time, is the better answer. It gives a total order and puts NaN at one end. For abs_max I didn’t need a sort at all, since a fold is O(n) and doesn’t allocate:

a.iter().fold(0.0_f64, |m, x| m.max(x.abs()))

abs_max ended up being the function most of my tests depended on. Nearly every numeric check in the repo looks like assert!((a - b).abs_max() < TOL).

Inverting a matrix by hand

The June commits were about the inverse. Building it from cofactors takes five steps:

  1. Take the submatrices: remove row i and column j. I used iterators: zip each row with a keep/drop flag, filter, flatten, reshape.
  2. Build the minors matrix, the determinant of every submatrix. I used ndarray-linalg’s det() here, so it isn’t entirely from scratch.
  3. Build the sign grid, a checkerboard of +1 and −1. When n is even, filling row by row in alternating order gives stripes instead of a checkerboard, so I build an (n+1)×(n+1) grid and slice off the top-left n×n corner.
  4. Get the cofactors by multiplying the minors by the grid element-wise.
  5. Get the inverse by transposing the cofactors and dividing by the determinant.

The test takes the inverse of the inverse and checks that it matches the original matrix. It also multiplies a random 10×10 matrix by its inverse and checks the result against the identity.

In July I added left and right inverses for non-square matrices: (AᵀA)⁻¹Aᵀ for tall ones and Aᵀ(AAᵀ)⁻¹ for wide ones. The tests use a 7×3 and a 3×8 matrix and check that multiplying gives the 3×3 identity.

What the types didn’t catch

The type system flagged every problem with the kind of number. It caught nothing about shape, and in linear algebra most mistakes are shape mistakes.

Array2<f64> says the array has two dimensions. It doesn’t say how many rows or columns. So my matrix multiply begins with an assertion that runs at runtime:

pub fn mat_mul(a: Array2<f64>, b: Array2<f64>) -> Array2<f64> {
    assert_eq!(a.ncols(), b.nrows());
    // ...
}

That’s the same check R does when %*% stops with “non-conformable arguments.” It just gets there sooner.

Rust can check shapes at compile time with const generics:

struct Matrix<const R: usize, const C: usize>([[f64; C]; R]);

impl<const R: usize, const K: usize> Matrix<R, K> {
    fn mat_mul<const C: usize>(&self, other: &Matrix<K, C>) -> Matrix<R, C> {
        // same triple loop
    }
}

Multiplying a 2×3 by a 2×3 is then a compile error, and the output type is 2×2 without me writing it anywhere:

error[E0308]: mismatched types
   |     let _bad = a.mat_mul(&a);
   |                ------- ^^ expected `3`, found `2`

The catch is that the sizes have to be known when you compile. That’s fine for a 3×3 rotation and no help when the matrix is a data frame you read from a CSV. For that case, the approach I’d take now is to check once where the data comes in and then keep the result in the type. is_symmetric returns a bool, and the caller has to remember what it said. A Symmetric wrapper you can only build by passing the check keeps that information around:

struct Symmetric(Array2<f64>);

impl Symmetric {
    fn new(a: Array2<f64>, tol: f64) -> Option<Symmetric> { /* check once */ }
}

A covariance matrix is symmetric by definition, so cov_mat could return Symmetric, and anything downstream that relies on symmetry could take a Symmetric as its argument. It’s the idea I took away from the project.

Checking against a library

The test that did the most for me was the laziest one. test_mat_mul fills a 50×150 and a 150×50 matrix with random walks, multiplies them with my triple loop and with ndarray’s .dot(), and checks the difference. The math is complicated, but comparing to a library you trust is easy.

What I’d carry forward

I’d keep the loops-first approach and the extension traits. I’d reach for total_cmp whenever I sort floats, and write every tolerance check as (a - b).abs() < TOL, since a one-sided comparison passes for any value that’s too small. And I’d make “this matrix is square” and “this matrix is symmetric” into types instead of assertions.