применить пользовательскую функцию numba njit к pandas подвижному объекту

Question

С выпуском pandas 1.0.0 .apply теперь обладает способностью использовать numba jit функций.

Есть ли способ воспользоваться этим .apply функциональность при использовании .rolling?

Например, у меня есть функция pandas.DataFrame и numba_mean, которую я хотел бы применить, чтобы получить скользящее среднее из 3 периодов.

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

df = pd.DataFrame({"A" : np.random.rand(10)})
@nb.jit 
def numba_mean(x): 
    return np.sum(x) / len(x) 

df.A.rolling(3).apply(numba_mean)

Однако я получаю приведенную ниже ошибку, которая означает, что она не работает в режиме nopython.

Компиляция возвращается к объектному режиму при включенном циклическом переключении, поскольку функция "numba_mean" не смогла сделать вывод типа из-за : неточный тип pyobject

Отслеживание ошибок при применении numba_mean с njit вместо jit

Эта ошибка могла быть вызвана следующими причинами аргумент (ы): - аргумент 0: невозможно определить тип Numba <class 'pandas.core.series.Series'>

Мне было интересно, есть ли способ использовать функции numbafied с pandas rolling объект?

Answer 1

Вам нужно указать ключевое слово engine , чтобы Pandas знал, что вы хотите использовать Numba:

df.A.rolling(3).apply(numba_mean, engine='numba', raw=True)

Pandas может jit функция для вас, но Я получаю более быстрые результаты, когда я делаю это сам. Возможно, Numba перекомпилирует для каждого вызова .apply(), а нижеприведенное время включает время компиляции.

При такой простой функции и небольшом количестве данных для каждого вызова (window = 3) маловероятно, что вы Вы получите ускорение по встроенной функции.

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

@nb.njit(nogil=True)
def numba_mean(x):
    return np.sum(x) / x.size

def numpy_mean(x): 
    return np.sum(x) / x.size

df = pd.DataFrame({"A" : np.random.rand(10000)})