Some context
Suppose a binary space partitioning on the unit square (like the cells of a k-d tree), represented as a numpy array of shape (n,4), where n denotes the number of (axis aligned, non-overlapping) bounding boxes, and each bbox is represented as [min_x, max_x, min_y, max_y].
Suppose further a set of p samples in the unit square, as an array of shape (p,2). More precisely, the samples are actually drawn from the (2,3-) halton sequence (though this isn't really all too relevant).
The problem
Now I would like to find, for each of those p samples an index into the array of bounding boxes along axis zero, which indicates the bbox inside of which the given sample falls. So I am trying to perform an orthogonal range search to find which of the bboxes encloses each given sample point.
A possible solution
My idea was to calculate the centroid of each bbox, make a k-d tree from the resulting points and query the k-d tree for the nearest neighbour (among those) for any of the p samples.
Using a k-d tree seems like a straightforward solution, yet trying to make use of scipy.spacial.KDTree for this gives me erroneous results (I am probably just using it the wrong way somehow, though).
What goes wrong
Here's a simple example to illustrate:
import numpy as np
from scipy.spatial import KDTree
# bbox coordinates as (min_x, max_x, min_y, max_y)
bbox = np.array([[0.75 1. 0.44 1. ]
[0. 0.1097561 0. 0.54666666]
[0.1097561 0.75 0. 0.54666666]
[0. 0.24264705 0.54666666 1. ]
[0.24264705 0.75 0.54666666 1. ]
[0.75 0.84090909 0. 0.44 ]
[0.84090909 1. 0. 0.44 ]])
# sample points
samples = np.array([[0. , 0. ],
[0.5 , 0.33333333],
[0.25 , 0.66666667],
[0.75 , 0.11111111],
[0.125 , 0.44444444]])
def find_bbox(bounding_boxes, samples):
'''docstring'''
# Calculate the centroid of each bounding box as (centroid_x, centroid_y).
centroids = np.vstack((np.mean(bounding_boxes[:,0:2], axis=1),
np.mean(bounding_boxes[:,2:4], axis=1))).T
# Build a scipy.spacial.KDTree from the bounding box centroids.
tree = KDTree(centroids, leafsize=1)
# For each XY-coordinate in (samples), find the nearest neighbour among
# (centroids).
distances, indices = tree.query(samples, k=1)
return indices
# Expected result would be: [1, 2, 4, 2, 2] or [1, 2, 4, 5, 2], instead I get
# [1, 2, 3, 5, 1]
print(find_bbox(bbox, samples))
Trying any combination of the optional keyword arguments to scipy.spatial.KDTree didn't help neither. Any hints would be much appreciated.
edit: Come to think about it, it must be a distance-metric problem, right? scipy.spatial.KDTree defaults to euclidean distance, and considering this, the results I get do make sense.
This doesn't seem to be an especially good fit for KDTree. I think this is theoretically possible, but it would require you to modify internal implementation details of KDTree. This would mean that your code might break on updating to new versions of SciPy.
Instead, I would suggest that you use an R Tree, which is a tree of bounding boxes. The advantage of this approach is that you still get O(log N) query times, but you avoid relying on internal implementation details that could change in a future version of SciPy.
SciPy does not implement an R Tree, but another library, rtree does.
Here is an example:
from rtree import index
import numpy as np
# bbox coordinates as (min_x, max_x, min_y, max_y)
bbox = np.array([
[0.75, 1., 0.44, 1. ],
[0., 0.1097561, 0., 0.54666666 ],
[0.1097561, 0.75, 0., 0.54666666 ],
[0., 0.24264705, 0.54666666, 1. ],
[0.24264705, 0.75, 0.54666666, 1. ],
[0.75, 0.84090909, 0., 0.44 ],
[0.84090909, 1., 0., 0.44 ],
])
# sample points
samples = np.array([[0. , 0. ],
[0.5 , 0.33333333],
[0.25 , 0.66666667],
[0.75 , 0.11111111],
[0.125 , 0.44444444]])
# interleaved is required because coordinates are in order (min_x, max_x, min_y, max_y)
idx = index.Index(interleaved=False)
for i, coord in enumerate(bbox):
idx.insert(i, coord)
indices = []
for sample in samples:
indices.append(list(idx.intersection(sample)))
print(indices)
Output:
[[1], [2], [4], [2, 5], [2]]
This shows that it can find both of the answers you find acceptable for the fourth search. If you don't care which one you find, you could modify it to stop at the first match:
for sample in samples:
indices.append(next(idx.intersection(sample)))
Output:
[1, 2, 4, 2, 2]