CODEWITHSUNDEEP
pythonnumpyaxis
Aleksejs Fomins
Aleksejs Fomins

Reputation: 900

How to apply an operation along an axis in numpy

I have an N-dimensional array, and I want to return another array of the same shape, where the values are progressively accumulated along a given axis. My attempt is below, but it throws an error because put_along_axis does not do exactly what I think it does. How to correctly write this?

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

Upvotes: 0

Views: 102

Answers (1)

Darina
Darina

Reputation: 1632

It sounds like you are just looking for np.cumsum:

data.cumsum(axis=axis)

will return just that.

Upvotes: 1

Related Questions