Python - Sum 4D Array

Question

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

unutbu · Accepted Answer · 2014-07-19T16:23:49.330

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

Compared with `sum(sum(M))` (answered by Cyber), is `np.einsum` faster? — Lei Yu, Jul 19 '14 at 14:25
I added a timeit benchmark for `sum(sum(M))`. Note that 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. — unutbu, Jul 19 '14 at 14:30
I suspect this depends on the size of M (some quick experiments gave results which went against my intuitions, though, so I can't trust them.) — DSM, Jul 19 '14 at 14:38
@DSM: Thanks for the warning; indeed `einsum` may not be the fastest. — unutbu, Jul 19 '14 at 16:16
Einsum being faster is likely the result of SSE or AVX instruction sets. Although Numpy 1.8+ appears to include this natively for all operations. See [this](http://stackoverflow.com/questions/18365073/why-is-numpys-einsum-faster-than-numpys-built-in-functions) question for more details. — Daniel, Jul 19 '14 at 22:32

Jaime · Answer 2 · 2014-07-19T18:36:26.737

3

By far the simplest in recent numpy (version 1.7 or newer) is to do:

np.sum(M, axis=(0, 1))

This will not build an intermediate array, as a duplicate call to np.sum would.

edited Jul 19 '14 at 18:36

answered Jul 19 '14 at 14:57

Jaime

65,696
17
124
159

score 1 · Answer 3 · answered Jul 19 '14 at 14:16

import numpy as np
l = np.array([[[[ 4,  1],
                [ 2,  1]],
               [[ 8,  2],
                [ 4,  2]]],
              [[[ 8,  2],
                [ 4,  2]],
               [[12,  3],
                [ 6,  3]]]])
sum(sum(l))

Output

array([[32,  8],
       [16,  8]])

Python - Sum 4D Array

3 Answers3

Linked

Related