TopolearnMost slow NumPy code is a hidden Python loop or a pile of large temporaries; these patterns fix both.
The goal of vectorization is to move every per-element loop into compiled code. A few patterns cover most cases:
- Conditional logic: replace
ifin a loop withnp.where, masks ornp.select(conditions, choices)for several branches. - Running state:
np.cumsum,np.cumprod,np.maximum.accumulatereplace accumulators;np.diffreplaces "current minus previous". - Group-by sums: with integer group labels,
np.bincount(labels, weights=values)sums per group in one call. For contiguous segments,np.add.reduceat(values, starts). - Pairwise work: broadcasting with
a[:, None]andb[None, :], chunked if the result would be huge. - Lookups: fancy indexing into a table, or
np.searchsortedinto sorted keys.
np.vectorize is not one of these patterns. It is a convenience wrapper that calls your Python function once per element; it runs at Python-loop speed.
Once code is vectorized, the next cost is temporaries. y = a * b + c * d allocates three full-size arrays. For big arrays memory bandwidth, not arithmetic, dominates, so reuse buffers: np.multiply(a, b, out=y); y += c * d. Chunking work into blocks that fit in cache (say 10-100k elements) can also help.
Always measure, ideally with %timeit or timeit, on realistic sizes. For small arrays (tens of elements) NumPy's per-call overhead of around a microsecond can make plain Python faster. And some algorithms are inherently sequential, like a recurrence where each step depends on the previous result in a non-linear way. For those, vectorizing creates contortions; reach for Numba, Cython or a library routine instead.
Example
import numpy as np
rng = np.random.default_rng(0)
# 1) branches: piecewise tax rate instead of if/elif in a loop
income = rng.integers(0, 200_000, size=1_000_000)
rate = np.select([income < 20_000, income < 80_000], [0.0, 0.2], default=0.4)
# 2) group sums with integer labels
store = rng.integers(0, 5, size=income.size)
per_store = np.bincount(store, weights=income, minlength=5)
print(per_store.shape) # (5,)
# 3) running state: drawdown from running peak
price = 100 + rng.standard_normal(1000).cumsum()
drawdown = price / np.maximum.accumulate(price) - 1
print(drawdown.min() <= 0) # True
# 4) fewer temporaries with out= and in-place ops
a, b, c = (rng.random(1_000_000) for _ in range(3))
y = np.empty_like(a)
np.multiply(a, b, out=y)
y += c # y = a*b + c, one buffer
print(np.allclose(y, a * b + c)) # TrueRemember
- Use masks,
whereandselectfor branches;cumsum/diff/accumulatefor running state bincount(labels, weights=...)is a fast group-sumnp.vectorizeis a loop in disguise- Reduce temporaries with
out=and in-place ops, and always measure
Watch out for
- Assuming vectorized is always faster: for tiny arrays or sequential recurrences it can be slower or impossible.