Python - Sum 4D Array

Given a 4D array M: (m, n, r, r), how can I sum all the m * n inner matrices (of shape (r, r)) to get a new matrix of shape (r * r)?

For example,

    M [[[[ 4,  1],
         [ 2,  1]],

        [[ 8,  2],
         [ 4,  2]]],

       [[[ 8,  2],
         [ 4,  2]],

        [[ 12, 3],
         [ 6,  3]]]]

I expect the result should be

      [[32, 8],
       [16, 8]]

You could use einsum:

In [21]: np.einsum('ijkl->kl', M)
Out[21]:
array([[32,  8],
       [16,  8]])

Other options include reshaping the first two axes into one axis, and then calling sum:

In [24]: M.reshape(-1, 2, 2).sum(axis=0)
Out[24]:
array([[32,  8],
       [16,  8]])

or calling the sum method twice:

In [26]: M.sum(axis=0).sum(axis=0)
Out[26]:
array([[32,  8],
       [16,  8]])

But using np.einsum is faster:

In [22]: %timeit np.einsum('ijkl->kl', M)
100000 loops, best of 3: 2.42 µs per loop

In [25]: %timeit M.reshape(-1, 2, 2).sum(axis=0)
100000 loops, best of 3: 5.69 µs per loop

In [43]: %timeit np.sum(M, axis=(0,1))
100000 loops, best of 3: 6.08 µs per loop

In [33]: %timeit sum(sum(M))
100000 loops, best of 3: 8.18 µs per loop

In [27]: %timeit M.sum(axis=0).sum(axis=0)
100000 loops, best of 3: 9.83 µs per loop

Caveat: timeit benchmarks can vary significantly due to many factors (OS, NumPy version, NumPy libraries, hardware, etc). The relative performance of various methods can sometimes also depend on the size of M. So it pays to do your own benchmarks on an M which is closer to your actual use case.

For example, for slightly larger arrays M, calling the sum method twice may be fastest:

In [34]: M = np.random.random((100,100,2,2))

In [37]: %timeit M.sum(axis=0).sum(axis=0)
10000 loops, best of 3: 59.9 µs per loop

In [39]: %timeit np.einsum('ijkl->kl', M)
10000 loops, best of 3: 99 µs per loop

In [40]: %timeit np.sum(M, axis=(0,1))
10000 loops, best of 3: 182 µs per loop

In [36]: %timeit M.reshape(-1, 2, 2).sum(axis=0)
10000 loops, best of 3: 184 µs per loop

In [38]: %timeit sum(sum(M))
1000 loops, best of 3: 202 µs per loop