Избежание циклов при использовании суммы NumPy

Question

Мне часто нужно произвести суммирование по определенным строкам или столбцам большего массива NumPy. Например, возьмите этот массив:

>>> c = np.arange(18).reshape(3, 6)
>>> print(c)
[[ 0  1  2  3  4  5]
 [ 6  7  8  9 10 11]
 [12 13 14 15 16 17]]

Предположим, я хочу суммировать только в тех случаях, когда индекс строки равен 0 или 2, И индекс столбца равен 0, 2, 4 или 5. Другими словами, Мне нужна сумма подмассива

[[ 0  2  4  5]
 [12 14 16 17]]

Я обычно делаю это с помощью невероятно полезного метода NumPy ix_; например,

>>> np.sum(c[np.ix_([0,2],[0,2,4,5])])
70

Пока все хорошо. Теперь, однако, предположим, что у меня есть другой массив, e, который похож на c, но имеет два ведущих измерения. Таким образом, его форма (2,3,3,6) вместо просто (3,6):

e = np.arange(108).reshape(2, 3, 3, 6)

(обратите внимание, что фактические массивы, с которыми я работаю, могут содержать любые случайные целые числа ; они не содержат последовательных целых чисел, как в этом примере.)

Я ищу вычисление выше для каждой комбинации строки / столбца. Следующее работает для этого простого примера, но для больших массивов с большим количеством измерений это может быть очень, очень медленным:

new_sum = np.empty((2,3))
for i in range(2):
   for j in range(3):
      temp_array = e[i,j,:,:]
      new_sum[i,j] = np.sum(temp_array[np.ix_([0,2],[0,2,4,5])])

Вопрос: Можно ли сделать это быстрее. , предположительно, не прибегая к использованию циклов?

В качестве сноски результат вышесказанного выглядит следующим образом:

>>> print(new_sum)
[[ 70. 214. 358.]
 [502. 646. 790.]]

Конечно, 70 в верхнем левом углу - это тот же результат мы получили раньше.

Answer 1

Вы можете создать логическую матрицу (маску), которая будет иметь True для значений, которые вы хотите сохранить, и False для тех, которые вам не нужны.

>>> mask = np.zeros((3,6), dtype='bool')
>>> mask[np.ix_([0,2],[0,2,4,5])] = True
>>> mask
array([[ True, False,  True, False,  True,  True],
       [False, False, False, False, False, False],
       [ True, False,  True, False,  True,  True]])

Затем вы можете взять Преимущество правил широковещательной рассылки numpy массива: применяют маску к массиву и суммируют по последним измерениям:

>>> new_sum = np.sum(e * mask.reshape(1,1,3,6), axis=(2,3))
>>> new_sum
array([[ 70, 214, 358],
       [502, 646, 790]])

Вот небольшой код, который сравнивает характеристики двух версий на большей матрице:

import numpy as np
import time

N, P = 200, 100
e = np.arange(18*N*P).reshape(N, P, 3, 6)

t_start = time.time()
new_sum = np.empty((N,P))
for i in range(N):
   for j in range(P):
      temp_array = e[i,j,:,:]
      new_sum[i,j] = np.sum(temp_array[np.ix_([0,2],[0,2,4,5])])
print(f'Timer 1: {time.time()-t_start}s')

t_start = time.time()
mask = np.zeros((3,6), dtype='bool')
mask[np.ix_([0,2],[0,2,4,5])] = True
new_sum_2 = np.sum(e * mask.reshape(1,1,3,6), axis=(2,3))
print(f'Timer 2: {time.time()-t_start}s')

print('Results are equal!' if np.allclose(new_sum, new_sum_2) else 'Results differ!')

Вывод:

% python3 script.py
Timer 1: 0.4343228340148926s
Timer 2: 0.002004384994506836s
Results are equal!

Как видите, вы получаете значительное улучшение с точки зрения времени вычислений.

Answer 2

Немного более быстрый подход (если ваш массив большой, а выбранные индексы малы), чем сумма (та же идея), что и решение @alaniwi:

  np.einsum('ijkl->ij',e[np.ix_(np.arange(e.shape[0]),np.arange(e.shape[1]),[0,2],[0,2,4,5])])

[[ 70, 214, 358],
 [502, 646, 790]]

Answer 3

Просто небольшое расширение вашей собственной идеи на самом деле: дайте еще несколько измерений np.ix_, а затем просуммируйте по последним двум осям.

import numpy as np

e = np.arange(108).reshape(2, 3, 3, 6)

indices = [[0,2], [0,2,4,5]]

print(
  np.sum(e[np.ix_(*[range(i) for i in e.shape[:-2]], *indices)], axis=(-2,-1))
)

Это дает:

array([[ 70, 214, 358],
       [502, 646, 790]])

Таким образом, аргументы для np.ix_ в этом случае:

range(2), range(3), [0,2], [0,2,4,5]

Для большей общности мы также могли бы воздержаться от предположений о количестве осей, используемых в индексах. list:

np.sum(e[np.ix_(*[range(i) for i in e.shape[:-len(indices)]], *indices)],
       axis=tuple(range(-len(indices),0)))

(Тип данных будет таким же, как для e, что имело бы место в примере в вопросе, если бы тот же тип данных был указан при вызове np.empty. Я предполагаю, что нет особой причины приводить к np.float.)