Кластеризация двумерных массивов numpy в более мелкие четырехмерные массивы

Question

Уважаемые пользователи Stackoverflow,

Мой скрипт Python сталкивается с проблемами производительности, потому что мне приходится разбивать 2D-таблицы с более чем 1 миллиардом элементов для потенциальных списков из сотен входных файлов.Я заменял свои вложенные циклы на вызовы манипулирования массивами, и в этом процессе я обнаружил, что numpy.take (который находит элементы в соответствии с набором индексов) и numpy.outer (который оценивает все возможные продукты между двумя элементами 1D массива) чрезвычайно полезны.Эти функции позволили мне умножить производительность моего кода на несколько сотен, где я мог бы их использовать.

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

Первый список индексов - th_t, второй список - dm_t,и матрица p_contact.4D массив кластерных элементов называется rc_p.Процедура кластеризации является следующей для цикла:

import numpy as np

th_t = [1, 3, 2, 1, 1, 3, 3, 0, 1, 0, 2, 1]
dm_t = [0, 0, 1, 0, 0, 1, 1, 1, 0, 0, 0, 0]
n_th = len(set(th_t))
n_dm = len(set(dm_t))
p_contact = [[0.0129, 0.0134, 0.0062, 0.0021, 0.0107, 0.0106, 0.0076, 0.0134, 0.0087, 0.0031, 0.0026, 0.0114]
[0.0123, 0.0021, 0.0033, 0.0120, 0.0099, 0.0125, 0.0001, 0.0018, 0.0030, 0.0059, 0.0038, 0.0125]
[0.0082, 0.0125, 0.0004, 0.0120, 0.0040, 0.0108, 0.0101, 0.0063, 0.0072, 0.0098, 0.0017, 0.0121]
[0.0096, 0.0008, 0.0073, 0.0100, 0.0123, 0.0104, 0.0077, 0.0025, 0.0106, 0.0126, 0.0031, 0.0033]
[0.0112, 0.0091, 0.0134, 0.0002, 0.0129, 0.0081, 0.0087, 0.0036, 0.0102, 0.0002, 0.0019, 0.0131]
[0.0099, 0.0081, 0.0037, 0.0004, 0.0135, 0.0005, 0.0025, 0.0086, 0.0091, 0.0016, 0.0130, 0.0011]
[0.0078, 0.0005, 0.0044, 0.0089, 0.0127, 0.0106, 0.0113, 0.0048, 0.0057, 0.0133, 0.0077, 0.0033]
[0.0017, 0.0010, 0.0048, 0.0052, 0.0113, 0.0066, 0.0133, 0.0092, 0.0020, 0.0125, 0.0011, 0.0023]
[0.0027, 0.0124, 0.0096, 0.0047, 0.0134, 0.0020, 0.0129, 0.0114, 0.0087, 0.0114, 0.0090, 0.0001]
[0.0032, 0.0014, 0.0038, 0.0114, 0.0058, 0.0017, 0.0089, 0.0057, 0.0022, 0.0056, 0.0046, 0.0094]
[0.0033, 0.0020, 0.0042, 0.0040, 0.0110, 0.0016, 0.0100, 0.0014, 0.0087, 0.0123, 0.0004, 0.0031]
[0.0010, 0.0029, 0.0054, 0.0015, 0.0064, 0.0060, 0.0131, 0.0064, 0.0073, 0.0097, 0.0132, 0.0092]]
n_sg = len(p_contact)
rc_p = np.zeros((n_th, n_dm, n_th, n_dm)) 
for i in range(n_sg): #n_sg can be about 40000
        for j in range(n_sg):
            rc_p[th_t[i]][dm_t[i]][th_t[j]][dm_t[j]] += p_contact[i][j]

Я пытался использовать различные функции numpy, чтобы избежать этого для цикла более миллиарда элементов, и в итоге я выполнил следующую процедуру:

import numpy as np

