加快滑动窗口平均计算

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)

Answer 1

这基本上是滑动窗口平均计算。这种平均可以被认为是在滑动窗口中求和并除以窗口大小的长度。所以，我们可以使用1D convolution 和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)