Vectorizing difficult problems

All distances between points

The naive version has two loops (each point against each point). The vectorized version uses broadcasting: pontos[:, None, :] (shape (n, 1, d)) less pontos[None, :, :] (shape (1, n, d)) gives all the differences in an array (n, n, d):

import numpy as np

rng = np.random.default_rng(0)
pontos = rng.random((4, 2))

dif = pontos[:, None, :] - pontos[None, :, :]
dist = np.sqrt((dif ** 2).sum(axis=-1))
print(dist.shape, np.allclose(np.diag(dist), 0), np.allclose(dist, dist.T))

lento = np.array([[np.linalg.norm(p - q) for q in pontos] for p in pontos])
print(np.allclose(dist, lento))

(4, 4) True True
True

The distance of each point to itself is zero (diagonal), and the matrix is symmetric. The vectorized version matches the loop version.

The cost: the intermediate array

The dif has n × n × d numbers. For 4 points, nothing. For 10,000 points in 2 dimensions:

n = 10_000
print(n * n * 2 * 8 / 1e9, "GB")

1.6 GB

It's 1.6 GB just for an intermediate, before any calculation, and the program runs out of memory. Vectorization traded time for memory, and this is the most important rule of this chapter: for each vectorization, ask what is the largest intermediate array.

An identity that avoids the intermediate

As |a − b|² = |a|² + |b|² − 2·a·b, it is possible to calculate all distances with one matrix multiplication, without the array (n, n, d). The intermediate result is n × n, and the multiplication uses the fastest linear algebra libraries:

def distancias(X):
    q = (X ** 2).sum(axis=1)
    d2 = q[:, None] + q[None, :] - 2 * X @ X.T
    return np.sqrt(np.maximum(d2, 0))


print(np.allclose(distancias(pontos), dist))

True

The np.maximum(d2, 0) exists because rounding errors can produce slightly negative numbers where the true value is zero, and the square root of a negative would give nan.

When it doesn't fit yet: processing in blocks

If the n × n also doesn't fit, process the problem in blocks: a short loop over slices of rows, with the vectorized operation inside. To find the nearest neighbor of each point in X in Y, the block limits the intermediate to bloco × len(Y) × d:

def mais_proximo(X, Y, bloco=500):
    resultado = np.empty(len(X), dtype=np.int64)
    for inicio in range(0, len(X), bloco):
        parte = X[inicio : inicio + bloco]
        d2 = ((parte[:, None, :] - Y[None, :, :]) ** 2).sum(axis=-1)
        resultado[inicio : inicio + bloco] = d2.argmin(axis=1)
    return resultado


X = rng.random((50, 3))
Y = rng.random((30, 3))
por_forca_bruta = np.array([np.argmin(((Y - x) ** 2).sum(axis=1)) for x in X])
print(np.array_equal(mais_proximo(X, Y, bloco=7), por_forca_bruta))

True

Here the block is 7 on purpose, to allow the loop to run several times and the test to exercise the division. In a real problem, the block size is adjusted to the memory budget.

Grouping without a loop: `reduceat`

To sum values by group when the groups are not small integers (the case of bincount), sort by group and use the add.reduceat, which sums consecutive segments:

grupos = np.array([2, 0, 1, 0, 2, 1])
valores = np.array([10.0, 20.0, 30.0, 40.0, 50.0, 60.0])
ordem = np.argsort(grupos, kind="stable")
g = grupos[ordem]
v = valores[ordem]
inicios = np.r_[0, np.flatnonzero(np.diff(g)) + 1]
print(g[inicios], np.add.reduceat(v, inicios))

[0 1 2] [60. 90. 60.]

After sorting, each group occupies a continuous segment. The np.diff(g) is different from zero where the group changes, and these are the starts of the segments.