th_t = [1, 3, 2, 1, 1, 3, 3, 0, 1, 0, 2, 1]
dm_t = [0, 0, 1, 0, 0, 1, 1, 1, 0, 0, 0, 0]
n_th = len(set(th_t))
n_dm = len(set(dm_t))
p_contact = [[0.0129, 0.0134, 0.0062, 0.0021, 0.0107, 0.0106, 0.0076, 0.0134, 0.0087, 0.0031, 0.0026, 0.0114]
[0.0123, 0.0021, 0.0033, 0.0120, 0.0099, 0.0125, 0.0001, 0.0018, 0.0030, 0.0059, 0.0038, 0.0125]
[0.0082, 0.0125, 0.0004, 0.0120, 0.0040, 0.0108, 0.0101, 0.0063, 0.0072, 0.0098, 0.0017, 0.0121]
[0.0096, 0.0008, 0.0073, 0.0100, 0.0123, 0.0104, 0.0077, 0.0025, 0.0106, 0.0126, 0.0031, 0.0033]
[0.0112, 0.0091, 0.0134, 0.0002, 0.0129, 0.0081, 0.0087, 0.0036, 0.0102, 0.0002, 0.0019, 0.0131]
[0.0099, 0.0081, 0.0037, 0.0004, 0.0135, 0.0005, 0.0025, 0.0086, 0.0091, 0.0016, 0.0130, 0.0011]
[0.0078, 0.0005, 0.0044, 0.0089, 0.0127, 0.0106, 0.0113, 0.0048, 0.0057, 0.0133, 0.0077, 0.0033]
[0.0017, 0.0010, 0.0048, 0.0052, 0.0113, 0.0066, 0.0133, 0.0092, 0.0020, 0.0125, 0.0011, 0.0023]
[0.0027, 0.0124, 0.0096, 0.0047, 0.0134, 0.0020, 0.0129, 0.0114, 0.0087, 0.0114, 0.0090, 0.0001]
[0.0032, 0.0014, 0.0038, 0.0114, 0.0058, 0.0017, 0.0089, 0.0057, 0.0022, 0.0056, 0.0046, 0.0094]
[0.0033, 0.0020, 0.0042, 0.0040, 0.0110, 0.0016, 0.0100, 0.0014, 0.0087, 0.0123, 0.0004, 0.0031]
[0.0010, 0.0029, 0.0054, 0.0015, 0.0064, 0.0060, 0.0131, 0.0064, 0.0073, 0.0097, 0.0132, 0.0092]]

#prepare the flattened list of index pairs
th_t        = np.asarray(th_t)
dm_t        = np.asarray(dm_t)
thdm_stack  = np.stack((th_t, dm_t))
thdm_stack  = np.transpose(thdm_stack)
thdm_table  = np.asarray(list(product(thdm_stack, thdm_stack)))
p_contact_f = p_contact.flatten()

#calculate clustered probabilities for each contact type
rc_p                 = np.zeros((n_th, n_dm, n_th, n_dm)) 

for th1 in range(n_th):
    for dm1 in range(n_dm):
        for th2 in range(n_th):
            for dm2 in range(n_dm):
                to_find                       = np.zeros((2, 2))
                to_find[0][0]                 = th1
                to_find[0][1]                 = dm1
                to_find[1][0]                 = th2
                to_find[1][1]                 = dm2
                condition                     = np.isin(thdm_table, to_find)
                condition                     = np.all(condition, axis=(1, 2))
                to_add                        = np.extract(condition, p_contact_f)
                rc_p[th1][dm1][th2][dm2]      = np.sum(to_add)

, который оказывается медленнее, чем исходная процедура, а не быстрее, вероятно, потому что мне нужно сгенерировать логическую матрицу размером в 1 миллиард и обрабатывать ее на каждом из тысяч шагов в цикле 4D for (который имеет тысячи меньшеэлементы, чем начальный цикл for, просто чтобы напомнить).

Итак, кто-нибудь из вас имеет представление о том, как я мог бы заменить этот дорогостоящий вложенный цикл for и максимально использовать базовый C-код numpy длякластеризовать эту большую 2D матрицу в гораздо меньший 4D массив?

Обратите внимание, что отдельные элементы в этих массивах являются вероятностями.Общая сумма всех элементов в 2D-массиве и в 4D-кластеризованном массиве равна 1, и под «кластеризацией» я подразумеваю группирование вероятностей по типам (все ячейки 2D-матрицы, которые отображают идентичные наборы индексов, получают свои вероятности, добавленные в одинэлементов кластерного массива 4D).

Всего наилучшего!

Answer 1

Я хочу подчеркнуть очень полезную функцию, предложенную Дэниелом Ф., которую я не знал и которая была ключевой для решения этой проблемы:

numpy.ravel_multi_index

Этоможет преобразовать индексные последовательности в 1D индексный список.Например, с парой индексов, основанной на двух списках из 2 и 9 индексов, индекс 1,4, выводимый этой функцией numpy, является 14-м индексом (9 индексов плюс 5).Это немного сложно понять, но очень мощно.

Answer 2

Вы на самом деле итерируете не по четырем измерениям, а по 2: i и j.Вы можете np.ravel_multi_index вместе с массивами th_t и dm_t уменьшить проблему до 2d, а в конце reshape вернуть до 4d:

idx = np.ravel_multi_index((th_t, dm_t), (n_th, n_dm))
rc_p = np.zeros((n_th * n_dm, n_th * n_dm))
for i in range (idx.size):
    np.add.at(rc_p[idx[i]], idx, p_contact[i])
rc_p = rc_p.reshape(n_th, n_dm, n_th, n_dm)

Или, если вы можете использовать numba, просто оберните ваш начальный зацикленный код в @jit, который скомпилирует его

from numba import jit

@jit
def foo(p_contact, th_t, dm_t, n_th, n_dm):
    n_sg = len(p_contact)
    rc_p = np.zeros((n_th, n_dm, n_th, n_dm)) 
    for i in range(n_sg): 
        for j in range(n_sg):
            rc_p[th_t[i]][dm_t[i]][th_t[j]][dm_t[j]] += p_contact[i][j]