Question

我在2D A和B中有两组点，我需要找到A中每个点的最小距离，到B中的一个点。到目前为止，我一直在使用SciPy的cdist，代码如下

import numpy as np
from scipy.spatial.distance import cdist

def ABdist(A, B):
    # Distance to all points in B, for each point in A.
    dist = cdist(A, B, 'euclidean')
    # Indexes to minimum distances.
    min_dist_idx = np.argmin(dist, axis=1)
    # Store only the minimum distances for each point in A, to a point in B.
    min_dists = [dist[i][md_idx] for i, md_idx in enumerate(min_dist_idx)]

    return min_dist_idx, min_dists

N = 10000
A = np.random.uniform(0., 5000., (N, 2))
B = np.random.uniform(0., 5000., (N, 2))

min_dist_idx, min_dists = ABdist(A, B)

适用于N的小值。但是现在这些集合的长度从N=10000增加到N=35000，而我正在进入

    dm = np.zeros((mA, mB), dtype=np.double)
MemoryError

我知道我可以使用for循环替换cdist，该循环只保留A中每个点与B中每个点的最小距离（和索引），因为这是我所需要的。我不需要完整的AxB距离矩阵。但我一直在使用cdist，因为它很快。

有没有办法用尽快（几乎？）的实现替换cdist，但这不会占用那么多内存？

Answer 1

最佳方法将涉及使用专门为最近邻搜索设计的数据结构，例如k-d tree。例如，SciPy的cKDTree允许您以这种方式解决问题：

from scipy.spatial import cKDTree
min_dists, min_dist_idx = cKDTree(B).query(A, 1)

在计算和内存使用方面，结果比任何基于广播的方法都高效得多。

例如，即使有1,000,000点，计算也不会耗尽内存，而且我的笔记本电脑只需几秒钟：

N = 1000000
A = np.random.uniform(0., 5000., (N, 2))
B = np.random.uniform(0., 5000., (N, 2))

%timeit cKDTree(B).query(A, 1)
# 3.25 s ± 17.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Answer 2

诀窍是在这里最大化计算与内存比率。输出的长度为N，A中每个pt的一个索引和距离。我们可以将它减少到一个循环，每次迭代有一个输出元素，这将在每次迭代中处理所有B个pts，这带来了高计算比率。

因此，利用来自this post的einsum和matrix-multiplication，pt中的每个点A，我们将获得欧氏距离的平方，如此 -

for pt in A:
    d = np.einsum('ij,ij->i',B,B) + pt.dot(pt) - 2*B.dot(pt)

因此，概括它涵盖A和预计算np.einsum('ij,ij->i',B,B)中的所有点，我们会有这样的实现 -

min_idx = np.empty(N, dtype=int)
min_dist = np.empty(N)
Bsqsum = np.einsum('ij,ij->i',B,B) 
for i,pt in enumerate(A):
    d = Bsqsum + pt.dot(pt) - 2*B.dot(pt)
    min_idx[i] = d.argmin()
    min_dist[i] = d[min_idx[i]]
min_dist = np.sqrt(min_dist)

以块的形式工作

现在，完全向量化的解决方案将是 -

np.einsum('ij,ij->i',B,B)[:,None] + np.einsum('ij,ij->i',A,A) - 2*B.dot(A.T)

因此，要以块的形式工作，我们会将A的行切掉，这样做更容易简单地重塑为3D，就像这样 -

chunk_size= 100 # Edit this as per memory setup available
                # More means more memory needed
A.shape = (A.shape[0]//chunk_size, chunk_size,-1)

min_idx = np.empty((N//chunk_size, chunk_size), dtype=int)
min_dist = np.empty((N//chunk_size, chunk_size))

Bsqsum = np.einsum('ij,ij->i',B,B)[:,None]
r = np.arange(chunk_size)
for i,chnk in enumerate(A):
    d = Bsqsum + np.einsum('ij,ij->i',chnk,chnk) - 2*B.dot(chnk.T)
    idx = d.argmin(0)
    min_idx[i] = idx
    min_dist[i] = d[idx,r]
min_dist = np.sqrt(min_dist)

min_idx.shape = (N,)
min_dist.shape = (N,)
A.shape = (N,-1)

在2D中找到点组之间的最小距离（快速且不太耗费内存）

2 个答案:

以块的形式工作