NumPy Performance Optimization

Measure Before Optimizing

timeit gives speed, tracemalloc gives memory — both from the standard library.

import timeit
import tracemalloc
import numpy as np

rng = np.random.default_rng(42)
A = rng.standard_normal((1000, 1000))
B = rng.standard_normal((1000, 1000))

best = min(timeit.repeat(lambda: A @ B, number=10, repeat=3)) / 10
print(f"A @ B: {best * 1000:.2f} ms per call")

tracemalloc.start()
C = A + B
_, peak = tracemalloc.get_traced_memory()
tracemalloc.stop()
print(f"A + B peak allocation: {peak / 1e6:.1f} MB")

$ python measure.py
A @ B: 14.20 ms per call
A + B peak allocation: 8.0 MB

Taking the minimum of several repeats filters out background noise. The allocation for A + B is exactly one new 1000 × 1000 float64 array: 8 MB. Numbers throughout this lesson come from one machine; measure on yours before acting on them.

float32 vs float64

import timeit
import numpy as np

X64 = np.random.default_rng(0).standard_normal((2000, 2000))
X32 = X64.astype(np.float32)
print(f"memory: {X64.nbytes / 1e6:.0f} MB vs {X32.nbytes / 1e6:.0f} MB")

t64 = min(timeit.repeat(lambda: X64 @ X64, number=3, repeat=3)) / 3
t32 = min(timeit.repeat(lambda: X32 @ X32, number=3, repeat=3)) / 3
print(f"matmul float64 {t64:.3f} s, float32 {t32:.3f} s, ratio {t64 / t32:.1f}x")

print(np.float32(1) + np.float32(1e-8) == 1, np.float64(1) + 1e-8 == 1)

$ python dtypes.py
memory: 32 MB vs 16 MB
matmul float64 0.111 s, float32 0.050 s, ratio 2.2x
True False

Half the bytes, and on this machine a noticeably faster matrix product — BLAS routines process more float32 values per instruction. The last line is the cost: float32 carries about 7 significant digits, so adding 1e-8 to 1 changes nothing, while float64 still registers it.

A Myth About Upcasting

A common tip says a float32 array silently turns into float64 in (x - x.mean()) / x.std(). Checking it on NumPy 2:

import numpy as np

x = np.arange(5, dtype=np.float32)
print((x - x.mean()) / x.std())
print(((x - x.mean()) / x.std()).dtype)
print((x * 2.5).dtype)
print((x + np.float64(1)).dtype, (x + np.ones(5)).dtype)

$ python promotion.py
[-1.4142135  -0.70710677  0.          0.70710677  1.4142135 ]
float32
float32
float64 float64

The whole expression stays float32: x.mean() of a float32 array is a float32 scalar, and a plain Python float (2.5) adopts the array’s dtype under NumPy 2’s promotion rules. What does upcast is mixing in an explicit NumPy float64 — a scalar or an array — as the last line shows.

Memory Layout and Copies

import timeit
import numpy as np

A = np.random.default_rng(0).standard_normal((3000, 3000))
F = np.asfortranarray(A)
print(A.flags["C_CONTIGUOUS"], F.flags["F_CONTIGUOUS"])

for name, arr in [("C-order", A), ("F-order", F)]:
    t = min(timeit.repeat(lambda: arr.copy(order="C"), number=5, repeat=3)) / 5
    print(f"{name} -> C copy: {t * 1000:.1f} ms")

half = A[:, ::2]
print(half.flags["C_CONTIGUOUS"], np.ascontiguousarray(half).flags["C_CONTIGUOUS"])

$ python layout.py
True True
C-order -> C copy: 18.9 ms
F-order -> C copy: 70.3 ms
False True

Copying a C-ordered array into C order is a straight memory copy; copying the Fortran-ordered version has to reorder every element, which jumps around memory and runs slower. Slicing with a step (A[:, ::2]) produces a non-contiguous view; np.ascontiguousarray makes a compact copy, worth doing once before an array is reused many times.

Reusing Buffers with out=

import timeit
import numpy as np

X = np.random.default_rng(42).standard_normal((2000, 2000))
buf = np.empty_like(X)

def fresh():
    return np.sqrt(X * X + 1.0)

def reuse():
    np.multiply(X, X, out=buf)
    np.add(buf, 1.0, out=buf)
    return np.sqrt(buf, out=buf)

print(np.allclose(fresh(), reuse()))
for f in (fresh, reuse):
    t = min(timeit.repeat(f, number=10, repeat=3)) / 10
    print(f"{f.__name__:5s} {t * 1000:.1f} ms")

$ python out_buffers.py
True
fresh 23.4 ms
reuse 8.2 ms

fresh allocates a new 32 MB array for each of its three steps; reuse writes all three into one preallocated buffer. Same result, and on this run nearly three times faster — for large arrays, allocating and first-touching fresh memory costs about as much as the arithmetic itself. It pays off most in loops that run the same computation many times.

einsum

import timeit
import numpy as np

rng = np.random.default_rng(0)
A = rng.standard_normal((64, 128, 128))
B = rng.standard_normal((64, 128, 128))

C1 = np.einsum("bij,bjk->bik", A, B)
print(C1.shape, np.allclose(C1, A @ B))
print(np.einsum("ii->", np.eye(5)), np.einsum("ij->", np.ones((3, 4))))

for name, f in [("A @ B", lambda: A @ B),
                ("einsum", lambda: np.einsum("bij,bjk->bik", A, B))]:
    print(f"{name:7s} {min(timeit.repeat(f, number=3, repeat=3)) / 3 * 1000:.1f} ms")

$ python einsum_demo.py
(64, 128, 128) True
5.0 12.0
A @ B   6.8 ms
einsum  35.0 ms

einsum writes contractions as index strings: bij,bjk->bik is a batch of matrix products, ii-> is a trace, ij-> sums everything. It is great for clarity, but for a plain batched matrix product the dedicated @ hands the work to BLAS and was about five times faster here. Reach for einsum when the contraction has no single-function equivalent (or try optimize=True, and measure).

When a Loop Is Unavoidable

Some algorithms — a recurrence where each step needs the previous result, or irregular per-item logic — don’t vectorize. The usual answers are a JIT compiler such as Numba (@numba.njit on a plain Python loop) or writing that one function in Cython. Install them separately, and measure them the same way as everything above.

Practice

1. Measure how long np.sort takes on one million float64 values versus the same values as float32.
2. Predict the dtype of np.ones(3, dtype=np.int8) + 300, then run it.
3. Rewrite y = 3 * x + 1 for a large x so that it allocates only one output array.

import timeit
import numpy as np

x64 = np.random.default_rng(0).standard_normal(1_000_000)
x32 = x64.astype(np.float32)
for name, x in [("float64", x64), ("float32", x32)]:
    t = min(timeit.repeat(lambda: np.sort(x), number=5, repeat=3)) / 5
    print(f"{name} sort {t * 1000:.1f} ms")

$ python practice_sort.py
float64 sort 11.6 ms
float32 sort 6.1 ms

import numpy as np

try:
    print((np.ones(3, dtype=np.int8) + 300).dtype)
except OverflowError as e:
    print("OverflowError:", e)

$ python practice_int8.py
OverflowError: Python integer 300 out of bounds for int8

import numpy as np

x = np.arange(5.0)
y = np.multiply(x, 3)
np.add(y, 1, out=y)
print(y)

$ python practice_inplace.py
[ 1.  4.  7. 10. 13.]