迭代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在我的机器上进行,因此使用多处理模块的时间不应超过一天.

迭代这些列(with 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.

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是没有这样的矢量化.根据文件

"如果阵列不是1-D,它们将被展平为1-D."

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

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