python - 使用 scipy 在其他点的截止距离内查找点

Question

假设我有两组点:

>>> points1.shape
(10000, 3)
>>> points2.shape
(1529, 3)

我想要 points1 的索引列表，points2 中的一个点的欧氏距离 cutoff 内。我可以像这样使用 scipy.spatial.distance.cdist 轻松做到这一点:

from scipy.spatial.distance import cdist
import numpy

indices = numpy.argwhere(cdist(points1, points2).min(axis=0) < cutoff)

然而，这似乎效率低下，因为我不需要知道点之间的距离有多远，只需要知道它们是否在截止距离内。 KDTree 可以帮助解决这个问题吗？

Answer 1

这里有 3 个备选方案，一个使用 cdist，两个使用 scipy.spatial.cKDTree :

import itertools as IT
import numpy as np
import scipy.spatial as spatial
import scipy.spatial.distance as dist
np.random.seed(2016)
points1 = np.random.randint(100, size=(10**5, 3))
points2 = np.random.randint(100, size=(1529, 3))
cutoff = 5

def using_cdist(points1, points2, cutoff):
    indices = np.where(dist.cdist(points1, points2) <= cutoff)[0]
    indices = np.unique(indices)
    return indices

def using_kdtree(points1, points2, cutoff):
    # build the KDTree using the *smaller* points array
    tree = spatial.cKDTree(points2)
    groups = tree.query_ball_point(points1, cutoff)
    indices = np.unique([i for i, grp in enumerate(groups) if len(grp)])
    return indices

def using_kdtree2(points1, points2, cutoff):
    # build the KDTree using the *larger* points array
    tree = spatial.cKDTree(points1)
    groups = tree.query_ball_point(points2, cutoff)
    indices = np.unique(IT.chain.from_iterable(groups))
    return indices

cdist_result = using_cdist(points1, points2, cutoff)
kdtree_result = using_kdtree(points1, points2, cutoff)
kdtree_result2 = using_kdtree2(points1, points2, cutoff)
assert np.allclose(cdist_result, kdtree_result)
assert np.allclose(cdist_result, kdtree_result2)

在这 3 个备选方案中，using_kdtree2 是最快的:

In [80]: %timeit using_kdtree3(points1, points2, cutoff)
10 loops, best of 3: 92.4 ms per loop

In [103]: %timeit using_kdtree(points1, points2, cutoff)
1 loops, best of 3: 938 ms per loop

In [104]: %timeit using_cdist(points1, points2, cutoff)
1 loops, best of 3: 1.51 s per loop

我关于什么是最快的直觉被证明是完全错误的。我认为使用较小的点数组构建 KDTree 是最快的。即使使用更大的点数组构建 KDTree 是有点慢，在较小的点数组上调用 tree.query_ball_point 是快得多:

In [68]: %timeit tree = spatial.cKDTree(points2)
1000 loops, best of 3: 312 µs per loop

In [69]: %timeit tree = spatial.cKDTree(points1)
10 loops, best of 3: 45.7 ms per loop

In [66]: %timeit tree = spatial.cKDTree(points2); groups = tree.query_ball_point(points1, cutoff)
1 loops, best of 3: 933 ms per loop

In [67]: %timeit tree = spatial.cKDTree(points1); groups = tree.query_ball_point(points2, cutoff)
10 loops, best of 3: 89.3 ms per loop

---

注意使用有一些问题

def orig(points1, points2, cutoff):
    return np.argwhere(dist.cdist(points1, points2).min(axis=0) < cutoff)

首先，通过调用 min(axis=0)，如果 points1 中有两个点，您将丢失信息都在 points2 中某个点的 cutoff 范围内。你只会得到索引的最近点。另一个问题是，通过调用 min0 轴，剩下的就是与 points2 关联的 1 轴。所以orig 将索引返回到 points2，而不是 points1。