高效的Numpy计算对两个向量进行平方差的方法。

Question

以下代码完美地实现了我想要的功能，即计算向量元素之间差的平方和（示例中长度为三），而我有一系列这样的向量（此处限制为五个）。期望结果显示在底部。 但是，实现感觉有两个问题：

1）需要添加虚拟维度，将形状从（5,3）更改为（5,1,3），以避免广播问题；

2）明显需要一个显式的 'for' 循环，这就是为什么在我的更大数据集上执行需要几个小时的原因（一百万个长度为2904的向量）。

是否有更有效和/或更符合Python风格的方法来实现相同的结果？

a = np.array([[ 4,  2,  3], [-1, -5,  4], [ 2,  1,  4], [-5, -1,  4], [6, -3,  3]])
a = a.reshape((5,1,3))

m = a.shape[0]
n = a.shape[2]
d = np.zeros((n,n))
for i in range(m):
    c = a[i,:] - np.transpose(a[i,:])
    c = c**2
    d += c

print d

[[   0.  118.  120.]
 [ 118.    0.  152.]
 [ 120.  152.    0.]]

- Grant Petty

Answer 1

如果您不介意依赖于 scipy，那么您可以使用来自scipy.spatial.distance库的函数：

In [17]: from scipy.spatial.distance import pdist, squareform

In [18]: a = np.array([[ 4,  2,  3], [-1, -5,  4], [ 2,  1,  4], [-5, -1,  4], [6, -3,  3]])

In [19]: d = pdist(a.T, metric='sqeuclidean')

In [20]: d
Out[20]: array([ 118.,  120.,  152.])

In [21]: squareform(d)
Out[21]: 
array([[   0.,  118.,  120.],
       [ 118.,    0.,  152.],
       [ 120.,  152.,    0.]])

- Warren Weckesser

---

这看起来很有趣。但是，正如在之前的问题回答中讨论的那样，实际数据集的大小很大（1M x 2094）。我不知道如何评估pdist对于这么大的数组（在我的代码中实现为np.memmap）的内存需求。对于这个问题有什么见解吗？ - Grant Petty

给定n个点，需要计算n*(n-1)/2个成对距离。因此，除非您可以分块处理（并且可以在计算下一块之前丢弃一块），否则您的内存需求为O(n^2)。pdist将在内存中创建此数组。如果要分块处理和/或直接将计算结果写入内存映射数组，则最好从@unutbu的答案开始。 - Warren Weckesser

Answer 2

您可以通过使用以下方法来消除for循环：

In [48]: ((a - a.swapaxes(1,2))**2).sum(axis=0)
Out[48]: 
array([[  0, 118, 120],
       [118,   0, 152],
       [120, 152,   0]])

请注意，如果a的形状为(N, 1, M)，则(a - a.swapaxes(1,2))的形状为(N, M, M)。请确保您有足够的RAM来容纳这么大的数组。页面交换也会使计算变得缓慢。

如果您的内存不足，您将不得不将计算分成块：

m, _, n = a.shape
chunksize = 10**4
d = np.zeros((n,n))
for i in range(0, m, chunksize):
    b = a[i:i+chunksize]
    d += ((b - b.swapaxes(1,2))**2).sum(axis=0)

这是在整个数组和逐行计算之间的一种折中方案。如果有一百万行，而块大小为10 ** 4，则循环只有100次迭代，而不是一百万次。因此，它应该比逐行计算快得多。选择最大的chunksize值，可以使计算在RAM中执行。

- unutbu

---

我没有提到的是（因为我认为它与我提到的特定问题无关），我正在使用的实际数据数组是一个 np.memmap 对象，占用了磁盘上的 12GB。我不确定它是 memmap 是否解决了您在顶部描述的特定操作的 RAM 问题，但我认为它应该可以。我想我可以尝试一下看看！ - Grant Petty

现在我想起来了，np.memmap 可能不会自动创建一个新的 memmap 对象来容纳 (N,M,M) 中间结果，这可能是个问题。而在我的应用中，M=2904，所以大小将达到 24 TB！ - Grant Petty