Numba

Compilar los bucles que escribiste

Algunos cálculos son naturalmente un bucle: cada paso depende del anterior, como en un promedio ponderado exponencialmente, un drawdown o el pago de una opción dependiente de la trayectoria. NumPy no puede vectorizarlos, y un bucle de Python sobre un millón de elementos es lento porque cada paso pasa por el intérprete. Numba es un compilador justo a tiempo: pones @njit sobre una función y, en la primera llamada, Numba traduce el bucle de Python a código máquina para los tipos exactos recibidos. Las llamadas posteriores corren a velocidad casi de C, a menudo de cincuenta a cien veces más rápido que el mismo bucle en Python puro, sin cambiar el código.

# An exponentially weighted average is a loop: each value depends on the one before
import numpy as np

try:
    from numba import njit                           # one decorator compiles the function when it is installed
except ImportError:
    def njit(func):                                  # without Numba the same code simply runs as Python
        return func

@njit
def ewma(values, alpha):
    out = np.empty(len(values))
    out[0] = values[0]
    for i in range(1, len(values)):
        out[i] = alpha * values[i] + (1.0 - alpha) * out[i - 1]
    return out

print(ewma(np.array([1.0, 2.0, 3.0]), 0.5))          # [1.   1.5  2.25]

El estilo que le gusta a Numba

En modo nopython, que selecciona @njit, Numba compila solo un subconjunto de Python y rechaza el resto en lugar de recaer en silencio en código lento. Entiende números, arrays de NumPy, tuplas, bucles simples, if y gran parte de las funciones de NumPy. No entiende pandas, objetos arbitrarios de Python ni diccionarios de tipos mezclados. Así que escribe la función crítica en estilo simple: recibe arrays y escalares, reserva la salida con np.empty o np.zeros, itera con for i in range(n) y mantén acumuladores escalares (total, peak). Llámala siempre con los mismos tipos de argumentos, porque una nueva combinación de tipos dispara una nueva compilación.

# A drawdown is a running peak and a running worst case: two scalars and one loop
import numpy as np

try:
    from numba import njit
except ImportError:
    def njit(func):
        return func

@njit
def max_drawdown(equity):
    peak = 0.0
    worst = 0.0
    for i in range(len(equity)):
        value = equity[i]
        if value > peak:
            peak = value
        drawdown = (peak - value) / peak
        if drawdown > worst:
            worst = drawdown
    return worst

curve = np.array([100.0, 120.0, 90.0, 130.0, 117.0])
print(max_drawdown(curve))                           # 0.25   from 120 down to 90
print(max_drawdown(np.array([1.0, 2.0, 3.0])))       # 0.0    a curve that only rises

Tiempo de compilación, caché y dónde Numba no ayuda

La primera llamada paga la compilación, a menudo un segundo o más, así que cronometra una función de Numba tras llamarla una vez con pocos datos, y usa @njit(cache=True) para conservar el código compilado entre ejecuciones. La aceleración necesita suficiente trabajo: una función llamada un millón de veces con tres números gasta su tiempo en la sobrecarga de la llamada, no en el bucle. Numba ayuda con bucles sobre arrays; no hace nada por código que ya es una llamada de NumPy, por esperar una base de datos ni por un mal algoritmo. @njit(parallel=True) con prange puede repartir iteraciones independientes entre núcleos, y fastmath=True cambia algo de rigor en coma flotante por velocidad, que solo debes activar tras comprobar que los resultados siguen dentro de la tolerancia.

# Counting crossings: Numba compiles any loop like this, so keep the logic in scalars
import numpy as np

try:
    from numba import njit
except ImportError:
    def njit(func):
        return func

@njit(cache=False)
def count_crossovers(fast, slow):
    count = 0
    previous = 0.0
    for i in range(len(fast)):
        gap = fast[i] - slow[i]
        if i > 0 and previous * gap < 0:
            count += 1
        previous = gap
    return count

fast = np.array([1.0, 3.0, 2.0, 1.0, 4.0])
slow = np.full(5, 2.0)
print(count_crossovers(fast, slow))                  # 2     below, above, touch, below, above: two strict sign changes

# Options change how a function is compiled; here they are ignored if Numba is not installed
import numpy as np

try:
    from numba import njit, prange
except ImportError:
    prange = range

    def njit(*args, **kwargs):
        if args and callable(args[0]):
            return args[0]
        return lambda func: func

@njit(cache=True)                                    # keep the compiled code on disk between runs
def total_notional(prices, quantities):
    total = 0.0
    for i in prange(len(prices)):                    # prange runs in parallel only with parallel=True
        total += prices[i] * quantities[i]
    return total

prices = np.array([184.25, 410.5, 99.0])
quantities = np.array([100.0, 50.0, 200.0])
print(total_notional(prices, quantities))            # 58750.0
print(total_notional(prices, quantities) == float(prices @ quantities))     # True  same answer as the vectorised dot product

• @njit compila una función de Python a código máquina en la primera llamada; los bucles que dependen de pasos anteriores se vuelven rápidos sin vectorizar.
• Escribe para el modo nopython: arrays y escalares, salida reservada, for i in range(n), acumuladores escalares; sin pandas ni contenedores de tipos mezclados.
• Cronometra tras una llamada de calentamiento, usa cache=True, no esperes ganancia en código que ya es una llamada de NumPy y comprueba los resultados con la versión Python simple.

A continuación: Cython, otra forma de compilar bucles de Python, añadiendo tipos y construyendo un módulo de extensión.

Lab: qué se califica

• Mantener un total corriente en rolling_mean25 pts
• Seguir el máximo en max_drawdown25 pts
• Contar los cruces en count_crossovers25 pts
• Rellenar un array de salida en ewma25 pts