加速滑动窗口平均值计算

Question

我有一些数据（股票数据）需要进行一些计算。我使用了numpy数组来完成这个操作。Numpy比Python内置函数要快得多。但是，我的代码执行时间比预期的要长。我的代码如下所示，并且我已经使用了ipython %timeit函数来测试它。结果如下：总执行时间为5.44毫秒，第二个‘for’循环花费了最多的时间，3.88毫秒，原因是该循环中的‘np.mean’函数。因此，寻找替代‘np.mean’函数以及其他加速执行时间的建议将会很有帮助。 代码

data = my_class.Project.all_data["AAP_data"]
data = np.array(data[["High", "Low", "Close"]])

true_range = np.empty((data.shape[0]-1, 1))
for i in range(1, true_range.shape[0]+1):
    true_range[i-1] = max((data[i, 0] - data[i, 1]), (abs(data[i, 0] - data[i-1, 2])),
                          (abs(data[i, 1] - data[i-1, 2])))

average_true_range = np.empty((true_range.shape[0]-13, 1))
for i in range(13, average_true_range.shape[0]+13):
    lastn_tr = true_range[(i-13):(i+1)]
    average_true_range[i-13] = np.mean(lastn_tr)

- Elgin Cahangirov

Answer 1

这基本上是滑动窗口平均计算。这种平均可以被认为是在滑动窗口中求和，然后除以窗口大小。因此，我们可以使用1D卷积和np.convolve的向量化解决方案来摆脱整个循环过程，从而得到average_true_range，如下所示 -

np.convolve(true_range,np.ones((14),dtype=int),'valid')/14.0

为进一步提高性能，我们可以从研究CPU执行乘法比除法更有效率中得到启示。因此，在这里采用这种方式以获得改进版本 -

r = 1.0/14
out = np.convolve(true_range,np.ones((14),dtype=int),'valid')*r

运行时测试 -

In [53]: def original_app(true_range):
    ...:     average_true_range = np.zeros((true_range.shape[0]-13, 1))
    ...:     for i in range(13, average_true_range.shape[0]+13):
    ...:         lastn_tr = true_range[(i-13):(i+1)]
    ...:         average_true_range[i-13] = np.mean(lastn_tr)
    ...:     return average_true_range
    ...: 
    ...: def vectorized_app(true_range):
    ...:     return np.convolve(true_range,np.ones((14),dtype=int),'valid')/14.0
    ...: 
    ...: def vectorized_app2(true_range):
    ...:     r = 1.0/14
    ...:     return np.convolve(true_range,np.ones((14),dtype=int),'valid')*r
    ...: 

In [54]: true_range = np.random.rand(10000) # Input array

In [55]: %timeit original_app(true_range)
1 loops, best of 3: 180 ms per loop

In [56]: %timeit vectorized_app(true_range)
1000 loops, best of 3: 446 µs per loop

In [57]: %timeit vectorized_app2(true_range)
1000 loops, best of 3: 401 µs per loop

速度大大提升了！

---

接下来，瓶颈可能会转移到获取 true_range 的第一部分。为了向量化这些东西，这里使用了切片方法 -

col0 = data[1:,0] - data[1:,1]
col1 = np.abs(data[1:,0] - data[:-1,2])
col2 = np.abs(data[1:,1] - data[:-1,2])
true_range = np.maximum(np.maximum(col0,col1),col2)

- Divakar

---

谢谢您的评论！我认为我应该阅读更多的numpy文档。 - Elgin Cahangirov

@ElginCahangirov NumPy文档就大部分东西而言都非常全面，所以这是个好主意！ - Divakar