scipy和numpy中的快速距离计算

fac*_*act 5 python performance numpy scipy multidimensional-array

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

例如:

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]))
Run Code Online (Sandbox Code Playgroud)

我用的地方scipy.spatial.distance.cdist

有更快的方法吗?理想情况下,我想获得output一个((day,observation,observation)),每一天包含观测值之间的成对距离阵列AB那一天,虽然在某种程度上避免了几天的循环。

Eel*_*aak 5

一种方法(尽管这将需要大量的内存)是巧妙地利用阵列广播:

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

编辑

但是经过一些测试,似乎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
Run Code Online (Sandbox Code Playgroud)

这是每天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
Run Code Online (Sandbox Code Playgroud)