TopolearnExplicit loops are Numba's strength, and a few habits decide whether they run at C speed.
With Numba, the loop you would avoid in plain Python becomes the fastest option. A typical compiled kernel looks like C written in Python syntax: preallocate the output with np.empty, loop over indices, and write results element by element.
Habits that keep loops fast:
- Preallocate arrays once instead of appending to lists.
- Loop in memory order. For a C-ordered 2D array, make the last index the inner loop:
for i in range(n): for j in range(m): a[i, j]. Striding across rows in the inner loop wastes cache. - Keep types stable. Initialise accumulators with the right type (
0.0for floats). A variable that is an int in one branch and a float in another forces type unification and can block optimisation. - Avoid temporary arrays inside hot loops.
a[i:i+3].sum()inside a loop allocates a view and calls a function; a short inner loop is often faster. - Fuse passes. Compute several statistics in one loop instead of calling
mean, thenstd, thenmax, each of which reads the whole array.
Loops with data dependencies are where Numba shines over NumPy: running sums with resets, state machines, recurrences like y[i] = a * y[i-1] + x[i], or algorithms with early exit. These are either impossible to vectorize or need several temporary arrays in NumPy.
Arrays passed in are not copied: writing to them inside the function modifies the caller's array, just as in Python.
Example
import numpy as np
from numba import njit
@njit
def ewma(x, alpha):
# y[i] depends on y[i-1]: hard to vectorize, easy to loop
y = np.empty_like(x)
y[0] = x[0]
for i in range(1, x.size):
y[i] = alpha * x[i] + (1 - alpha) * y[i - 1]
return y
@njit
def row_sums(a):
out = np.zeros(a.shape[0])
for i in range(a.shape[0]): # outer: rows
for j in range(a.shape[1]): # inner: contiguous columns
out[i] += a[i, j]
return out
print(ewma(np.array([1.0, 2.0, 3.0, 4.0]), 0.5)) # [1. 1.5 2.25 3.125]
print(row_sums(np.arange(6.0).reshape(2, 3))) # [ 3. 12.]Remember
- Preallocate outputs and write by index
- Make the innermost loop walk contiguous memory
- Recurrences and early exits are where loops beat vectorized NumPy
Watch out for
- Iterating column-first over a C-ordered array in the inner loop.