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

Question

У меня есть 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): 

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

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

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

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

Answer 1

Общий способ векторизации чего-либо подобного заключается в использовании радиовещания для создания декартова произведения набора с самим собой. В вашем случае у вас есть массив 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 не векторизации, как это. Согласно the docs

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

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

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

Answer 2

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