Why are np.hypot and np.subtract.outer very fast compared to vanilla broadcast and are there faster ways to calculate distance matrix?

distance-matrix, numpy, numpy-ndarray, python, vectorization

Solution

First of all, `d0` and `d1` takes each `50000 x 30000 x 8 = 12 GB` which is pretty big. Make sure you have more than 100 GB of memory because this is what the whole script requires! This is a huge amount of memory. If you do not have enough memory, the operating system will use a storage device (eg. swap) to store excess data which is much slower. Actually, there is no reason Cell-4 is slower than Cell-3 and I guess that you already do not have enough memory to (fully) store `d1` in RAM while `d0` seems to fit (mostly) in memory. There is not difference on my machine when both can fit in RAM (one can also reverse the order of the operations to check this). This also explain why further operation tends to get slower.

That being said, Cells 8+9 are also slower because they create temporary arrays and need more memory passes to compute the result than Cells 3+4+5. Indeed, the expression `np.sqrt(d0**2 + d1**2)` first compute `d0**2` in memory resulting in a new 12 GB temporary array, then compute `d1**2` resulting in another 12 GB temporary array, then perform the sum of the two temporary array to produce another new 12 GB temporary array, and finally compute the square-root resulting in another 12 GB temporary array. This can required up to 48 GB of memory and require 4 read-write memory-bound passes. This is not efficient and do not use the CPU/RAM efficiently (eg. CPU cache).

There is a much faster implementation consisting in doing the whole computation in 1 pass and in parallel using the Numba's JIT. Here is an example:

import numba as nb
@nb.njit(parallel=True)
def distanceMatrix(a, b):
    res = np.empty((a.shape[0], b.shape[0]), dtype=a.dtype)
    for i in nb.prange(a.shape[0]):
        for j in range(b.shape[0]):
            res[i, j] = np.sqrt((a[i, 0] - b[j, 0])**2 + (a[i, 1] - b[j, 1])**2)
    return res

This implementation use 3 times less memory (only 12 GB) and is much faster than the one using `subtract.outer`. Indeed, due to swapping, Cell 3+4+5 takes few minutes while this one takes 1.3 second!

The takeaway is that memory accesses are expensive as well as temporary array. One need to avoid using multiple passes in memory while working on huge buffers and take advantage of CPU caches when the computation performed is not trivial (for example by using array chunks).

Problem

Original source