NumPy vectorization - replacing Python loops in scientific computing

I wrote this in 2017 and rewrote it in 2022 after re-running every benchmark. The prose, the numbers and the code below, including the default_rng seeding, are from 2022 pass

I summed a million floats in a Python loop and it took 9.8 ms. The same task as a.sum() took 0.16 ms and matched to every digit I checked. That is 60x for deleting three lines, which made me suspect my timing code before digging into what numpy was really doing under the hood.

To dig deeper, I wrote the same pair for five more operations : six loops I would plausibly have typed, six one-line array calls, keeping identical data on both sides. The measured speedups came back at 3.2x, 9x, 28x, 35x, 60x and 382x!

Timing process

All comparisons are part of bench.py python script in the companion repo at the end. It seeds np.random.default_rng(20170204) and re-seeds before each operation, so the loop and the array call see the same numbers. Timing is a small autotimer rather than timeit: it doubles the batch size until a batch clears 0.05s on time.perf_counter, warms both versions once, then keeps the best of 7 iterations. Top numbers are at n = 1,000,000; the size sweep runs a * b + c alone, from n of 10 up to four million.

The timed loop iterates a plain Python list, a.tolist(), not the array. So every pair is a list loop against one call on the array. Iterating a Numpy array in Python is a third thing, slower again, because every element comes back wrapped in a fresh scalar object.

The price of one add

The sum loop does one add per element and there’s nothing obviously wasteful in there, so the 60x has to come from somewhere other than the arithmetic. At 9.8 ms over a million elements it spends about 9.8 ns per element, and a float add is well under a nanosecond. The interpreter fetches a pointer, checks the object’s type, unboxes it, sends + through dynamic dispatch, boxes the result into a fresh heap object, and steps the iterator. A Python list is a block of pointers to floats scattered across the heap, so each fetch risks a cache miss.

A Numpy array is one contiguous typed buffer :

import numpy as np

rng = np.random.default_rng(20170204)
a = rng.standard_normal(1000000)
a_list = a.tolist()

total = 0.0
for x in a_list:      # 9.8 ms
    total += x

total = a.sum()       # 0.16 ms

The fixed cost per call

The array call is not free either. Before Numpy touches any data it parses arguments, resolves dtypes and shapes, works out broadcasting, allocates an output buffer and sets up the ufunc. That bill is fixed per call, so there is a size below which it swamps the work. For a * b + c at n of 10 the array version measures 0.9x, slower than the loop.

$C_{\text{fixed}}$ is that setup and $c_{\text{py}}$ is the loop’s cost per element, the 9.8 ns from the last section. $c_{\text{elem}}(n)$, Numpy’s cost per element, changes with $n$ : flat while the working set fits in cache, rising toward the bandwidth-bound value once the arrays spill to main memory.

The sweep runs 0.9x at n of 10, 7.5x at 100, 44.6x at 1000, 41.3x at 10,000, a peak of 57.9x near 100,000, then 29.7x at a million and 27.4x at four million. The number I act on is the low end : below a few thousand elements the fixed cost is most of what I am paying.

More data buying less speedup was not what I expected from the sweep. My best guess is cache. Near $10^5$ the operands and the temporary NumPy allocates for a * b still fit a fast cache level, so $c_{\text{elem}}$ sits near its floor. Past that the arrays spill, $c_{\text{elem}}$ climbs toward the rate memory can stream them at, and the ratio settles.

Four of the six hit the same ceiling

Four of the six operations sit under one ceiling. Once the interpreter is out of the way, memory sets the pace.

a * b + c is the plain case at 28x, and Numpy does not fuse it : two ufuncs, a multiply then an add, with a full n-element temporary between them, so at a million elements the expression holds a, b, c, the temporary and the result at once. If the footprint bites, np.multiply(a, b, out=a) then a += c, an explicit out= buffer, or numexpr, which does fuse it. The loop allocates none of that and carries one float at a time.

That 28x and the 29.7x from the sweep are the same operation at the same $n$, two independent runs about 1.7 apart in the multiplier. I have not reconciled them beyond noting that they are two runs.

Cumulative sum is 9x, 30 ms in for-loop against 3.4 ms in Numpy. np.cumsum has to run left to right: out[i] needs out[i-1], so there is one pass, no room to reorder, and a full n-element output written on top of the input read.

