遍历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在[0,20),xj在[0,20))的组合。这将需要进行(600 choose 2) * 400个单独的计算, 但由于我的机器每个计算只需约0.002 s, 所以使用多进程模块不会花费太长时间，大约一天左右。

如何迭代这些列（其中j）最好的方法是什么？我认为应该避免像下面这样的操作：

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

什么是最符合Python风格的方式来处理这个问题？ 编辑：我改变了标题，因为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进行向量化处理。

- wflynny

---

你想了解递归。这个问题已经在Java中得到了回答：https://dev59.com/4XRC5IYBdhLWcg3wCMrX - Mike Vella

我认为递归不会按照预期使用numpy。 - wflynny

你为什么认为numpy不能与递归一起使用？ - Mike Vella

在Python中，迭代似乎比递归更受推荐。Numpy通过向量化将其提升到了另一个层次。当然，递归也可以在Numpy中使用，但我认为还有一种更具“Python风格”的迭代方法。 - wflynny

你可以看一下这个讨论：https://dev59.com/CGQo5IYBdhLWcg3wfPjK，但我认为你不应该太纠结于此，itertools.combinations很好用！当它导致问题时再担心它——过早优化是万恶之源。 - Mike Vella

这个线程很有帮助。谢谢！ - wflynny

Answer 1

向量化这种情况的一般方法是使用广播来创建集合与自身的笛卡尔积。在您的情况下，您有一个形状为(200, 600, 20)的数组arr，因此您需要获取它的两个视图：

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，它在除了最后一个维度以外的所有维度上都进行广播，那么你可以这样做：

如果您拥有一个向量化函数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循环或函数调用所带来的速度提高所抵消。

- Jaime

Answer 2

如果您不想使用递归，通常应该使用itertools.combinations。没有特定的原因（据我所知）会导致您的代码运行变慢。计算密集型部分仍然由numpy处理。Itertools还具有可读性的优点。

- Mike Vella

---

是的，我几分钟前编辑了我的帖子，加入了那个选项。 - wflynny