Я вижу некоторую необычную производительность с использованием Numba, а также хочу оптимизировать цикл JIT.
Запустите и сгенерируйте некоторые практически важные данные:
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from time import time
import numba
times = np.arange(datetime(2000, 1, 1), datetime(2020, 2, 1), timedelta(minutes=10)).astype(np.datetime64)
tlen = len(times)
A, Z = np.array(['A', 'Z']).view('int32')
symbol_names = np.random.randint(low=A, high=Z, size=1 * 7, dtype='int32').view(f'U{7}')
times = np.concatenate([times] * 1)
names = np.array([y for x in [[s] * tlen for s in symbol_names] for y in x])
open_column = np.random.randint(low=40, high=60, size=len(times), dtype='uint32')
high_column = np.random.randint(low=50, high=70, size=len(times), dtype='uint32')
low_column = np.random.randint(low=30, high=50, size=len(times), dtype='uint32')
close_column = np.random.randint(low=40, high=60, size=len(times), dtype='uint32')
df = pd.DataFrame({'open': open_column, 'high': high_column, 'low': low_column, 'close': close_column}, index=[names, times])
df.index = df.index.set_names(['Symbol', 'Date'])
df['entry'] = np.select( [df.open > df.open.shift(), False], (df.close, -1), np.nan)
df['exit'] = df.close.where(df.high > df.open*1.33, np.nan)
Функция синхронизации:
def timing(f):
def wrap(*args):
time1 = time()
ret = f(*args)
time2 = time()
print('{:s} function took {:.3f} s'.format(f.__name__, (time2-time1)))
return ret
return wrap
JIT-скомпилированная функция:
@numba.jit(nopython=True)
def entry_exit(arr, limit=0, stop=0, tbe=0):
is_active = 0
bars_held = 0
limit_target = np.inf
stop_target = -np.inf
result = np.empty(arr.shape[0], dtype='float32')
for n in range(arr.shape[0]):
ret = 0
if is_active == 1:
bars_held += 1
if arr[n][2] < stop_target:
ret = stop_target
is_active = 0
elif arr[n][1] > limit_target:
ret = limit_target
is_active = 0
elif bars_held >= tbe:
ret = arr[n][3]
is_active = 0
elif arr[n][5] > 0:
ret = arr[n][3]
is_active = 0
if is_active == 0:
if arr[n][4] > 0:
is_active = 1
bars_held = 0
if stop != 0:
stop_target = arr[n][3] * stop
if limit != 0:
limit_target = arr[n][3] * limit
result[n] = ret
return result
Тесты:
@timing
def run_one(arr):
entry_exit(arr, limit=1.20, stop=0.50, tbe=5)
@timing
def run_ten(arr):
for _ in range(10):
entry_exit(arr, limit=1.20, stop=0.50, tbe=5)
arr = df[['open', 'high', 'low', 'close', 'entry', 'exit']].to_numpy()
run_one(arr)
run_ten(arr)
При запуске этого в родном Python я получаю:
Имеет смысл.
Когда я запускаю то же самое в JIT, я получаю совершенно другое:
Почему это происходит? Мне также очень интересно узнать, как еще больше ускорить функцию, поскольку текущий прирост скорости, хотя и значительный, недостаточен.
numba.jit
скомпилирует функцию при первом использовании. Это делает первое выполнение функции дорогим, а последующие намного дешевле.
Ваш тест предположительно работает run_one
, который вызывает entry_exit
, который компилируется numba, поэтому компилируется медленно, но быстро запускается. Затем он вызывает run_ten
, но entry_exit
уже скомпилирован, поэтому повторно используется скомпилированная форма — так это быстро.
Таким образом, я ожидаю, что разбивка будет примерно такой
run_one: 0.74s compile + 1 x 0.01s run
run_ten: no compile + 10 x 0.01s run
Чтобы проверить это, вам просто нужно убедиться, что вы вызываете функцию один раз (чтобы numba скомпилировала ее), прежде чем вы начнете тестировать ее скорость. Или вы можете установить флаги, чтобы сообщить numba о компиляции заранее.
Все, что вам нужно сделать, чтобы убедиться в этом, это изменить тестовый сценарий следующим образом:
@timing
def run_one(arr):
entry_exit(arr, limit=1.20, stop=0.50, tbe=5)
@timing
def run_ten(arr):
for _ in range(10):
entry_exit(arr, limit=1.20, stop=0.50, tbe=5)
arr = df[['open', 'high', 'low', 'close', 'entry', 'exit']].to_numpy()
# Run it once so that numba compiles it
entry_exit(arr, limit=1.20, stop=0.50, tbe=5)
# Use the compiled version
run_one(arr)
run_ten(arr)
Любые предложения о том, как написать это лучше, кроме использования предварительной компиляции?
Если моя гипотеза верна, значит, вы все делаете правильно. Вы просто не учитываете время компиляции. Если вам не нужны накладные расходы на компиляцию один раз за запуск, то numba может быть не для вас. Тем не менее, похоже, что у numba есть некоторые дополнительные параметры постоянного кеша, но я их не использовал. numba.pydata.org/numba-doc/latest/developer/caching.html
Вы можете использовать
cache=True
. Это кэширует скомпилированную функцию на вашем жестком диске. Но инициализация кэша также может занять некоторое время, но при наличии более одной кэшированной функции это будет намного быстрее, чем компиляция функции каждый раз при перезапуске интерпретатора.