Multidimensional cumulative sum in numpy

Question

I want to be able to calculate the cumulative sum of a large n-dimensional numpy array. The value of each element in the final array should be the sum of all elements which have indices greater than or equal to the current element.

2D: xᶦʲ = ∑xᵐⁿ ∀ m ≥ i and n ≥ j

3D: xᶦʲᵏ = ∑xᵐⁿᵒ ∀ m ≥ i and n ≥ j and o ≥ k

Examples in 2D:

1 1 0       2  1  0
1 1 1  ->   5  3  1
1 1 1       8  5  2

1 2 3       6  5  3
4 5 6  ->  21 16  9
7 8 9      45 33 18

Example in 3D:

1 1 1       3   2   1
1 1 1       6   4   2
1 1 1       9   6   3

1 1 1       6   4   2
1 1 1  ->  12   8   4
1 1 1      18  12   6

1 1 1       9   6   3
1 1 1      18  12   6
1 1 1      27  18   9

Divakar · Accepted Answer · 2018-05-22 09:12:40Z

5

Flip along the last axis, cumsum along the same, flip it back and finally cumsum along the second last axis onwards until the first axis -

def multidim_cumsum(a):
    out = a[...,::-1].cumsum(-1)[...,::-1]
    for i in range(2,a.ndim+1):
        np.cumsum(out, axis=-i, out=out)
    return out

Sample 2D case run -

In [107]: a
Out[107]: 
array([[1, 1, 0],
       [1, 1, 1],
       [1, 1, 1]])

In [108]: multidim_cumsum(a)
Out[108]: 
array([[2, 1, 0],
       [5, 3, 1],
       [8, 5, 2]])

Sample 3D case run -

In [110]: a
Out[110]: 
array([[[1, 1, 1],
        [1, 1, 1],
        [1, 1, 1]],

       [[1, 1, 1],
        [1, 1, 1],
        [1, 1, 1]],

       [[1, 1, 1],
        [1, 1, 1],
        [1, 1, 1]]])

In [111]: multidim_cumsum(a)
Out[111]: 
array([[[ 3,  2,  1],
        [ 6,  4,  2],
        [ 9,  6,  3]],

       [[ 6,  4,  2],
        [12,  8,  4],
        [18, 12,  6]],

       [[ 9,  6,  3],
        [18, 12,  6],
        [27, 18,  9]]])

edited May 22, 2018 at 9:12

answered May 22, 2018 at 8:47

Divakar

222k19 gold badges273 silver badges374 bronze badges

Sign up to request clarification or add additional context in comments.

3 Comments

user2685230 Over a year ago

Of course, thanks! For completeness the solution for my use case in 4D is: a[...,::-1].cumsum(-1)[...,::-1].cumsum(-2).cumsum(-3).cumsum(-4)

Divakar Over a year ago

@user2685230 Edited with a generic solution based on the same idea.

jtlz2 Over a year ago

This is flipping amazing, and should really be called a rubik's cube :D

ch271828n · Accepted Answer · 2020-05-13 01:04:47Z

5

For those who want a "numpy-like" cumsum where the top-left corner is smallest:

def multidim_cumsum(a):
    out = a.cumsum(-1)
    for i in range(2,a.ndim+1):
        np.cumsum(out, axis=-i, out=out)
    return out

Modified from @Divakar (thanks to him!)

edited May 13, 2020 at 1:04

answered Mar 12, 2020 at 3:23

ch271828n

17.9k5 gold badges71 silver badges108 bronze badges

1 Comment

jtlz2 Over a year ago

Perfect, invaluable!

Paul Panzer · Accepted Answer · 2018-05-22 09:16:17Z

2

Here is a general solution. I'm going by the description, not the examples, i.e. order of vertical display is top down not bottom up:

import itertools as it
import functools as ft

ft.reduce(np.cumsum, it.chain((a[a.ndim*(np.s_[::-1],)],), range(a.ndim)))[a.ndim*(np.s_[::-1],)]

Or in-place:

for i in range(a.ndim):
    b = a.swapaxes(0, i)[::-1]
    b.cumsum(axis=0, out=b)

edited May 22, 2018 at 9:16

answered May 22, 2018 at 9:04

Paul Panzer

53.3k3 gold badges60 silver badges103 bronze badges

3 Comments

shaunc Over a year ago

Seems your in-place version has a bug? b is overwritten on subsequent iterations.

Paul Panzer Over a year ago

@shaunc "[...] bug? bis overwritten [...]" ??? Isn't that the meaning of in-place?

shaunc Over a year ago

I mean -- suppose a.ndim > 1. On each iteration, b doesn't depend on the previously calculated value of b, but it gets a new copy from a, ignoring the previous work.

Collectives™ on Stack Overflow

Multidimensional cumulative sum in numpy

3 Answers 3

3 Comments

1 Comment

3 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

3 Comments

1 Comment

3 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related