Эффективно работает на подматрицах двумерных массивов

Question

У меня проблемы с выполнением следующей операции матрицы. Начиная с квадратной матрицы (двумерный массив) и групп, которые охватывают каждый индекс матрицы (словарь: ключи - это группы, значения - это списки индексов матрицы в группе), мне нужно получить новую меньшую матрицу, которая содержит сумму элементов в каждой подматрице исходной матрицы. Подматрицы определяются в соответствии с индексами групп. Поэтому новая матрица также будет квадратной, но с количеством групп в качестве ее размеров.

Давайте рассмотрим следующий пример:

import numpy as np

X = np.arange(49).reshape((7, 7))

d = {0: [0, 1], 1: [2, 3, 4], 2: [5, 6]}

def get_new_matrix(matrix, groups_indexes):
    groups_number = len(groups_indexes)
    new_matrix = np.zeros((groups_number, groups_number))
    for i in range(groups_number):
        for j in range(groups_number):
            new_matrix[i][j] = np.sum(matrix[groups_indexes[i]][:,groups_indexes[j]])
    return new_matrix

Z = get_new_matrix(X, d)
print(Z)

[[ 16  39  36]
 [129 216 159]
 [156 249 176]]

Глядя на результат, например, во (второй) строке 1 и (третьем) столбце 2, мы видим, что результат равен 159, это:

Z[1,2]

Это означает, что в исходной матрице подматрица, определенная группами 1 в строках и 2 в столбцах, это строки 2, 3 и 4 и столбцы 5 и 6, явно:

X[[2, 3, 4]][:,[5, 6]]

, а сумма всех элементов в подматрице равна 19 + 20 + 26 + 27 + 33 + 34 = 159.

явно:

np.sum(X[[2, 3, 4]][:,[5, 6]])

Есть ли способ написать более питонический код, избегая двух циклов for для получения новой матрицы и повышая общую эффективность? Полагаю, это должно быть что-то вроде необычного индексирования, вещания и т. Д. , но я не смог найти лучшего решения.

Мой текущий код ужасно масштабируется для больших начальных матриц (и потенциально большого начального числа групп), и, поскольку я буду запускать его не только для произвольных больших начальных квадратных матриц, но и на многих итерациях, мне действительно нужно улучшить Это. Или, может быть, нет способа улучшить код, и объяснение будет очень полезно:)

Answer 1

Если ваши групповые индексы охватывают всю матрицу и являются смежными, вы можете хранить их как просто индексы вместо словаря. Поскольку каждая группа заканчивается началом следующей группы, вам нужно только сохранить начальный индекс. Ваш текущий d может быть переписан как

d = sorted(val[0] for val in d.values())

Или, если вы не привязаны к формату словаря, просто

d = np.array([0, 2, 5])

Я рекомендую применить np.add.reduceat дважды, по одному разу для каждого измерения, по существу, как вы делаете в текущем цикле, но при этом внутренне управляйте циклом:

result = np.add.reduceat(np.add.reduceat(X, d, axis=0), d, axis=1)

Результат для ввода в вопросе:

array([[ 16,  39,  36],
       [129, 216, 159],
       [156, 249, 176]])

159 действительно является элементом индекса [1, 2].

Это, кажется, масштабируется довольно хорошо. Работа с X = np.arange(10**6).reshape(10**3, 10**3) и d = np.arange(0, 10**3, 10) занимает около 2,27 мс на моем не слишком мощном ноутбуке. Я не думаю, что этот фрагмент кода может стать узким местом для всего, что вы делаете.