如何沿矩阵轴进行滚动求和？

Question

给定一个矩阵 X，它有 T 行和 k 列：

T = 50
H = 10
k = 5 
X = np.arange(T).reshape(T,1)*np.ones((T,k))

如何在行的方向上，对 X 进行一个带有延迟 H 的滚动累加呢？

Xcum = np.zeros((T-H,k))
for t in range(H,T):
    Xcum[t-H,:] = np.sum( X[t-H:t,:], axis=0 )

注意，最好避免使用步幅和卷积，遵循广播/向量化的最佳实践。

Answer 1

这里有一个步长的解决方案。我知道这不是你想要的，但我想看看它的表现如何。

def foo2(X):
    temp = np.lib.stride_tricks.as_strided(X, shape=(H,T-H+1,k), 
        strides=(k*8,)+X.strides))
    # return temp.sum(0)
    return np.einsum('ijk->jk', temp)

这个方法的运行时间是35微秒，而Jaime的cumsum方法是22微秒。einsum比sum(0)稍微快一点。temp使用了X的数据，所以没有额外的内存开销。不过，这个方法比较难理解。

Answer 2

其实你在计算滚动总和的时候，最后一行数据是缺失的，正确的结果应该是这样的：

Xcum = np.zeros((T-H+1, k))
for t in range(H, T+1):
    Xcum[t-H, :] = np.sum(X[t-H:t, :], axis=0)

如果你想在任意的轴上使用numpy来实现这个功能，最简单的方法就是沿着那个轴使用 np.cumsum，然后通过计算两个切片的差值来得到结果。根据你的示例数组和轴：

temp = np.cumsum(X, axis=0)
Xcum = np.empty((T-H+1, k))
Xcum[0] = temp[H-1]
Xcum[1:] = temp[H:] - temp[:-H]

另外一个选择是使用pandas，它的 rolling_sum 函数，虽然看起来不太可能，但它确实可以在二维数组上正常工作，正好符合你的需求：

import pandas as pd
Xcum = pd.rolling_sum(X, 10)[9:] # first 9 entries are NaN

Answer 3

听起来你想要的是下面这个：

import scipy.signal
scipy.signal.convolve2d(X, np.ones((H,1)), mode='valid')

这当然是用到了卷积（convolve），不过，问题本身就是在讨论卷积操作。如果用广播（broadcasting），会导致算法变得更慢，而且占用更多内存。