pythonforbio.
[WASM idle]
NumPy15/18

Vectorization

Starting sandbox…
Beginnerlesson3 tests

Vectorization

In one line: replace the Python loop with a whole-array expression so the loop happens in compiled code instead of the interpreter.

Why the loop is slow

A Python for loop over a NumPy array pays, per element:

  1. an interpreter bytecode dispatch
  2. boxing the C value into a Python object
  3. a dynamic type lookup for the operation
  4. allocating a result object
  5. unboxing to store it back

That is ~50–100 ns of overhead per element against ~1 ns of actual arithmetic. NumPy pays those costs once for the whole array and then runs a tight, often SIMD-vectorised, C loop.

x = rng.random(1_000_000)

out = np.empty_like(x)                 # ~400 ms
for i in range(len(x)):
    out[i] = x[i] * 2 + 1

out = x * 2 + 1                         # ~2 ms

Note the loop version is also paying indexing overhead on both read and write, which is why it is worse than a pure-Python list comprehension in some cases.

The translation table

Loop pattern Vectorised form
for x in a: out.append(f(x)) f(a) if f is a ufunc
if/else per element np.where(cond, a, b)
accumulate a total a.sum(), a.sum(axis=...)
running total np.cumsum(a)
count matching mask.sum()
find first match np.argmax(mask) (0 if none — check mask.any())
build a filtered list a[mask]
lookup table table[codes] — see Fancy Indexing
pairwise over two arrays broadcasting — see Broadcasting
grouped sum np.bincount(codes, weights=vals)
sliding window np.lib.stride_tricks.sliding_window_view
chained if/elif/else np.select([c1, c2], [v1, v2], default=v3)

Worked conversions

Conditional classification:

# loop
labels = []
for fc, p in zip(log2fc, padj):
    if p < 0.05 and fc > 1:      labels.append("up")
    elif p < 0.05 and fc < -1:   labels.append("down")
    else:                        labels.append("ns")

# vectorised
sig = padj < 0.05
labels = np.select([sig & (log2fc > 1), sig & (log2fc < -1)],
                   ["up", "down"], default="ns")

Sliding window (GC content in 100 bp windows):

from numpy.lib.stride_tricks import sliding_window_view

is_gc = np.isin(seq_arr, ["G", "C"])
windows = sliding_window_view(is_gc, 100)      # (L-99, 100) — a VIEW, no copy
gc = windows.mean(axis=1)

sliding_window_view creates the windows by stride manipulation, so it costs nothing until you reduce over it. See Memory Layout and Strides.

Grouped counting:

# variants per 1 Mb bin
bins = positions // 1_000_000
counts_per_bin = np.bincount(bins, minlength=n_bins)

# summed coverage per bin
depth_per_bin = np.bincount(bins, weights=depths, minlength=n_bins)

When NOT to vectorise

Vectorization is not free and not always right.

  • Memory. A vectorised expression materialises every intermediate. ((a + b) * c - d) / e allocates four temporaries the size of a. On a 10 GB array, a chunked loop wins. Use out= on ufuncs to reuse buffers.
  • Genuinely sequential logic. Recurrences where step n depends on step n−1 (HMM forward algorithm, alignment DP) cannot be vectorised along that axis. Vectorise the other axis — all states at once, stepping through positions — or use numba.
  • Early exit. A loop that breaks on the first match beats scanning the whole array.
  • Readability. A five-line loop that runs in 2 ms on 100 elements does not need to become an unreadable one-liner. Profile before you obfuscate.
  • The array is small. Below ~1000 elements, NumPy's per-call overhead can exceed the savings.

The escape hatches, in order of preference

  1. Restructure into existing NumPy operations. Usually possible with np.where, np.select, bincount, reduceat, or sliding_window_view.
  2. scipyscipy.ndimage, scipy.signal, scipy.spatial have compiled implementations of many "obviously loopy" operations.
  3. numba@njit on a plain Python loop, JIT-compiled. Often the best answer for DP algorithms. Frequently faster than vectorised NumPy because it avoids temporaries.
  4. Cython / C extension — when numba cannot express it.

np.vectorize is not on this list. It is a Python loop. See ufuncs.

Profiling before optimising

%timeit expr                          # IPython
python -m cProfile -s cumtime script.py

The usual finding is that 90% of the time is in one line, and it is not the line you expected. In bioinformatics it is very often file parsing, not computation.

Common mistakes

  • Vectorising the wrong axis on a recurrence, then wondering why the answer is wrong.
  • Building huge intermediates. Watch (n, n, d) shapes from broadcasting.
  • np.vectorize, thinking it compiles.
  • Appending inside a loop even in the "vectorised" version. See Array Creation.
  • .apply() in pandas as a vectorisation strategy. df.apply(f, axis=1) is a Python loop over rows. See pandas Performance.
  • Optimising before measuring. The bottleneck is rarely where you think.

See also

ufuncs · Broadcasting · Axes and Reductions · Fancy Indexing · Memory Layout and Strides · pandas Performance · Split-Apply-Combine

Test cases · 3

#viainputexpected stdout
1stdinATGCATGC0.5
2stdinGGGGCCCC1.0
3stdinATATATAT0.0

Hints · 2

01Hint

np.frombuffer(seq.encode(), dtype='S1') turns a string into an array of bytes.

02Hint

Compare the byte array to b'G' and b'C', then use np.mean on the OR.

starter

No output yet — run the code to populate this drawer.

Structured ArraysViews vs Copies