How to apply cumulative sum over an axis in a 2d Numpy array in slices with condition to previous va-CodePudding

For a 2d numpy array, I would like to have an array with the cumulative sum for every row (signal time series). But the sum needs to restart every time the signal changes. Below, you can find an example of the input and output of the function.

    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 inputs

def group_cumsum2d(s, axis=1):

For a 1d array, the following function will work:

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)

The 1d function was found in this question: How to apply cummulative sum in Numpy in slices with condition to previous value?

CodePudding user response：

Just apply your solution along an axis with 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)

gives True. I.e. the result matches cum_sum in all entries.