大熊猫滚动最大值的numpy版本

Question

TL;DR：我的问题是如何改进我的功能以超越 pandas 自身的最大移动功能？

---

背景信息：

所以我正在处理大量移动平均线、移动最大值和移动最小值等，到目前为止，我发现的唯一移动窗口之类的功能位于 pandas.rolling method。问题是：我拥有的数据是 numpy 数组，我想要的最终结果也必须在 numpy 数组中；尽管我想简单地将其转换为 pandas 系列并返回 numpy 数组来完成这样的工作：

result2_max = pd.Series(data_array).rolling(window).max().to_numpy()

，转换数据类型似乎没有必要，而且在 numpy 实现中可能有一些方法可以做完全相同的事情。

然而，尽管它看起来很不合 Python，但它比我想出的或在网上看到的任何方法都快。我将在下面给出小基准：

import numpy as np
import pandas as pd

def numpy_rolling_max(data, window):

    data = data[::-1]
    data_strides = data.strides[0]

    movin_window = np.lib.stride_tricks.as_strided(data, 
                                                    shape=(data.shape[0] - window +1, window), 
                                                    strides = (data_strides ,data_strides)
                                                    )[::-1]
    max_window =np.amax(movin_window, axis = 1)#this line seems to be the bottleneck


    nan_array = np.full(window - 1, np.nan)
    return np.hstack((nan_array, max_window))


def pandas_rolling_max(data, window):
    return pd.Series(data).rolling(window).max().to_numpy()

length = 120000
window = 190
data = np.arange(length) + 0.5

result1_max = numpy_rolling_max(data, window)#21.9ms per loop
result2_max = pandas_rolling_max(data, window)#5.43ms per loop

result_comparision = np.allclose(result1_max, result2_max, equal_nan = True)

在 arraysize = 120k，window = 190 的情况下，pandas 滚动最大值比 numpy 版本快大约 3 倍。我不知道从哪里开始，因为我已经尽可能地矢量化了我自己的函数，但它仍然比 pandas 版本慢得多，我真的不知道为什么。

提前谢谢你

编辑：我找到了瓶颈，就是这一行：

max_window =np.amax(movin_window, axis = 1)

但是看到已经是向量化的函数调用，我还是不知道怎么继续。

Answer 1

我们可以使用1D max filter from Scipy 来复制与pandas 相同的行为，并且仍然更高效 -

from scipy.ndimage.filters import maximum_filter1d

def max_filter1d_same(a, W, fillna=np.nan):
    out_dtype = np.full(0,fillna).dtype
    hW = (W-1)//2 # Half window size
    out = maximum_filter1d(a,size=W, origin=hW)
    if out.dtype is out_dtype:
        out[:W-1] = fillna
    else:
        out = np.concatenate((np.full(W-1,fillna), out[W-1:]))
    return out

示例运行 -

In [161]: np.random.seed(0)
     ...: a = np.random.randint(0,999,(20))
     ...: window = 3

In [162]: a
Out[162]: 
array([684, 559, 629, 192, 835, 763, 707, 359,   9, 723, 277, 754, 804,
       599,  70, 472, 600, 396, 314, 705])

In [163]: pd.Series(a).rolling(window).max().to_numpy()
Out[163]: 
array([ nan,  nan, 684., 629., 835., 835., 835., 763., 707., 723., 723.,
       754., 804., 804., 804., 599., 600., 600., 600., 705.])

In [164]: max_filter1d_same(a,window)
Out[164]: 
array([ nan,  nan, 684., 629., 835., 835., 835., 763., 707., 723., 723.,
       754., 804., 804., 804., 599., 600., 600., 600., 705.])

# Use same dtype fillna for better memory efficiency
In [165]: max_filter1d_same(a,window,fillna=0)
Out[165]: 
array([  0,   0, 684, 629, 835, 835, 835, 763, 707, 723, 723, 754, 804,
       804, 804, 599, 600, 600, 600, 705])

实际测试用例大小的时间安排 -

In [171]: # Actual test-cases sizes
     ...: np.random.seed(0)
     ...: data_array = np.random.randint(0,999,(120000))
     ...: window = 190

In [172]: %timeit pd.Series(data_array).rolling(window).max().to_numpy()
100 loops, best of 3: 4.43 ms per loop

In [173]: %timeit max_filter1d_same(data_array,window)
100 loops, best of 3: 1.95 ms per loop