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):

最麻烦的方式是什么？

编辑：我改变了标题，因为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没有这样的矢量化。根据{{3}}

“如果数组不是1-D，它们将被展平为1-D。”

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

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

Answer 2

如果您不想使用递归，通常应该使用itertools.combinations.没有具体原因（afaik）为什么这会导致代码运行速度变慢。计算密集型部分仍然由numpy处理。 Itertools还具有可读性的优势。

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

2 个答案: