Pandas使用滑动窗口将函数应用于多个列



我需要在数据框上使用滑动窗口来计算一些度量。如果metric只需要1列,我就用rolling。但有些如何它不工作与2+列。下面是我如何使用常规周期计算度量。

def mean_squared_error(aa, bb):
return np.sum((aa - bb) ** 2) / len(aa)
def rolling_metric(df_, col_a, col_b, window, metric_fn):
result = []
for i, id_ in enumerate(df_.index):
if i < (df_.shape[0] - window + 1):
slice_idx = df_.index[i: i+window-1]
slice_a, slice_b = df_.loc[slice_idx, col_a], df_.loc[slice_idx, col_b]
result.append(metric_fn(slice_a, slice_b))
else:
result.append(None)
return pd.Series(data = result, index = df_.index)
df = pd.DataFrame(data=(np.random.rand(1000, 2)*10).round(2), columns = ['y_true', 'y_pred'] )
%time df2 = rolling_metric(df, 'y_true', 'y_pred', window=7, metric_fn=mean_squared_error)

对于1000行,这将花费接近1秒的时间。

请建议更快的矢量化方法来计算滑动窗口上的度量。

在本例中:

您可以事先计算平方误差,然后使用.Rolling.mean():

df['sq_error'] = (df['y_true'] - df['y_pred'])**2
%time df['sq_error'].rolling(6).mean().dropna()

请注意,在您的示例中,实际窗口大小为6(打印切片长度),这就是我在代码片段中将其设置为6的原因。

你甚至可以这样写:

%time df['y_true'].subtract(df['y_pred']).pow(2).rolling(6).mean().dropna()
一般

:

如果您不能将其减少到单个列,那么从pandas 1.3.0开始,您可以使用method='table'参数将函数应用于整个DataFrame。然而,这有以下要求:

  • 这仅在使用numba引擎时实现。因此,您需要在apply中设置engine='numba'并安装它。
  • 您需要在apply中设置raw=True:这意味着在您的函数中您将操作numpy数组而不是DataFrame。这是前一点的结果。

因此,您的计算可能是这样的:

WIN_LEN = 6
def mean_sq_err_table(arr, min_window=WIN_LEN):
if len(arr) < min_window:
return np.nan
else:
return np.mean((arr[:, 0] - arr[:, 1])**2)

df.rolling(WIN_LEN, method='table').apply(mean_sq_err_table, engine='numba', raw=True).dropna()

因为它使用numba,所以这也相对较快。

最新更新