python – 迭代numpy数组列的所有成对组合

我有一个大小的numpy数组

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

我想在最后两个维度的每个成对组合上计算scipy.stats.kendalltau.例如：

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

这样我用arr [：,j,xj]覆盖arr [：,i,xi]的所有组合,其中i< j和xi在[0,20]中,xj在[0,20]中.这是(600选择2)* 400个人计算,但由于我的机器上每个计算大约需要0.002秒,因此使用多处理模块的时间不应超过一天. 迭代这些列的最佳方法是什么(使用i< j)？我想我应该避免像

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

这种做法最简单的方法是什么？

编辑：我更改了标题,因为Kendall Tau对这个问题并不重要.我意识到我也可以做类似的事情

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

但是必须有一个更好的,更加矢量化的方式与numpy.

解决方法

矢量化这样的东西的一般方法是使用广播来创建集合的笛卡尔积.在你的情况下,你有一个形状的阵列arr(200,20),所以你会有两个视图：

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

为了清楚起见,上面两行已经扩展,但我通常会写相同的：

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

如果你有一个矢量化函数f,它在除了最后一个维度之外的所有维度上进行广播,那么你可以这样做：

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

然后out将是一个形状数组(200,600),out [i,k,l]保持f的值(arr [i,j],arr [k,l]) .例如,如果您想计算所有成对内部产品,您可以：

from numpy.core.umath_tests import inner1d

out = inner1d(arr[:,arr)

不幸的是,scipy.stats.kendalltau没有这样的矢量化.根据the docs

“If arrays are not 1-D,they will be flattened to 1-D.”

所以你不能这样做,你最终会做Python嵌套循环,无论是使用itertools还是在np.vectorize下伪装它都明确写出来.这会很慢,因为Python变量的迭代,并且因为每个迭代步骤都有一个Python函数,这些都是昂贵的操作.

请注意,当你可以采用矢量化方式时,有一个明显的缺点：如果你的函数是可交换的,即如果f(a,b)== f(b,a),那么你正在进行两次所需的计算.根据实际计算的成本,这通常会因没有任何Python循环或函数调用而增加速度.