NumPy
Advanced 14 min read NumPy

NumPy Performance

From working code to fast working code

Performance comes from vectorization, dtype and alloc discipline: avoid Python loops, downcast types, reuse buffers, watch copies, and know when to move to numba, sparse or GPU kernels.

performancevectorizationdtypecopynumbasparseSIMD
Prerequisites: Vectorized Operations
Before we start

What you will learn

  • Profile numpy code and find the hot spots
  • Choose dtype and memory layout for speed
  • Recognize the escape hatches (numba, sparse, GPU)
Introduction

What is this?

numpy performance is almost entirely about moving work into C loops: vectorize whole-array ops, keep dtype compact, avoid accidental copies and Python-level indirection. When those stop delivering — exotic algorithms, jagged data, per-element logic — the honest next steps are numba/JIT, sparse arrays, or GPU vectors (torch/cupy).

Motivation

Why do you need it?

  • A single per-row np.where save minutes on 10M-row jobs, and that's real money in batch.
  • Float32 over float64 halves bandwidth — the bottleneck for most memory-bound ops.
  • Understanding copies tells you whether your 'fast' code secretly allocates.
  • Senior DS roles are judged on scaling data work without scaling the bill.
Think about it

Mental model

Imagine

A superhighway vs a stop-and-go street grid

Think of

Vectorized numpy vs per-element Python

Vector ops ride the expressway (compiled, SIMD); Python per-element work stops at every intersection. The speed difference is the road network, not the cars.

Mechanics

How it works

  • Compiled ufuncs + blocking + SIMD turn whole-array ops into a tight loop over cache-resident blocks.
  • Smaller dtypes transfer half the bytes from RAM — memory bandwidth is the real limit.
  • Views avoid copies; operations like +, *, where, concatenate allocate fresh output (copy semantics).
  • np.where, advanced indexing and concat create new arrays; favor in-place ops (np.add(out=...)) in loops.
  • Beyond numpy: numba JIT-compiles loops; scipy.sparse stores zeros implicitly; cupy/torch offload to GPU.

The performance ladder.

Profile
Vectorize
Right dtype
No copies
Escape hatch
Hands-on

Examples

Profile-then-fix a hot loop

Profile-then-fix a hot loop
import numpy as np, time
rows = np.random.default_rng(0).lognormal(0, 1, (2_000_000, 5))
 
t0 = time.perf_counter()
tot = 0
for i in range(len(rows)):          # python per-row
    for v in rows[i]:
        tot += v * 2
t1 = time.perf_counter()
fast = (rows * 2).sum()
t2 = time.perf_counter()
print(tot == fast, round(t1-t0, 3), round(t2-t1, 4))

The nested loop touches every element from Python; the ufunc version sums the whole array in C. Ratios of 10-100× are normal at this scale — this single swap is usually the entire performance fix.

Real world

In production

  • A batch feature job downcasts float64 to float32, cutting memory and infer time in half.
  • A serving path reuses an output buffer via out= to stop the allocation chatter at 10k req/s.
  • A sparse click matrix (0.1% density) uses scipy.sparse, turning 10GB dense into 20MB.
Guidelines

Best practices

  • Profile before optimizing — guesswork wastes hours.
  • Measure with large arrays; small sizes hide alloc and dispatch overhead.
  • Prefer out= and in-place ufuncs inside hot loops.
  • Know when to stop: if you're writing Python loops for correctness-critical logic, call numba.
Watch out

Common mistakes

  • Micro-optimizing irrelevant code while the real cost is copies in a hot loop.
  • Forgetting that each numpy op that returns a new array doubles peak memory.
  • Blindly torch-ifying on CPU where numpy + numba wins.
When things break

Debugging tips

  • timeit per op on realistic size; a 1%-of-time op is not the bug.
  • Watch arr.nbytes and peak RSS — memory growth suggests hidden copies.
  • Compare loop vs vectorized numerically (np.allclose) before trusting the fast path.
Interviews

Interview checkpoint

Test yourself

Quick quiz

0/2
1

The biggest single performance win in numpy is usually:

2

Each new array-returning numpy op can double peak memory usage.

Go deeper

Official documentation

NumPy — how to optimize
docs
numba documentation
docs