Numba

编译你写的循环

有些计算天生就是循环：每一步依赖上一步，比如指数加权平均、回撤，或路径依赖期权的收益。NumPy 无法把它们向量化，而对一百万个元素的 Python 循环很慢，因为每一步都要经过解释器。Numba 是一个即时编译器：在函数上方加 @njit，首次调用时 Numba 会针对实际传入的类型把 Python 循环翻译成机器码。之后的调用以接近 C 的速度运行，通常比纯 Python 中的同一个循环快五十到一百倍，而代码无需改动。

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

Numba 喜欢的风格

在 @njit 选择的 nopython 模式下，Numba 只编译 Python 的一个子集，对其余部分直接拒绝，而不是悄悄退回到慢代码。它理解数字、NumPy 数组、元组、简单循环、if 以及大部分 NumPy 函数；不理解 pandas、任意 Python 对象或混合类型的字典。所以要用朴素的风格来写热点函数：接收数组和标量，用 np.empty 或 np.zeros 分配输出，用 for i in range(n) 循环，并保持标量累加器（total、peak）。每次用相同类型的参数调用，因为新的类型组合会触发重新编译。

# 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

编译耗时、缓存，以及 Numba 帮不上忙的地方

第一次调用要付出编译代价，往往一秒以上，所以给 Numba 函数计时应在它用小数据调用过一次之后，并用 @njit(cache=True) 在多次运行之间保留编译结果。加速需要足够的工作量：一个对三个数调用一百万次的函数，时间都花在调用开销上，而不是循环里。Numba 对数组上的循环有帮助；对本来就是一次 NumPy 调用的代码、等待数据库、或糟糕的算法，它都无能为力。@njit(parallel=True) 配合 prange 可以把相互独立的迭代分散到多个核心，fastmath=True 用一点浮点严格性换取速度，只有在确认结果仍在容差内之后才应启用。

# 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 在首次调用时把 Python 函数编译成机器码；依赖前面步骤的循环无需向量化也能变快。
• 为 nopython 模式编写：数组与标量、预先分配的输出、for i in range(n)、标量累加器；不用 pandas 或混合类型容器。
• 预热调用之后再计时，使用 cache=True，对本就是一次 NumPy 调用的代码不要期待收益，并用朴素的 Python 版本核对结果。

下一节：Cython，通过添加类型并构建扩展模块来编译 Python 循环的另一种方式。

实验：评分内容

• 在 rolling_mean 中维护滚动总和25 分
• 在 max_drawdown 中跟踪峰值25 分
• 在 count_crossovers 中统计交叉次数25 分
• 在 ewma 中填充输出数组25 分