The smallest win, and the most instructive number of the six experiments, is the filter. The loop takes 16 ms, numpy takes 5.0 ms, 3.2x.

mask = a > 0.0    # allocation 1 : one bool per element
a[mask]           # allocation 2, after a pass to count the True entries

a > 0.0 allocates a full boolean array first, then boolean indexing then counts the True entries to size the output, allocates that, and makes a second pass to gather the survivors: two allocations, several passes over memory, and almost no arithmetic for a C loop to be faster. Once the mask is an array rather than a single True, the rest follows : mask.sum() counts matches, np.where(mask, a, b) chooses elementwise, and boolean and fancy indexing return a copy where a slice or .T returns a view.

Reductions and broadcasting

The z-score is the last of the bandwidth-bound four, 35x, and it has two modes rather than just one. Reductions are the first. sum, mean, std, var, argmax and argmin are one operation with an axis argument under six names : each collapses the axis its operation runs over.

The mnemonic I use is that the axis you name is the axis that vanishes.

X = rng.standard_normal((1000, 50))   # (n, d)
X.sum(axis=0)   # -> (50,)
X.sum(axis=1)   # -> (1000,)

dtype=np.float32 on a reduction down-casts the accumulator, which a million values will notice; integer sums up-cast to int64 to dodge silent overflow.

Second mode is that of Broadcasting, and it is where the loop model differs : the input array shapes differ and NumPy invents the missing axis.

The rule : compare shapes from the trailing axis, treat a missing leading axis as length 1, and call two axes compatible when they are equal or one of them is 1. So (2,3) with (3,) works and (2,3) with (2,) fails, even though a left-to-right reading says it should work. A stretched axis is a stride-0 axis: the iterator re-reads one address as it walks that dimension instead of copying anything.

Z-scoring a matrix by column needs both shapes at once. The loop is two nested passes and easy to index wrong :

# for-loop version : lists of lists
for j in range(d):
    col = [row[j] for row in rows]
    m = sum(col) / n
    s = (sum((v - m) ** 2 for v in col) / n) ** 0.5
    for i in range(n):
        out[i][j] = (rows[i][j] - m) / s


# Numpy version

(X - X.mean(0)) / X.std(0)      # (n, d) - (1, d), then / (1, d)

169 ms in for loop against 4.8 ms in vectorized version.

Where BLAS takes over

The dot product is the outlier, 382x, 27.5 ms ( for loop ) against 0.072 ms ( numpy ), and it is fast for a different reason than everything above.

Every cheat-sheet I learned from conflated two products. a * b is elementwise, one ufunc writing an n-element output; a @ b contracts over the shared axis. I pick between them by asking what the loop body would have been : one multiply per cell means *, a sum of products means @. (np.dot is the legacy spelling and diverges above two dimensions.)

a * b     # (n,) * (n,) -> (n,)
a @ b     # (n,) @ (n,) -> scalar

@ hands that contraction to a BLAS kernel : cache tiling, SIMD, and a lot of tuning. BLAS libraries are commonly multi-threaded as well, but the repo records neither the library nor a thread count, so I would not lean on threading to explain this run. Part of the 382x is also a heavier baseline : the dot loop does a multiply and a running add per step, so it has further to fall than the elementwise loop.

Floating-point addition is not associative, so a left fold, np.sum and a BLAS dot all round differently. np.sum has been pairwise since Numpy 1.9, its error growing like $O(\varepsilon \log n)$ against $O(\varepsilon n)$ for the naive fold, so over a million standard-normal draws the two agree to about fifteen significant digits and part ways after that, the array result being the more accurate one. For a reduction or a dot, “same computation, just faster” is slightly wrong: it is the same math up to reassociation.

One ratio behind the spread

One observation accounts for all six experiments : the win scales with how much arithmetic the operation does per byte it moves. The dot does a million multiply-adds and returns one number, so it runs compute-bound and gets 382x. The filter reads n floats, writes a mask, reads it back and writes the survivors, doing almost no arithmetic on the way, so it stays bandwidth-bound at 3.2x.

The $n$ itself is a separate lever and a larger one - a double sum over $x_i y_j$ costs $O(n^2)$ as two nested loops and $O(n)$ written as $(\sum_i x_i)(\sum_j y_j)$, because the sums factor. Vectorizing the nested loop keeps the $n^2$; while rewriting the expression removes it.

Browse the code and reproduce the benchmarks NumPy broadcasting docs