Numba Jit Acceleration

Numba — Aceleración JIT de código numérico Python

Resumen

Numba (11.1k⭐, BSD-2, mantenido por Anaconda) es un compilador Just-In-Time que traduce un subconjunto de Python numérico (con NumPy) a máquina nativa vía LLVM. Acelera bucles y cálculo numérico 10-1000x sin salir de Python ni escribir C.

Cuándo usarlo: simulaciones Monte Carlo, física/cálculos por píxel, procesamiento de nubes de puntos (LiDAR), sombras solares, estadística sobre arrays grandes — todo lo CPU-bound en los proyectos de datos de David.

Instalación

pip install numba
# Versiones de numba están casadas con versiones de Python y NumPy:
# comprobar tabla de compatibilidad si hay ImportError tras actualizar numpy

Patrones principales

1. Compilar una función (@njit)

from numba import njit
import numpy as np

@njit(fastmath=True)          # fastmath: reordena ops flotantes (más rápido, menos estricto)
def shadow_loop(sun_az, sun_el, grid):
    n = grid.shape[0]
    out = np.zeros(n)
    for i in range(n):
        out[i] = grid[i] * np.sin(sun_el) / np.tan(np.radians(sun_az + i))
    return out

• @njit = @jit(nopython=True): si la función usa algo no soportado, falla en tiempo de compilación en vez de degradar a modo object lento. Siempre njit, nunca jit pelado.
• Primera llamada compila (~0.1-1s). Para producción: cache=True.

2. Paralelizar bucles (parallel=True + prange)

from numba import prange

@njit(parallel=True)
def pairwise_distances(X, Y):
    n, m = X.shape[0], Y.shape[0]
    D = np.empty((n, m))
    for i in prange(n):                    # prange, NO range
        for j in range(m):
            D[i, j] = np.sqrt(((X[i] - Y[j]) ** 2).sum())
    return D

• Usa prange para el bucle exterior independiente. NUMBA_NUM_THREADS limita hilos.
• Reducciones (sumas, máx) soportadas; escrituras con índice derivado del bucle (out[i] = ...) soportadas; escrituras aleatorias = race condition.

3. UFuncs personalizados (@vectorize / @guvectorize)

from numba import vectorize, int64

@vectorize([int64(int64, int64)])
def gcd(a, b):
    while b:
        a, b = b, a % b
    return a

gcd(np.array([12, 18]), np.array([8, 24]))   # funciona con broadcasting como ufunc NumPy

4. GPU CUDA

from numba import cuda

@cuda.jit
def add_kernel(a, b, out):
    i = cuda.grid(1)
    if i < a.size:
        out[i] = a[i] + b[i]

d_a = cuda.to_device(a); d_out = cuda.device_array_like(a)
add_kernel[blocks, threads_per_block](d_a, d_b, d_out)

• Requiere NVIDIA GPU + driver CUDA. @cuda.jit no compila si hay objeto no soportado: mismo principio nopython.
• Alternativa más simple sin escribir kernels: @jit(target_backend='cuda') para bucles paralelizables.

Pitfalls

• No soporta casi nada del ecosistema: nada de pandas, dicts heterogéneos, funciones lambda dentro, atributos de objetos arbitrarios. Solo NumPy, scalars, tuplas/listas homogéneas, arrays tipados.
• np.linalg limitado: np.dot sí; svd/lstsq solo en versiones recientes y con restricciones.
• Compilación por signatura: cada combinación de tipos de argumentos compila una variante. Pasar arrays no-contiguos (A.T) cambia la signatura y recompila → usar np.ascontiguousarray.
• fastmath=True rompe IEEE: 0*inf y comparaciones con NaN pueden cambiar de resultado.
• NumPy 2.x: numba antiguos (<0.59) reventaban con numpy≥2. Actualizar numba si se actualiza numpy.
• Medir siempre: time.perf_counter() sobre 2ª llamada (la primera paga la compilación), o warmup antes de cronometrar.

Verificación

python -c "import numba, numpy as np; from numba import njit; f = njit(lambda x: x.sum()); print('OK', f(np.arange(1000)), numba.__version__)"

Cuándo NO usar Numba

• Código I/O-bound o con pandas en el hot path → vectorizar primero con NumPy/pandas.
• Necesitas depuración interactiva → separar el kernel numérico en su propia función @njit.
• Cálculo con tipos arbitrarios/objetos → Cython o Python puro.

Referencia

• Repo: https://github.com/numba/numba
• Docs: https://numba.readthedocs.io/en/stable/
• Examples: https://github.com/numba/numba-examples (notebooks en Binder)