Ускорение финансовых индикаторов на Python с помощью Numba JIT

Суть: Использование декоратора @njit из библиотеки numba позволяет компилировать циклы Python в машинный код LLVM, ускоряя расчет сложных финансовых индикаторов (таких как ATR, RSI или кастомные скользящие средние) до 100 раз по сравнению с чистым Pandas.

Содержание

Исходный код

В количественном анализе и алгоритмической торговле скорость расчетов имеет решающее значение. При тестировании стратегий (backtesting) на миллионах баров стандартные средства Pandas, такие как .rolling().apply(), создают огромные накладные расходы. Особенно это критично для рекуррентных индикаторов (например, Wilder’s MA, EMA, ATR), где значение текущего бара зависит от предыдущего, что делает невозможной простую векторизацию в NumPy.

Для решения этой проблемы мы используем Numba — JIT-компилятор (Just-In-Time), который транслирует аннотированный код Python в оптимизированные машинные инструкции LLVM. Это позволяет выполнять циклы со скоростью языков C/C++.

Если вам нужно более сложное динамическое сглаживание ценовых рядов без запаздывания, обратите внимание на Фильтр Калмана на Python: Сглаживание ценовых рядов с pykalman. Быстрые циклы Numba также критически важны при симуляциях портфеля, например, когда выполняется Расчет Value at Risk (VaR) через Монте-Карло и GARCH на Python. А для работы с терабайтами тиковых данных лучше использовать распределенные системы, как описано в статье Как вычислить профиль объема (VPVR) на 50 ГБ данных с Dask.

Ниже представлен готовый к продакшену код для расчета индикатора Average True Range (ATR) с использованием Numba JIT и сравнением производительности с классической реализацией на Pandas.

import numpy as np
import numba as nb
import pandas as pd
import time

@nb.njit(fastmath=True, cache=True)
def calculate_atr_numba(high, low, close, period):
    n = len(close)
    tr = np.empty(n, dtype=np.float64)
    atr = np.empty(n, dtype=np.float64)
    
    # Первый элемент True Range
    tr[0] = high[0] - low[0]
    
    # Расчет True Range через цикл
    for i in range(1, n):
        hl = high[i] - low[i]
        hc = abs(high[i] - close[i-1])
        lc = abs(low[i] - close[i-1])
        tr[i] = max(hl, hc, lc)
        
    # Расчет первого значения ATR (простое среднее за период)
    sum_tr = 0.0
    for i in range(period):
        sum_tr += tr[i]
    atr[period-1] = sum_tr / period
    
    # Заполнение начальных значений NaN
    for i in range(period-1):
        atr[i] = np.nan
        
    # Расчет Wilders Moving Average для ATR
    alpha = 1.0 / period
    for i in range(period, n):
        atr[i] = tr[i] * alpha + atr[i-1] * (1.0 - alpha)
        
    return atr

# Функция сравнения производительности
def benchmark():
    # Генерируем синтетические данные (1 000 000 баров)
    np.random.seed(42)
    n = 1000000
    close = 100.0 + np.cumsum(np.random.randn(n) * 0.1)
    high = close + np.random.rand(n) * 0.5
    low = close - np.random.rand(n) * 0.5
    period = 14

    # 1. Запуск Numba (первый запуск включает компиляцию)
    calculate_atr_numba(high, low, close, period)
    
    # Замер времени Numba
    start_time = time.time()
    atr_numba = calculate_atr_numba(high, low, close, period)
    numba_duration = time.time() - start_time
    print(f'Время выполнения Numba: {numba_duration:.6f} сек')

    # 2. Запуск Pandas (классический подход)
    df = pd.DataFrame({'high': high, 'low': low, 'close': close})
    
    start_time = time.time()
    df['h-l'] = df['high'] - df['low']
    df['h-pc'] = (df['high'] - df['close'].shift(1)).abs()
    df['l-pc'] = (df['low'] - df['close'].shift(1)).abs()
    df['tr'] = df[['h-l', 'h-pc', 'l-pc']].max(axis=1)
    # Реализация Wilders MA в Pandas через ewm
    atr_pandas = df['tr'].ewm(alpha=1/period, adjust=False).mean().values
    pandas_duration = time.time() - start_time
    print(f'Время выполнения Pandas: {pandas_duration:.6f} сек')
    
    print(f'Ускорение в {pandas_duration / numba_duration:.2f} раз!')

if __name__ == '__main__':
    benchmark()

Разбор параметров

high: Одномерный массив NumPy (np.ndarray) типа float64, содержащий максимальные цены баров.
low: Одномерный массив NumPy (np.ndarray) типа float64, содержащий минимальные цены баров.
close: Одномерный массив NumPy (np.ndarray) типа float64, содержащий цены закрытия баров.
period: Целое число (int), задающее период сглаживания для индикатора ATR.
@nb.njit(fastmath=True, cache=True): Декоратор Numba. Режим nopython=True (сокращенно njit) гарантирует компиляцию без участия интерпретатора Python. Флаг fastmath разрешает оптимизацию операций с плавающей точкой, а cache=True сохраняет скомпилированный код на диск для ускорения последующих запусков.

Как запустить

1. Установите необходимые библиотеки в вашу виртуальную среду:

pip install numba numpy pandas

2. Скопируйте приведенный выше исходный код в файл numba_atr.py.

3. Запустите скрипт через терминал: python numba_atr.py. Вы увидите сравнение времени выполнения расчета ATR с помощью Numba и стандартного подхода на базе Pandas. Разница в скорости на больших датасетах обычно составляет от 50 до 150 раз в пользу Numba.