Numba

Compiling the loops you wrote

Some calculations are naturally a loop: each step depends on the one before, as in an exponentially weighted average, a drawdown or a path-dependent option payoff. NumPy cannot vectorise these, and a Python loop over a million elements is slow because every step goes through the interpreter. Numba is a just-in-time compiler: you put @njit above a function, and on its first call Numba translates the Python loop to machine code for the exact types it was given. Later calls run at roughly C speed, often fifty to a hundred times faster than the same loop in plain Python, with no change to the 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]

The style Numba likes

In nopython mode, which @njit selects, Numba compiles only a subset of Python and refuses the rest rather than silently falling back to slow code. It understands numbers, NumPy arrays, tuples, simple loops, if and a large part of NumPy's functions. It does not understand pandas, arbitrary Python objects, or dictionaries of mixed types. So write the hot function in the plain style: take arrays and scalars, allocate the output with np.empty or np.zeros, loop with for i in range(n), and keep scalar accumulators (total, peak). Call it with the same argument types each time, because a new combination of types triggers a new 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

Compile time, caching and where Numba does not help

The first call pays for compilation, often a second or more, so time a Numba function after calling it once on small data, and use @njit(cache=True) to keep the compiled code between runs. The speed-up needs enough work: a function called a million times on three numbers spends its time in call overhead, not in the loop. Numba helps with loops over arrays; it does nothing for code that is already one NumPy call, for waiting on a database, or for a bad algorithm. @njit(parallel=True) with prange can spread independent iterations over cores, and fastmath=True trades a little floating-point strictness for speed, which you should only enable after checking that results stay within tolerance.

# 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 compiles a Python function to machine code on its first call; loops that depend on earlier steps become fast without vectorising.
• Write for nopython mode: arrays and scalars, a pre-allocated output, for i in range(n), scalar accumulators; no pandas or mixed-type containers.
• Time after a warm-up call, use cache=True, expect no gain on code that is already one NumPy call, and check results against the plain Python version.

Next up: Cython, another way to compile Python loops, by adding types and building an extension module.

Lab · what's graded

• Keep a running total in rolling_mean25 pts
• Track the peak in max_drawdown25 pts
• Count the crossings in count_crossovers25 pts
• Fill an output array in ewma25 pts