8

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?

Hooked
  • 84,485
  • 43
  • 192
  • 261
Numby
  • 125
  • 1
  • 1
  • 8
  • 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. – Hannes Ovrén Dec 17 '14 at 12:50
  • 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. – Numby Dec 17 '14 at 13:04
  • 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. – Hannes Ovrén Dec 17 '14 at 14:26

2 Answers2

14

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.

Hooked
  • 84,485
  • 43
  • 192
  • 261
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])
U3.1415926
  • 812
  • 12
  • 30