如何在条件为先前值的切片中的 2d Numpy 数组中的轴上应用累积和？

Question

对于 2d numpy 数组，我想要一个包含每一行（信号时间序列）的累积和的数组。 但每次信号变化时，求和都需要重新开始。 下面，您可以找到 function 的输入和 output 的示例。

    import numpy as np
    signal = np.array([
        [1, 1, 1, -1, -1, -1, -1], 
        [1, 1, 1, -1, -1, -1, 1], 
        [1, 1, 1, -1, -1, 1, 1]
    ])

Output

cum_sum = np.array([
        [1, 2, 3, -1, -2, -3, -4], 
        [1, 2, 3, -1, -2, -3, 1], 
        [1, 2, 3, -1, -2, 1, 2]
    ])

Function 输入

def group_cumsum2d(s, axis=1):

对于一维数组，以下 function 将起作用：

def group_cumsum(s):
    # make a copy and ensure np.array (in case list was given)
    s = np.array(s).copy()
    idx = np.nonzero(np.diff(s))[0]  # last of each group
    off = np.diff(np.concatenate(([0], np.cumsum(s)[idx])))
    s[idx + 1] -= off
    return np.cumsum(s)

在这个问题中找到了 1d function： How to apply cummulative sum in Numpy in slice with condition to previous value?

Answer 1

只需使用np.apply_along_axis沿轴应用您的解决方案。

def group_cumsum2d(s, axis=1):
    return np.apply_along_axis(group_cumsum, axis, s)
    
np.all(group_cumsum2d(signal) == cum_sum)

给出True 。 即结果与所有条目中的 cum_sum 匹配。