Question

我有三个numpy数组中包含一堆坐标：xarr，yarr和zarr（每个数组中的相应位置属于同一个点 - 即第一个点位于xarr[0]，yarr[0]，zarr[0]）。鉴于P(x,y,z)处的空间中的另一个点，我想找到P(x,y,z)距离 r 内的所有点。

我目前（并且非常低效）的方法是简单地迭代并计算到每个点的距离，看它是否在P(x,y,z)的 r 之内。

但是，我想使用SciPy的k-d树算法来做到这一点，但我不确定如何开始实现它（我对Python很新）。我真的很感激，如果有人可以简要地概述一些代码，这些代码演示如何以我所拥有的格式给出数据来设置 k-d树。

我知道SciPy documentation of its k-d tree implementation，我已经查看了它，但我仍然对如何创建树感到困惑，因为我已经调用了np.mgrid和ravel()格式的数据，我不太明白为什么。

谢谢！

Answer 1

以下是对scipy docs中提供的示例的解释：

from scipy import spatial
x, y = np.mgrid[0:4, 0:4]

np.mgrid创建一个从0到4的网格x，y。由于您已经有x，y，z坐标，因此您将跳过此步骤。

points = zip(x.ravel(), y.ravel())
points = zip(xarr.ravel(), yarr.ravel(), zarr.ravel()) #in your case
points = zip(xarr, yarr, zarr) # if x,y,z are already 1-d

zip创建一个包含每个x，y点对的元组列表（将每个点的坐标关联在一起）。 ravel展平x，y网格网格（将n-d数组转换为1-d），以便可以使用zip。在您的情况下，如果ravel，xarr，yarr还不是1-d，则只会使用zarr。

tree = spatial.KDTree(points)

创建索引点以提供快速邻居查找。

tree.query_ball_point([2, 0], 1)

在点r=1

的[2,0]内查找点数

希望这会有所帮助。

Answer 2

省略进口......我没有你的数据，所以我不得不伪造一些......

In [40]: x = np.random.random(100)
In [41]: y = np.random.random(100)    
In [42]: z = np.random.random(100)    
In [43]: p = np.random.random(3)    
In [44]: p
Out[44]: array([ 0.60515083,  0.39263392,  0.36129813])

即三个坐标数组和一个我将搜索邻居的点。

接下来，让我们看看如何构建一个包含与不同数据点和三个colunms一样多行的数组...

In [45]: np.vstack((x,y,z)).T.shape
Out[45]: (100, 3)

嗯，这是正确的。

我们使用KDTree

中的scipy.spatial来构建k-d树

In [46]: tree = KDTree(np.vstack((x,y,z)).T)

然后我们使用树的一种方法，即恰当命名的.query_ball_point()来查找p附近点的索引

In [47]: indices = tree.query_ball_point(p, 0.33)

我在任意地使用半径等于1/3。

最终我们希望看到这些邻居，所以我将使用树的.data属性以及我刚才计算的索引

In [48]: tree.data[indices]
Out[48]: 
array([[ 0.4117843 ,  0.21440852,  0.3352732 ],
       [ 0.48921727,  0.13855976,  0.43331816],
       [ 0.71598133,  0.32270361,  0.20292187],
       [ 0.71761991,  0.27309708,  0.12670474],
       [ 0.6282775 ,  0.13752325,  0.4143872 ],
       [ 0.55995847,  0.31302848,  0.2780926 ],
       [ 0.75896359,  0.16043536,  0.33530071],
       [ 0.81138529,  0.64635994,  0.33819097],
       [ 0.43537193,  0.5353203 ,  0.52095431],
       [ 0.66996807,  0.48346547,  0.52761835],
       [ 0.69426851,  0.24725511,  0.57650329],
       [ 0.5350322 ,  0.23155768,  0.62545958],
       [ 0.51228139,  0.38078056,  0.61246054]])

那就是......

使用k-d树搜索邻近点

2 个答案: