如何加速 pandas.rolling 加权平均或其他功能？

Question

我想计算一个长数组的加权平均值。

在固定大小的窗口中滚动。

例如：

In [1]: import pandas as pd

In [2]: a = pd.Series(range(int(1e8)))

In [5]: import numpy as np; w = np.array(list(range(10)));

In [6]: a.rolling(10).apply(lambda x: (x * w).sum())

我试过了，这很慢，我读了一些博客，有时可以通过以下方式加快速度：

 a.rolling(10).apply(np.argmax, engine='numba', raw=True)

但这只能在内置函数中使用，对于某些自定义函数，似乎不起作用。

你知道如何让它在花费时间方面经得起考验吗？

Answer 1

np.convolve 的解决方案

我们可以使用valid模式卷积运算对序列a上的权重w进行卷积，本质上这与计算滚动加权和的效果相同。

s = np.convolve(a, w[::-1], 'valid')
s = [np.nan] * (len(w) - 1) + list(s)

解决方案sliding_window_view

或者，我们也可以使用sliding_window_view来加快滚动加权和计算

from numpy.lib.stride_tricks import sliding_window_view

s = (sliding_window_view(a, len(w)) * w).sum(1)
s = [np.nan] * (len(w) - 1) + list(s)

时间

a = pd.Series(range(int(1e4)))

%%timeit
s = np.convolve(a, w[::-1], 'valid')
s = [np.nan] * (len(w) - 1) + list(s)
# 1000 loops, best of 5: 626 µs per loop

%%timeit
s = (sliding_window_view(a, len(w)) * w).sum(1)
s = [np.nan] * (len(w) - 1) + list(s)    
# 1000 loops, best of 5: 1.2 ms per loop

%%timeit
s = a.rolling(10).apply(lambda x: (x * w).sum())
# 1 loop, best of 5: 3.6 s per loop

从性能测试中可以看出，使用 np.convolve 的速度大约是 5750x，而使用 sliding_window_view 的速度大约是 2880x熊猫rolling + apply 方法。