Numba

Compiler les boucles que vous avez écrites

Certains calculs sont naturellement une boucle : chaque étape dépend de la précédente, comme dans une moyenne pondérée exponentiellement, un drawdown ou le gain d'une option dépendant du chemin. NumPy ne sait pas vectoriser cela, et une boucle Python sur un million d'éléments est lente car chaque étape passe par l'interpréteur. Numba est un compilateur à la volée : vous mettez @njit au-dessus d'une fonction, et au premier appel Numba traduit la boucle Python en code machine pour les types exacts reçus. Les appels suivants tournent à peu près à la vitesse du C, souvent cinquante à cent fois plus vite que la même boucle en Python pur, sans changer le code.

# 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]

Le style qu'aime Numba

En mode nopython, que sélectionne @njit, Numba ne compile qu'un sous-ensemble de Python et refuse le reste plutôt que de retomber silencieusement sur du code lent. Il comprend les nombres, les tableaux NumPy, les tuples, les boucles simples, if et une grande part des fonctions de NumPy. Il ne comprend ni pandas, ni les objets Python quelconques, ni les dictionnaires de types mélangés. Écrivez donc la fonction critique dans le style simple : prendre des tableaux et des scalaires, allouer la sortie avec np.empty ou np.zeros, boucler avec for i in range(n) et garder des accumulateurs scalaires (total, peak). Appelez-la avec les mêmes types d'arguments à chaque fois, car une nouvelle combinaison de types déclenche une nouvelle compilation.

# 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

Temps de compilation, cache et limites de Numba

Le premier appel paie la compilation, souvent une seconde ou plus : chronométrez donc une fonction Numba après l'avoir appelée une fois sur de petites données, et utilisez @njit(cache=True) pour garder le code compilé d'une exécution à l'autre. L'accélération demande assez de travail : une fonction appelée un million de fois sur trois nombres passe son temps dans le surcoût d'appel, pas dans la boucle. Numba aide pour les boucles sur des tableaux ; il ne fait rien pour du code qui est déjà un appel NumPy, pour l'attente d'une base de données ou pour un mauvais algorithme. @njit(parallel=True) avec prange peut répartir des itérations indépendantes sur les cœurs, et fastmath=True échange un peu de rigueur flottante contre de la vitesse, à n'activer qu'après avoir vérifié que les résultats restent dans la tolérance.

# 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 compile une fonction Python en code machine au premier appel ; les boucles qui dépendent des étapes précédentes deviennent rapides sans vectorisation.
• Écrivez pour le mode nopython : tableaux et scalaires, sortie préallouée, for i in range(n), accumulateurs scalaires ; ni pandas ni conteneurs de types mélangés.
• Chronométrez après un appel d'échauffement, utilisez cache=True, n'attendez aucun gain sur du code déjà réduit à un appel NumPy, et vérifiez les résultats contre la version Python simple.

À suivre : Cython, une autre façon de compiler des boucles Python, en ajoutant des types et en construisant un module d'extension.

Lab : ce qui est noté

• Tenir un total courant dans rolling_mean25 pts
• Suivre le sommet dans max_drawdown25 pts
• Compter les croisements dans count_crossovers25 pts
• Remplir un tableau de sortie dans ewma25 pts