Numba

Die Schleifen kompilieren, die Sie geschrieben haben

Manche Berechnungen sind von Natur aus eine Schleife: Jeder Schritt hängt vom vorigen ab, etwa bei einem exponentiell gewichteten Mittel, einem Drawdown oder dem Payoff einer pfadabhängigen Option. NumPy kann das nicht vektorisieren, und eine Python-Schleife über eine Million Elemente ist langsam, weil jeder Schritt durch den Interpreter geht. Numba ist ein Just-in-Time-Compiler: Sie setzen @njit über eine Funktion, und beim ersten Aufruf übersetzt Numba die Python-Schleife für die übergebenen Typen in Maschinencode. Spätere Aufrufe laufen mit ungefähr C-Geschwindigkeit, oft fünfzig- bis hundertmal schneller als dieselbe Schleife in reinem Python, ohne den Code zu ändern.

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

Der Stil, den Numba mag

Im nopython-Modus, den @njit wählt, kompiliert Numba nur eine Teilmenge von Python und weist den Rest zurück, statt stillschweigend auf langsamen Code zurückzufallen. Es versteht Zahlen, NumPy-Arrays, Tupel, einfache Schleifen, if und einen großen Teil der NumPy-Funktionen. Es versteht weder pandas noch beliebige Python-Objekte noch Dictionaries gemischter Typen. Schreiben Sie die heiße Funktion also im schlichten Stil: Arrays und Skalare entgegennehmen, die Ausgabe mit np.empty oder np.zeros anlegen, mit for i in range(n) iterieren und skalare Akkumulatoren (total, peak) führen. Rufen Sie sie jedes Mal mit denselben Argumenttypen auf, denn eine neue Typkombination löst eine neue Kompilierung aus.

# 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

Kompilierzeit, Caching und wo Numba nicht hilft

Der erste Aufruf bezahlt die Kompilierung, oft eine Sekunde oder mehr: Messen Sie eine Numba-Funktion daher, nachdem Sie sie einmal mit kleinen Daten aufgerufen haben, und nutzen Sie @njit(cache=True), um den kompilierten Code zwischen Läufen zu behalten. Der Speed-up braucht genug Arbeit: Eine Funktion, die eine Million Mal mit drei Zahlen aufgerufen wird, verbringt ihre Zeit im Aufrufaufwand, nicht in der Schleife. Numba hilft bei Schleifen über Arrays; es nützt nichts bei Code, der schon ein NumPy-Aufruf ist, beim Warten auf eine Datenbank oder bei einem schlechten Algorithmus. @njit(parallel=True) mit prange kann unabhängige Iterationen auf Kerne verteilen, und fastmath=True tauscht etwas Gleitkomma-Strenge gegen Tempo, was Sie erst aktivieren sollten, nachdem Sie geprüft haben, dass die Ergebnisse in der Toleranz bleiben.

# 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 kompiliert eine Python-Funktion beim ersten Aufruf zu Maschinencode; Schleifen, die von früheren Schritten abhängen, werden ohne Vektorisierung schnell.
• Schreiben Sie für den nopython-Modus: Arrays und Skalare, vorab angelegte Ausgabe, for i in range(n), skalare Akkumulatoren; kein pandas, keine Container gemischter Typen.
• Messen Sie nach einem Aufwärmaufruf, nutzen Sie cache=True, erwarten Sie keinen Gewinn bei Code, der schon ein NumPy-Aufruf ist, und prüfen Sie Ergebnisse gegen die einfache Python-Version.

Als Nächstes: Cython, ein weiterer Weg, Python-Schleifen zu kompilieren, indem man Typen ergänzt und ein Erweiterungsmodul baut.

Lab: Was bewertet wird

• In rolling_mean eine laufende Summe führen25 Pkt.
• In max_drawdown den Höchststand verfolgen25 Pkt.
• In count_crossovers die Kreuzungen zählen25 Pkt.
• In ewma ein Ausgabe-Array füllen25 Pkt.