From double for-loops to vectorization — a recipe anyone can use
Date: 2026-08-31 | Author: Xuesi Ma |
Bottom line: If you have nested Python for loops filling an array, and every cell is computed from the same standalone math (no cell needs another cell's result), you can usually replace the whole loop nest with one NumPy expression — and it runs ~100–1000× faster for free. The 5-step recipe below (SPOT → RESHAPE → BROADCAST → MASK → LOOP LEFT) shrank a reflector-design search in this project from ~2 days of compute to ~40 minutes, with identical results. When a loop can't be broadcast (recurrences), numba @njit gives the same kind of speed-up — see the BONUS section.
1. Why plain Python loops are slow
NumPy is fast because it runs C loops over whole arrays. A Python loop is slow because the interpreter executes bytecode per iteration, and every small NumPy call on a scalar carries a big fixed overhead (measured here: ~2.8×10-5 s per scalar call). A typical 1000×1000 grid means 106 cells; any per-cell Python work adds up to tens of minutes to hours, with most of that time spent not doing math.
The insight: your loop body usually computes the same expression for every cell. Write that expression once for whole arrays, and NumPy's C loop does the repetition — the “loop” disappears.
2. The recipe — 5 steps
STEP 1 — SPOT it. Look for nested loops whose body is the same expression every iteration, with no iteration using another iteration's output. Quick check: if you could reorder the loops and get the same answer, it vectorizes.
STEP 2 — RESHAPE. Give each loop variable its own array axis with NumPy “newaxis”:
col = x[:, None] # shape (N, 1) -> a COLUMN vector
row = y[None, :] # shape (1, M) -> a ROW vector
If your data is already 2-D (e.g. point coordinates), you add one axis per index: A[:, None, :] and B[None, :, :].
STEP 3 — BROADCAST. Write your expression ONCE. NumPy's broadcasting rule automatically repeats each operand along the axes it's missing (when shapes are compatible), filling the whole N×M table in a single call:
out = np.empty((N, M))
for i in range(N): # BEFORE
for j in range(M):
out[i, j] = f(x[j], y[i])
out = f(x[None, :], y[:, None]) # AFTER (one call)
Watch the argument order: if f's first argument is x, the row ( x[None,:]) must come first — a common gotcha.
STEP 4 — MASK. Inner if/else decisions are just per-element choices; express them with array logic instead:
# BEFORE
if v < 0:
v = 0.0
# AFTER
out = np.where(out < 0, 0.0, out) # or: out[out < 0] = 0.0
This works for any decision that depends only on the cell's own value. A loop break / early exit does not map to broadcasting — that is control flow, handle it separately.
STEP 5 — LOOP LEFT. Only genuinely sequential dimensions must stay loops: recurrences (s[i] needs s[i-1]), cumulative scans, iterative refinement. For those, keep the loop — see the numba bonus below.
Rule of thumb: “If reordering the iterations gives the same answer, it vectorizes.”
3. Always verify — and measure honestly
- Verify the vectorized result against the loop before believing it:
assert np.allclose(out_loop, out_vec). Same math, same answers — this catches argument-order bugs instantly.
- Measure with best-of-N timing, not one call (
time.perf_counter(); take the min of a few repeats). Ignore the first call, which may include import / JIT warm-up.
- When it won't fit in memory (the broadcast table is huge), chunk the loops over one axis — vectorize per block.
4. Example — what it did for the reflector search (z0_parallel.py)
The reflector-candidate search brute-forces a 1000 (r0) × 1000 (φ0) grid = 106 points, solving two equations per point, 2000 times per sweep. It was a textbook case: a nested loop calling the same standalone expression per cell. Measured profile:
| Phase | BEFORE | AFTER (broadcast + batched refine) |
|---|
| grid sampling (80 scalar evals × 106 points) | ~2260 s per call | ~3 s |
| one solver call, end to end | ~46 min | ~1.2–1.8 s |
| full design sweep (2000 calls) | ~2 days | ~30–40 min |
Results were identical — same candidate sets, with z0 agreeing to Δz0 ≤ 2.3×10-11 m (just bisection-vs-brentq round-off, far below the 1e-4 match tolerance).
5. BONUS — numba: when broadcasting can't help
Some loops are genuinely sequential and cannot broadcast — the classic is a recurrence:
s[i] = 0.1*x[i] + 0.9*s[i-1] # s[i] needs s[i-1]
Do not try the tempting one-liner s[1:] = 0.1*x[1:] + 0.9*s[:-1]. It runs without an error yet silently produces garbage, because the right-hand side reads s values that haven't been computed yet (verified: correct [1.063, 1.05, 1.077, …] vs garbage [1.063, 1.05, 0.073, 0.393, …]).
The fix for recurrences, branch-heavy kernels, and awkward loop logic: keep the loop, but compile it with numba @njit — an identical pure-Python loop JIT-compiled to machine code (install with conda install -c conda-forge numba).
Head-to-head on the same 90 000-cell grid (three implementations, all verified identical):
| op | python loop | broadcast (numpy) | numba @njit | winner |
|---|
transcendental exp/cos | 188 ms | 0.34 ms (557×) | 1.04 ms (181×) | broadcast ~3× |
| plain arithmetic | 61 ms | 0.58 ms (105×) | 0.14 ms (438×) | numba ~4× |
Broadcasting wins transcendental-heavy math (NumPy's exp/cos are SIMD-optimized); numba wins plain arithmetic (no temporary arrays, register-resident loop). Both are ≥100× faster than the plain Python loop — the catastrophic gap is Python-loop vs everything else, not broadcast-vs-numba. Practical rule: vectorize by default; reach for @njit when the loop can't be reordered or you need the last factor on arithmetic-heavy work.
6. Resources
An interactive companion notebook and a slide (PNG) walk through the same material live — SPOT/RESHAPE/BROADCAST, the failure case, and the numba benchmark:
- Notebook:
loops_to_vectorization_demo.ipynb (run “Run All” and watch the numbers appear)
- Slide:
loops_to_vectorization_slide.png
Tip: always include an np.allclose assertion and a before/after timing in your own change when you apply this to a new workflow. |