kdtree.py

#!/usr/bin/env python

"""
  kdtree.py
"""

import numpy as np
import pandas as pd
from copy import copy
from itertools import chain
from scipy.spatial import KDTree

def get_boxes(node, bbox):
    bbox_less = copy(bbox)
    bbox_greater = copy(bbox)
    if isinstance(node, KDTree.innernode):
        if node.split_dim == 0:
            bbox_less['top'] = node.split
            bbox_greater['bottom'] = node.split
        else:
            bbox_less['right'] = node.split
            bbox_greater['left'] = node.split
        
        less_gen = get_boxes(node.less, bbox_less)
        greater_gen = get_boxes(node.greater, bbox_greater)
        
        return less_gen + greater_gen
    else:
        return [bbox]

# --
# Run

x = np.random.normal(0, 1, (10000, 2))

bbox = {
    "top"    : x[:,0].max(),
    "bottom" : x[:,1].min(),
    "left"   : x[:,0].min(),
    "right"  : x[:,0].max()
}

tree = KDTree(x, leafsize=100)
boxes = pd.DataFrame(get_boxes(tree.tree, bbox))

# Check that it worked
for _ in range(100):
    p = np.random.normal(0, 1, 2)
    
    sel = (
        (boxes.top > p[1]) & 
        (boxes.bottom < p[1]) & 
        (boxes.left < p[0]) & 
        (boxes.right > p[0])
    )
    
    assert sel.sum() == 1