How to find set of points in x,y grid using KDTree.query_ball_tree

北城以北 提交于 2020-06-24 14:41:27

问题


I am working in python and I have a x,y mesh grid which are numpy arrays. I need to find for each point (x1,y1) in the grid, the points which are present at a distance r from (x1,y1). Scipy has a function KDTree.query_ball_tree which takes as input, a KD Tree object (which can be constructed from the numpy arrays) and a distance r, but I am not able to understand how it works.

For example, consider the following points below:

[(1, 1), (2, 1), (3, 1), (4, 1), (1, 2), (2, 2), (3, 2), (4, 2), (1, 3), (2, 3), (3, 3), (4, 3), (1, 4), (2, 4), (3, 4), (4, 4)]`

I want to find all the points which are at a distance 2 from (1,1). The the output should be:

[(1,2),(1,3),(2,1),(3,1)]

I am using KDTree because, I want to avoid for loops for traversing the grid, because the mesh grid is 601x90 (YxX) and it will not be optimum in time, if for loops are used. Can someone provide me with an example illustrating KDTree.query_ball_tree for my situation?


回答1:


If you are looking for all points close within a distance of a single point, use scipy.spatial.KDTree.query_ball_point not query_ball_tree. The latter when you need to compare sets of points against each other.

import numpy as np
from scipy.spatial import KDTree

pts = np.array([(1, 1), (2, 1), (3, 1), (4, 1), (1, 2), (2, 2), (3, 2), (4, 2), (1, 3), (2, 3), (3, 3), (4, 3), (1, 4), (2, 4), (3, 4), (4, 4)])

T = KDTree(pts)
idx = T.query_ball_point([1,1],r=2)
print pts[idx]

This returns

[[1 1]
 [2 1]
 [1 2]
 [2 2]
 [1 3]
 [3 1]]

Note that your output must include the point (1,1) as well since that is a distance of zero from your target.




回答2:


Building on @Hooked 's answer, the following finds data points with two known coordinates in a data set which has three coordinate values.

import numpy as np
from scipy.spatial import KDTree

pts = np.array([[1, 1, 0], [2, 1, 1], [3, 1, 2], [4, 1, 3], [1, 2, 4], [2, 2, 5], [3, 2, 6], 
                [4, 2, 7], [1, 3, 8], [2, 3, 9], [3, 3, 10], [4, 3, 11], [1, 4, 12], [2, 4, 13], [3, 4, 14], [4, 4, 15]])
pts_cut=[]
pts_cut=pts[:,0:2]

T = KDTree(pts_cut)
idx = T.query_ball_point([1,1], r=2)
print(pts[idx])


来源:https://stackoverflow.com/questions/27523982/how-to-find-set-of-points-in-x-y-grid-using-kdtree-query-ball-tree

易学教程内所有资源均来自网络或用户发布的内容,如有违反法律规定的内容欢迎反馈
该文章没有解决你所遇到的问题?点击提问,说说你的问题,让更多的人一起探讨吧!