如何在numpy中沿轴应用操作

How to apply an operation along an axis in numpy

我有一个 N 维数组,我想 return 另一个相同形状的数组,其中的值沿给定轴逐渐累积。我的尝试在下面,但它抛出了一个错误,因为 put_along_axis 并没有完全按照我的想法去做。如何正确书写?

def accumulate(data, axis):
    rez = np.zeros(data.shape)
    for i in range(1, data.shape[axis]):
        tmp1 = np.take(rez, i-1, axis=axis)
        tmp2 = np.take(rez, i, axis=axis)
        tmpsum = tmp1 + tmp2
        np.put_along_axis(rez, i, tmpsum, axis=axis)
    return rez

听起来你只是在寻找 np.cumsum:

data.cumsum(axis=axis)

会return就是这样。