将自定义numba njit函数应用于熊猫滚动对象

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模式下不起作用。

由于函数“ numba_mean”由于以下原因导致类型推断失败，编译已回到对象模式并启用了循环提升：>

将numba_mean与njit而不是jit一起应用时的错误回溯>

此错误可能是由以下参数引起的：-参数0：无法确定<class 'pandas.core.series.Series'>

的Numba类型

我想知道是否有一种方法可以将numbafied

函数与pandas rolling对象一起使用？

随着pandas 1.0.0发行版，.apply现在可以使用numba jit函数。使用.rolling时是否有办法利用此.apply功能？例如我有一只熊猫....

Answer 1

您需要指定engine keyword以便让熊猫知道您要使用Numba：

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