Numba

Compilar os laços que você escreveu

Alguns cálculos são naturalmente um laço: cada passo depende do anterior, como numa média ponderada exponencialmente, num drawdown ou no payoff de uma opção dependente do caminho. O NumPy não consegue vetorizá-los, e um laço Python sobre um milhão de elementos é lento porque cada passo passa pelo interpretador. O Numba é um compilador just-in-time: você coloca @njit sobre uma função e, na primeira chamada, o Numba traduz o laço Python em código de máquina para os tipos exatos recebidos. As chamadas seguintes rodam em velocidade próxima à do C, muitas vezes cinquenta a cem vezes mais rápido que o mesmo laço em Python puro, sem mudar o 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]

O estilo de que o Numba gosta

No modo nopython, que o @njit seleciona, o Numba compila apenas um subconjunto do Python e recusa o resto em vez de recair silenciosamente em código lento. Ele entende números, arrays do NumPy, tuplas, laços simples, if e grande parte das funções do NumPy. Não entende pandas, objetos Python arbitrários nem dicionários de tipos mistos. Então escreva a função crítica no estilo simples: receba arrays e escalares, aloque a saída com np.empty ou np.zeros, itere com for i in range(n) e mantenha acumuladores escalares (total, peak). Chame-a sempre com os mesmos tipos de argumentos, pois uma nova combinação de tipos dispara uma nova compilação.

# 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

Tempo de compilação, cache e onde o Numba não ajuda

A primeira chamada paga a compilação, muitas vezes um segundo ou mais, então cronometre uma função Numba depois de chamá-la uma vez com poucos dados, e use @njit(cache=True) para guardar o código compilado entre execuções. O ganho exige trabalho suficiente: uma função chamada um milhão de vezes com três números gasta o tempo no custo da chamada, não no laço. O Numba ajuda em laços sobre arrays; não faz nada por código que já é uma chamada NumPy, por esperar um banco de dados nem por um mau algoritmo. @njit(parallel=True) com prange pode espalhar iterações independentes pelos núcleos, e fastmath=True troca um pouco do rigor de ponto flutuante por velocidade, o que só deve ser ativado após verificar que os resultados seguem dentro da tolerância.

# 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 uma função Python em código de máquina na primeira chamada; laços que dependem de passos anteriores ficam rápidos sem vetorizar.
• Escreva para o modo nopython: arrays e escalares, saída pré-alocada, for i in range(n), acumuladores escalares; sem pandas nem contêineres de tipos mistos.
• Cronometre após uma chamada de aquecimento, use cache=True, não espere ganho em código que já é uma chamada NumPy e confira os resultados com a versão Python simples.

A seguir: Cython, outra forma de compilar laços Python, adicionando tipos e construindo um módulo de extensão.

Lab: o que é avaliado

• Manter um total corrente em rolling_mean25 pts
• Acompanhar o pico em max_drawdown25 pts
• Contar os cruzamentos em count_crossovers25 pts
• Preencher um array de saída em ewma25 pts