如何对2D Numpy数组按轴计算分段累计和，信号变化时自动重置？

需求说明 #

对二维numpy数组的指定轴（默认行方向，每一行对应一条信号时间序列）计算累计和，当连续的信号值发生变化时，累计和重新从当前值开始计数。

输入输出示例 #

输入示例 #

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]
])

期望输出 #

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]
])

实现代码 #

def group_cumsum2d(s, axis=1):
    s = np.asarray(s).copy()
    # 沿指定轴识别数值变化的位置
    diff_mask = np.diff(s, axis=axis) != 0
    # 构造和原数组同维度的掩码，标记数值变化后的第一个元素位置
    pad_shape = list(s.shape)
    pad_shape[axis] = 1
    pos_mask = np.concatenate([np.zeros(pad_shape, dtype=bool), diff_mask], axis=axis)
    # 计算全局累计和
    global_cumsum = np.cumsum(s, axis=axis)
    # 计算每个分段需要扣除的偏移量
    offset = np.cumsum(pos_mask * global_cumsum, axis=axis)
    # 应用偏移量得到分段累计和结果
    return global_cumsum - offset

功能说明 #

• 支持任意维度的numpy数组输入，可通过axis参数指定累计和的计算方向
• 全程使用numpy向量化运算，无显式循环，计算效率高

测试验证 #

cum_sum = group_cumsum2d(signal)
print(cum_sum)

运行后输出和期望结果完全一致。

内容的提问来源于stack exchange，提问作者J. Daay