python - 将自定义 numba njit 函数应用于 pandas 滚动对象

Question

与 pandas 1.0.0 版本，.apply现在有ability使用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模式。

Compilation is falling back to object mode WITH looplifting enabled because Function "numba_mean" failed type inference due to: non-precise type pyobject

应用时的错误回溯numba_mean与 njit而不是jit

This error may have been caused by the following argument(s): - argument 0: cannot determine Numba type of <class 'pandas.core.series.Series'>

我想知道是否有办法将 numbafied 函数与 pandas 一起使用rolling对象？

Answer 1

您需要指定 engine keyword让 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)})