Panda的BaseIndexer非常方便,虽然需要一些琢磨才能正确使用。
接下来,我使用np.searchsorted快速查找每个窗口的索引(开始、结束):
from pandas.api.indexers import BaseIndexer
class RangeWindow(BaseIndexer):
def __init__(self, val, width):
self.val = val.values
self.width = width
def get_window_bounds(self, num_values, min_periods, center, closed):
if min_periods is None: min_periods = 0
if closed is None: closed = 'left'
w = (-self.width/2, self.width/2) if center else (0, self.width)
side0 = 'left' if closed in ['left', 'both'] else 'right'
side1 = 'right' if closed in ['right', 'both'] else 'left'
ix0 = np.searchsorted(self.val, self.val + w[0], side=side0)
ix1 = np.searchsorted(self.val, self.val + w[1], side=side1)
ix1 = np.maximum(ix1, ix0 + min_periods)
return ix0, ix1
一些豪华选项:min_periods
、center
和 closed
是根据 DataFrame.rolling 指定实现的。
应用:
df = pd.DataFrame([
[4.5, 10],
[4.6, 11],
[4.8, 9],
[5.5, 6],
[5.6, 6],
[8.1, 10],
[8.2, 13]
], columns='a b'.split())
df.b.rolling(RangeWindow(df.a, width=1.0), center=True, closed='both').mean()
0 10.0
1 10.0
2 10.0
3 6.0
4 6.0
5 11.5
6 11.5
Name: b, dtype: float64
时间:
df = pd.DataFrame(
np.random.uniform(0, 1000, size=(1_000_000, 2)),
columns='a b'.split(),
)
df = df.sort_values('a').reset_index(drop=True)
%%time
avg = df.b.rolling(RangeWindow(df.a, width=1.0)).mean()
CPU times: user 133 ms, sys: 3.58 ms, total: 136 ms
Wall time: 135 ms
性能更新:
在@anon01的评论后,我想知道当滚动窗口较大时是否可以更快地进行滚动。结果我应该先测量Pandas的滚动均值和求和性能...(过早优化,有人吗?)请看最后为什么。
无论如何,想法是只需做一次cumsum
,然后取由窗口端点解除引用的元素的差异:
def fast_rolling_sum(a, b, width):
z = np.concatenate(([0], np.cumsum(b)))
ix0 = np.searchsorted(a, a - width/2, side='left')
ix1 = np.searchsorted(a, a + width/2, side='right')
return z[ix1] - z[ix0]
def fast_rolling_mean(a, b, width):
z = np.concatenate(([0], np.cumsum(b)))
ix0 = np.searchsorted(a, a - width/2, side='left')
ix1 = np.searchsorted(a, a + width/2, side='right')
return (z[ix1] - z[ix0]) / (ix1 - ix0)
有了这个(以及上面的100万行df
),我看到:
%timeit fast_rolling_mean(df.a.values, df.b.values, width=100.0)
对比:
%timeit df.rolling(RangeWindow(df.a, width=100.0), min_periods=1).mean()
然而!!! Pandas很可能已经在做这样的优化(这是一个非常明显的优化)。随着窗口大小的增加,时间并不会增加(这就是为什么我说我应该先检查一下)。