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?

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.

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