numba

Numba - High-Performance Python with JIT

Numba makes Python code go fast. It works by decorating your functions with decorators that tell Numba to compile them. It is particularly effective for code that involves heavy numerical loops and NumPy array manipulations.

When to Use

• When NumPy's built-in vectorization isn't enough for your specific algorithm.
• You have complex nested loops that are slow in standard Python.
• You need to write custom "ufuncs" (universal functions) that operate element-wise on arrays.
• High-performance physical simulations (Monte Carlo, N-body, Grid-based solvers).
• Accelerating code for execution on NVIDIA GPUs (CUDA).
• Creating parallelized code that utilizes all CPU cores without the overhead of multiprocessing.

Reference Documentation

Official docs: https://numba.pydata.org/numba-doc/latest/index.html
User Guide: https://numba.pydata.org/numba-doc/latest/user/index.html
Search patterns: @njit, @vectorize, prange, cuda.jit, numba.typed

Core Principles

nopython Mode (@njit)

This is the "gold standard" for Numba. In this mode, Numba compiles the code without using the Python C-API, resulting in maximum speed. If it can't compile (e.g., because of unsupported Python objects), it throws an error.

Just-In-Time (JIT) Compilation

Compilation happens the first time you call the function. The machine code is then cached for subsequent calls.

Array-Oriented

Numba is designed to work with NumPy arrays. It understands their memory layout and can generate highly optimized loops over them.

Quick Reference

Installation

pip install numba

Standard Imports

import numpy as np
from numba import njit, prange, vectorize, guvectorize, cuda

Basic Pattern - Accelerating a Loop

import numpy as np
from numba import njit

# 1. Apply the @njit decorator (alias for @jit(nopython=True))
@njit
def sum_array(arr):
    res = 0.0
    # Standard Python loop that would be slow is now fast as C
    for i in range(arr.shape[0]):
        res += arr[i]
    return res

# 2. Execute
data = np.random.random(1_000_000)
result = sum_array(data) # First call compiles, then runs