Introduction
Numba is a powerful library that speeds up Python code by applying just-in-time (JIT) compilation under the hood of your typical NumPy and Python loops. This makes it an appealing choice for statisticians and data scientists who want to push pure Python code close to C-level speeds.
🔑 Key Idea
- Numba uses JIT to compile and optimize functions for significant code runtime improvements.
In this article, we’ll look at the basics of JIT compilation, and how Numba extends NumPy’s universal functions (ufuncs) for often dramatic computational speedups.
Basic JIT Compilation with Numba
JIT, or just-in-time compilation, translates chunks of Python code into optimized machine code right when it’s needed. Rather than relying on Python’s dynamic interpretation at runtime, Numba intercepts certain function calls and compiles them to improve performance significantly.
💡 Critical Concept: @njit vs @jit
- @njit: Short for “no python mode,” this decorator offers the highest performance. It requires strict typing and allows no Python object features, leading to faster execution.
- @jit: Provides more flexibility with Python objects but may be slower because it can’t optimize as aggressively.
Before writing any code, let’s first install Numba with your package manager of choice. Here’s how to do so with pip:
pip install numba
Below is a basic example that computes the standard deviation of a random sample. It compares the speed of Numba’s JIT-compiled function to NumPy’s built-in std. Note that all code in this tutorial is being run on recent consumer grade MacBook.
import time
import numpy as np
from numba import njit
@njit
def fast_std(x):
mean = 0.0
n = len(x)
for i in range(n):
mean += x[i] / n
variance = 0.0
for i in range(n):
diff = x[i] - mean
variance += (diff * diff) / (n - 1)
return np.sqrt(variance)
# Generate data outside timing
x = np.random.randn(100_000_000)
# Warm-up call (compilation happens here)
_ = fast_std(x)
# Now time it
start = time.perf_counter()
numba_result = fast_std(x)
numba_time = time.perf_counter() - start
start = time.perf_counter()
np_result = np.std(x)
np_time = time.perf_counter() - start
print(f"NumPy time: {np_time:.4f}s")
print(f"Numba time: {numba_time:.4f}s")
print(f"Speedup: {np_time/numba_time:.1f}x")
The function fast_std(x) uses a standard two-pass algorithm for computing the mean and variance, illustrating how numerical stability can be maintained while achieving acceleration through Numba’s JIT compiler.
And the result:
NumPy time: 0.3915s Numba time: 0.1733s Speedup: 2.3x
Well, that’s modestly better, but doesn’t seem like the speedup we might have been hoping for. Let’s try another example, this time computing pairwise distances between points in a dataset. This kind of nested-loop problem is harder for NumPy’s built-ins to vectorize cleanly, making it a better candidate for Numba acceleration.
The code:
import time
import numpy as np
from numba import njit
def python_pairwise_distance(A):
"""
Naive Python approach with explicit loops.
A is assumed to be a NumPy array of shape (n, d).
Returns an n x n distance matrix.
"""
n = A.shape[0]
dist_mat = np.zeros((n, n), dtype=np.float64)
for i in range(n):
for j in range(n):
# Compute Euclidean distance
sq_sum = 0.0
for k in range(A.shape[1]):
diff = A[i, k] - A[j, k]
sq_sum += diff * diff
dist_mat[i, j] = np.sqrt(sq_sum)
return dist_mat
@njit
def numba_pairwise_distance(A):
"""
Numba-optimized approach with explicit loops.
Uses @njit to compile the function.
"""
n = A.shape[0]
dist_mat = np.zeros((n, n), dtype=np.float64)
for i in range(n):
for j in range(n):
# Compute Euclidean distance
sq_sum = 0.0
for k in range(A.shape[1]):
diff = A[i, k] - A[j, k]
sq_sum += diff * diff
dist_mat[i, j] = np.sqrt(sq_sum)
return dist_mat
# Generate a random dataset: 2000 points, each with 5 dimensions
np.random.seed(42)
points = np.random.randn(2000, 5)
# Warm up the Numba function to account for compilation overhead
_ = numba_pairwise_distance(points)
# Time the Python version
start = time.perf_counter()
dist_python = python_pairwise_distance(points)
python_time = time.perf_counter() - start
# Time the Numba version
start = time.perf_counter()
dist_numba = numba_pairwise_distance(points)
numba_time = time.perf_counter() - start
print(f"Naive Python time: {python_time:.4f}s")
print(f"Numba time: {numba_time:.4f}s")
print(f"Speedup: {python_time / numba_time:.1f}x")
And the result:
Naive Python time: 6.1277s Numba time: 0.0101s Speedup: 607.8x
There we go, victory achieved! If your problem doesn’t fit neatly into a single vectorized NumPy call or you have lots of nested loops, Numba can be a real game-changer.
Extending NumPy UFuncs with Numba
Numba integrates nicely with NumPy’s vectorized operations through the @vectorize or @guvectorize decorators. By leveraging these, you can write custom universal functions (ufuncs) for elementwise operations in pure Python and still get the speed of optimized code.
💡 Critical Concept: @vectorize vs @guvectorize
- @vectorize creates a ufunc that applies an elementwise operation. Use for straightforward elementwise computations.
- @guvectorize extends the idea of ufuncs to generalized broadcasting. Use for more complex array processing where you need custom logic over sub-arrays and a more flexible mapping between inputs and outputs.
For example, here is a custom elementwise function compared to a naive Python implementation:
import time
import numpy as np
from numba import vectorize
@vectorize
def numba_func(x, y):
"""
Computes x^2 + y^2 elementwise using Numba's @vectorize.
"""
return x**2 + y**2
def python_func(arr1, arr2):
"""
Naive Python loop version (no Numba).
"""
n = len(arr1)
out = np.empty(n, dtype=np.float64)
for i in range(n):
out[i] = arr1[i]**2 + arr2[i]**2
return out
# Generate large random arrays
arr_size = 1_000_000
arr1 = np.random.randn(arr_size)
arr2 = np.random.randn(arr_size)
# Warm up the Numba function once (to pay the compilation cost up front)
_ = numba_func(arr1, arr2)
# Measure time for naive Python loop
start = time.perf_counter()
result_python = python_func(arr1, arr2)
python_time = time.perf_counter() - start
# Measure time for Numba vectorized approach
start = time.perf_counter()
result_numba = numba_func(arr1, arr2)
numba_time = time.perf_counter() - start
# Print results
print(f"Naive Python time: {python_time:.4f}s")
print(f"Numba time: {numba_time:.4f}s")
print(f"Speedup: {python_time / numba_time:.1f}x")
And the results:
Naive Python time: 0.2160s Numba time: 0.0009s Speedup: 246.6x
Here, numba_func is automatically compiled and applied across the input arrays, often outperforming plain Python loops by a large margin.
Review & Looking Ahead
Numba’s just-in-time compilation offers a practical route to speeding up Python and NumPy code for statistical computations. By decorating functions with @njit or writing custom ufuncs, you can quickly transform slow Python loops into efficient, near-C-level routines.
With a careful approach, Numba can be a powerful ally in high-performance data science and statistical computing.
