python - 计算以每个数据点为中心的固定半径球内数据点数的有效方法

Question

我有一个包含许多数据点的数据库，每个数据点都有一个 x,y,z 坐标。我想计算与相邻点一定距离内的点数。有些点会有一对在半径 R 内，有些则没有。我只是想计算一定距离内的对数。我可以轻松地编写一个算法来执行此操作，但效率不够(因为我会遍历每个数据点)。

这似乎在 astropy、scipy 等中肯定已经存在，但我似乎找不到我要找的东西。有什么东西可以做到这一点吗？

Answer 1

正如@Davis Herring 在评论中提到的，一个有效的选择是 k-d 树。

k-d 树是一种避免蛮力方法并允许有效距离计算的算法*(背景见答案底部)。

有几种 Python 实现，其中之一是通过 SciPy :

SciPy k-d tree in Cython (更快，因为它使用 C/Cython)

SciPy k-d tree in pure Python

您可以通过首先为您的 xyz 数据构建一个 k-d 树来使用它:

import numpy as np  #for later code
from scipy.spatial import cKDTree

kdtree = cKDTree(xyzData)

然后，你必须用一个点point查询k-d树计算 point 之间的距离及其最近的邻居。此查询的输出是距离 NN_dist在 point 之间及其最近的邻居和索引 NN_idx那个邻居的。为了对所有点进行计算，我们需要一个 for 循环，但考虑到 k-d 树算法，这比蛮力计算快得多:

NN_dists = np.zeros(numPoints)  #pre-allocate an array to store distances
for i in range(numPoints):
    point = xyzData[i]

    NN_dist, NN_idx = kdtree.query(point,k=[1])

    #Note: 'k' specifies the kth neighbor distance to compute, 
    #so set k=2 if you end up finding the point as its own "neighbor":
    if NN_dist == 0:
        NN_dist, NN_idx = targetTree.query(curCoord,k=[2])
    
    NN_dists[i] = NN_dist

(有关详细信息，请参阅 k-d tree query)。

然后，要找到低于某个阈值的距离，您可以在使用比较运算符(如 < )时使用 NumPy 数组的内置实用程序:

distanceThres = 10
goodIdx = NN_dists < distanceThres
goodPoints = xyzData[goodIdx]

这将为您提供索引 goodIdx和点goodPoints在您指定的距离阈值内 distanceThres (尽管您可能必须根据 xyz 坐标数据的形状/格式更改此代码)。

---

*关于 k-d 树的浅色背景(掩盖细节——更多信息请参见引用资料):k-d 树方法以一种避免计算每个点之间的距离的方式划分数据集(即蛮力法) ).它通过将数据集划分为二进制空间分区来构建 k-d 树来实现这一点。这些分区使得距离计算(例如，最近邻搜索)可以忽略远处分区中的数据点。此外，每个点都会重复使用同一个 k-d 树。

一般来说，网上有很多关于 k-d 树的资源。当我学习这个算法时，我发现这些引用资料最有帮助:Stanford k-d trees或 Princeton k-d trees .

如果您有任何问题，请告诉我 - 我自己在天文学项目中遇到过这个问题，所以我可以提供更多帮助。