Перебирать все попарные комбинации столбцов массива numpy

у меня есть массив numpy размера

arr.size = (200, 600, 20). 

Я хочу, чтобы вычислить scipy.stats.kendalltau на каждой попарной комбинации последних двух измерений. Например:

kendalltau(arr[:, 0, 0], arr[:, 1, 0])
kendalltau(arr[:, 0, 0], arr[:, 1, 1])
kendalltau(arr[:, 0, 0], arr[:, 1, 2])
...
kendalltau(arr[:, 0, 0], arr[:, 2, 0])
kendalltau(arr[:, 0, 0], arr[:, 2, 1])
kendalltau(arr[:, 0, 0], arr[:, 2, 2])
...
...
kendalltau(arr[:, 598, 20], arr[:, 599, 20])

такие, что я покрываю все комбинации arr[:, i, xi] С arr[:, j, xj] С i < j и xi in [0,20), xj in [0, 20). Это (600 choose 2) * 400 индивидуальные расчеты, но так как каждый принимает о 0.002 s на моей машине это не должно занять больше дня с модулем многопроцессорной обработки.

что лучший способ итерации по этим столбцам (с i<j)? Я считаю, что я должен избегать чего-то вроде

for i in range(600):
    for j in range(i+1, 600):
        for xi in range(20):
            for xj in range(20):

каков самый нумпитонный способ сделать это?

Edit: я изменил название, так как Кендалл Тау не очень важен для вопроса. Я понимаю, что я также мог бы сделать что-то вроде

import itertools as it
for i, j in it.combinations(xrange(600), 2):
    for xi, xj in product(xrange(20), xrange(20)):

но должен быть лучший, более векторизованный способ с numpy.

2 ответов


общий способ векторизации чего-то подобного-использовать широковещание для создания декартового произведения множества с самим собой. В вашем случае у вас есть массив arr формы (200, 600, 20), так что вы бы взять два взгляда на это:

arr_x = arr[:, :, np.newaxis, np.newaxis, :] # shape (200, 600, 1, 1, 20)
arr_y = arr[np.newaxis, np.newaxis, :, :, :] # shape (1, 1, 200, 600, 20)

выше две линии были расширены, для ясности, но я обычно пишу эквивалент:

arr_x = arr[:, :, None, None]
arr_y = arr

если у вас есть функция векторизации, f, что вещание на всех, кроме последнего измерения, можно тогда сделай:

out = f(arr[:, :, None, None], arr)

а то out будет массив формы (200, 600, 200, 600) С out[i, j, k, l] держа стоимостью f(arr[i, j], arr[k, l]). Например, если вы хотите вычислить все попарные внутренние продукты, вы можете сделать:

from numpy.core.umath_tests import inner1d

out = inner1d(arr[:, :, None, None], arr)

к сожалению scipy.stats.kendalltau не векторизован, как это. Согласно документы

"если массивы не являются 1-D, они будут сплющены до 1-D."

поэтому вы не можете пойти об этом это, и вы собираетесь в конечном итоге делать вложенные циклы Python, будь то явное их написание, используя itertools или маскировка под np.vectorize. Это будет медленно, из-за итерации по переменным Python и потому, что у вас есть функция Python на шаг итерации, которые являются дорогостоящими действиями.

обратите внимание, что, когда вы можете пойти векторизованным способом, есть очевидный недостаток: если ваша функция коммутативна, т. е. если f(a, b) == f(b, a), тогда вы делаете дважды, необходимые расчеты. В зависимости от того, насколько дорого ваше фактическое вычисление, это очень часто компенсируется увеличением скорости от отсутствия циклов Python или вызовов функций.


Если вы не хотите использовать рекурсию следует использовать itertools.сочетания. нет никакой конкретной причины (afaik), почему это должно заставить ваш код работать медленнее. Вычислительно-интенсивные части все еще обрабатываются numpy. Itertools также имеет преимущество читаемости.