Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

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]]
like image 435
Lei Yu Avatar asked Jul 19 '14 14:07

Lei Yu


2 Answers

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
like image 192
unutbu Avatar answered Oct 06 '22 00:10

unutbu


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.

like image 29
Jaime Avatar answered Oct 06 '22 00:10

Jaime