在scipy和numpy中快速计算距离

Question

设A,B为((day,observation,dim))数组。每个数组在给定的一天中包含相同数量的观测值，观测值是一个具有dim维度的点(即dim浮点数)。对于每一天，我想计算A和B当天所有观测值之间的空间距离。

例如:

import numpy as np
from scipy.spatial.distance import cdist
A, B = np.random.rand(50,1000,10), np.random.rand(50,1000,10)
output = []
for day in range(50):
    output.append(cdist(A[day],B[day]))

我使用scipy.spatial.distance.cdist .

有更快的方法吗?理想情况下，我希望为output获得((day,observation,observation))数组，该数组每天包含A和B当天的观测值之间的成对距离，同时以某种方式避免几天的循环。

Answer 1

一种方法(尽管它需要大量内存)是巧妙地使用数组广播:

output = np.sqrt( np.sum( (A[:,:,np.newaxis,:] - B[:,np.newaxis,:,:])**2, axis=-1) )

---

编辑

但是经过一些测试，似乎scikit-learn的euclidean_distances是大型数组的最佳选择。(注意，我已经将你的循环重写为一个列表推导式。)

这是每天100个数据点:

# your own code using cdist
from scipy.spatial.distance import cdist
%timeit dists1 = np.asarray([cdist(x,y) for x, y in zip(A, B)])
100 loops, best of 3: 8.81 ms per loop
# pure numpy with broadcasting
%timeit dists2 = np.sqrt( np.sum( (A[:,:,np.newaxis,:] - B[:,np.newaxis,:,:])**2, axis=-1) )
10 loops, best of 3: 46.9 ms per loop
# scikit-learn's algorithm
from sklearn.metrics.pairwise import euclidean_distances
%timeit dists3 = np.asarray([euclidean_distances(x,y) for x, y in zip(A, B)])
100 loops, best of 3: 12.6 ms per loop

，这是每天2000个数据点:

In [5]: %timeit dists1 = np.asarray([cdist(x,y) for x, y in zip(A, B)])
1 loops, best of 3: 3.07 s per loop
In [7]: %timeit dists3 = np.asarray([euclidean_distances(x,y) for x, y in zip(A, B)])
1 loops, best of 3: 2.94 s per loop

Answer 2

编辑:我是个白痴，忘记了python的map是惰性求值的。我的"更快"代码实际上并没有做任何工作!强制评估取消了性能提升。

我认为你的时间将被花在scipy函数内部的时间所支配。我将使用map而不是循环，因为我认为它更简洁，但我不认为有任何神奇的方法可以在这里获得巨大的性能提升。也许用cython编译代码或使用numba会有一点帮助。