Многомерная накопленная сумма в numy

Question

Я хочу иметь возможность рассчитать совокупную сумму большого n-мерного массива numpy.Значение каждого элемента в конечном массиве должно быть суммой всех элементов, индексы которых больше или равны текущему элементу.

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

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

Примеры в 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

Пример в 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

Answer 1

Вот общее решение.Я иду по описанию, а не по примерам, то есть порядок вертикального отображения сверху вниз, а не снизу вверх:

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

Или на месте:

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

Answer 2

Перевернуть по последней оси, собрать по этой же длине, перевернуть обратно и, наконец, переместить по второй по последней оси вперед до первой оси -

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

Образец двухмерного прогона случая -

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

Пример запуска в 3D-случае -

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