2

I would like to get an array which is the sum of the sub matrices in an array.

For example, lets say we have a 3x3 matrix, where each item contains a 2x2 sub matrix:

matrix = np.array([[[[1,2],[3,4]],     [[5,6],[7,8]],     [[9,10],[11,12]]],
                   [[[13,14],[15,16]], [[17,18],[19,20]], [[21,22],[23,24]]],
                   [[[25,26],[27,28]], [[29,30],[31,32]], [[33,34],[35,36]]]])

which looks like:

[[[[ 1  2]
   [ 3  4]]

  [[ 5  6]
   [ 7  8]]

  [[ 9 10]
   [11 12]]]


 [[[13 14]
   [15 16]]

  [[17 18]
   [19 20]]

  [[21 22]
   [23 24]]]


 [[[25 26]
   [27 28]]

  [[29 30]
   [31 32]]

  [[33 34]
   [35 36]]]]

One way to get the answer is using list comprehension

ans = [ [ np.sum(sub_matrices) for sub_matrices in row ] for row in matrix ]

which will be:

[[10, 26, 42], [58, 74, 90], [106, 122, 138]]

I was wondering if there is a better way to get this result. Maybe using some numpy function built in?

user1179317
  • 2,693
  • 3
  • 34
  • 62

2 Answers2

4

Using einsum which is faster in this case than conventional sum:

np.einsum('ijkl->ij',matrix)

array([[ 10,  26,  42],
       [ 58,  74,  90],
       [106, 122, 138]])
Space Impact
  • 13,085
  • 23
  • 48
  • It really is faster. – klim Nov 06 '18 at 08:49
  • @klim Yes, it is. – Space Impact Nov 06 '18 at 08:49
  • Seriously, `einsum` is magic. Is there any place you would recommend (besides the official documentation) for reading up on it? I would never have thought of using it here. – Dominik Stańczak Nov 06 '18 at 09:20
  • 2
    I am surprised that `einsum` is faster than `.sum(axis=(2,3))` even in this case where no temporaries are involved. – MB-F Nov 06 '18 at 10:39
  • Sorry still trying to get familiar with einsum. Is it true that i and j is the rows and cols of the 3x3, and k and l is the rows and cols for the 2x2 matrix? And since k and l are omitted on the output side, we sum those together? Thanks for the very efficient answer btw. – user1179317 Nov 06 '18 at 14:20
  • @user1179317 Yes, and `einsum` is a pretty complicated function and it is not easy to master. Please check this [`link`](http://ajcr.net/Basic-guide-to-einsum/) and read the docs link provided in the answer and check its different functionalities to get familiar. – Space Impact Nov 06 '18 at 14:26
  • This `einsum` speed up is nice, but hardly ground breaking (it's not even 2x). It's advantage decreases with array size, and may even be reversed for some combinations of dimensions. – hpaulj Nov 07 '18 at 04:47
2

You can use .sum twice:

>>> matrix.sum(axis=2).sum(axis=2)
array([[ 10,  26,  42],
       [ 58,  74,  90],
       [106, 122, 138]])

Or, more elegantly, as suggested by @hpaulj, you can simply pass a tuple to the axis argument:

>>> matrix.sum(axis=(2,3))
array([[ 10,  26,  42],
       [ 58,  74,  90],
       [106, 122, 138]])
juanpa.arrivillaga
  • 88,713
  • 10
  • 131
  • 172