How to vectorise triple for looped cumulative sum

I want to vectorise the triple sum

\sum_{i=1}^I\sum_{j=1}^J\sum_{m=1}^J a_{ijm}

such that I end up with a matrix

A \in \mathbb{R}^{I \times J}

where A_{kl} = \sum_{i=1}^k\sum_{j=1}^l\sum_{m=1}^l a_{ijm} for k = 1,...,I and l = 1, ...,J

carrying forward the sums to avoid pointless recomputation.

I currently use this code:
np.cumsum(np.cumsum(np.cumsum(a, axis = 0), axis = 1), axis = 2).diagonal(axis1 = 1, axis2 = 2)
but it is inefficient as it does lots of extra work and extracts the correct matrix at the end with the diagonal method. I can't think of how to make this faster.

The main challenge here is to compute the inner two sums, i.e. the sum of the square slices of a matrix originating from the top left. The final sum is just a cumsum on top of that along the 0th axis.

Setup:

import numpy as np

I, J = 100, 100
arr = np.random.rand(I, J, J)

Your implementation:

%%timeit
out = np.cumsum(np.cumsum(np.cumsum(arr, axis = 0), axis = 1), axis = 2).diagonal(axis1 = 1, axis2 = 2)
# 10.9 ms ± 162 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Your implementation improved by taking the diagonal before cumsumming over the 0th axis:

%%timeit
out = arr.cumsum(axis=1).cumsum(axis=2).diagonal(axis1=1, axis2=2).cumsum(axis=0)
# 6.25 ms ± 34.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Finally, some tril/triu trickery:

%%timeit
out = np.cumsum(np.cumsum(np.tril(arr, k=-1).sum(axis=2) + np.triu(arr).sum(axis=1), axis=1), axis=0)
# 3.15 ms ± 71.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

which is already better, but admittedly still not ideal. I don't see a better way to compute the inner two sums noted above with pure numpy.

You can use Numba so to produce a very fast implementation. Here is the code:

import numba as nb
import numpy as np

@nb.njit('(float64[:,:,::1],)', parallel=True)
def compute(arr):
    ni, nj, nk = arr.shape
    assert nj == nk
    result = np.empty((ni, nj))
    # Parallel cumsum along the axis 1 and 2 + extraction of the diagonal
    for i in nb.prange(ni):
        tmp = np.zeros(nk)
        for j in range(nj):
            for k in range(nk):
                tmp[k] += arr[i, j, k]
            result[i, j] = np.sum(tmp[:j+1])
    # Cumsum along the axis 0
    for i in range(1, ni):
        for k in range(nk):
            result[i, k] += result[i-1, k]
    return result

result = compute(a)

Here are performance results on my 6-core i5-9600KF with a 100x100x100 float64 input array:

Initial code:      12.7 ms
Chryophylaxs v1:    7.1 ms
Chryophylaxs v2:    5.5 ms
Numba:              0.2 ms

This implementation is significantly faster than all others. It is about 64 times faster than the initial implementation. It is also actually optimal on my machine since it completely saturate the bandwidth of my RAM only for reading the input array (which is mandatory). Note that it is better not to use multiple threads for very small arrays.

Note that this code also use far less memory as it only need 8 * nk * num_threads bytes of temporary storage as opposed to 16 * ni * nj * nk bytes for the initial solution.