Существует ли функция numpy для получения суммы подматриц?

Question

Я хотел бы получить массив, который является суммой подматриц в массиве.

Например, допустим, у нас есть матрица 3x3, где каждый элемент содержит субматрицу 2x2:

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

, который выглядит как:

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

Один из способов получить ответ - использовать понимание списка

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

что будет:

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

Мне было интересно, есть ли лучший способ получить этот результат. Может быть, с помощью какой-нибудь встроенной функции numpy?

Answer 1

Использование einsum, что в этом случае быстрее, чем обычного sum:

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

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

Answer 2

Вы можете использовать .sum дважды:

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

Или, что более элегантно, как предлагает @hpaulj, вы можете просто передать кортеж аргументу axis:

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