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

Question

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?

Octave

I dont' understand, if you have a symmetric/uniform mesh, then finding the indices is more or less trivial, since the element index is related to its position. A Kd-tree is totally overkill in this scenario. – Ettie

No, its not a symmetric mesh. the distance between the y points and the x points are not the same. Y distance is 0.5 and X is 3.4. The above grid was an example. the actual grid looks somewhat like this.. [(-150,-150),(-149.5,-146.6),(-149,-143.2),(-148.5,-139.8).....]... So, what we want to find here is for any point say (-149,-143.2), we want to find points which are within a radius of 2. – Octave

If the X and Y distance is the same, and the points are ordered, then you still don't need a Kd-tree. And you have not specified where (in which module) this KDTree.query_ball_tree function is defined, which makes it hard to answer the question. – Ettie

Answer 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.

Barbet

Answer 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])

Annapolis