Distributed KDTree using Dask

dkdtree.py

from functools import partial
import numpy as np
from scipy.spatial import cKDTree
from dask.delayed import delayed
import dask.bag as db

class DaskKDTree(object):
    """
    
    Usage Example:
    --------------
    
    from dask.distributed import Client
    client = Client('52.91.203.58:8786')
    tree = DaskKDTree(client, leafsize=1000)
    tree.load_random(num_points=int(1e8), chunk_size=int(1e6))
    
    # find all points within 10km of Bourbon Street
    bourbon_street = (-10026195.958134, 3498018.476606)
    radius = 10000 # meters
    result = tree.query_ball_point(x=bourbon_street, r=radius)
    """
    
    
    def __init__(self, client, leafsize):
        self.client = client
        self.leafsize = leafsize
        self.trees = []
    
    
    def load_random(self, num_points=int(1e6), chunk_size=300):
        parts = int(num_points / chunk_size)
        self.trees = [delayed(DaskKDTree._run_load_random)(int(chunk_size), leafsize=self.leafsize) for f in range(parts)]
        self.trees = self.client.persist(self.trees)
        
        
    @staticmethod
    def _run_load_random(count, leafsize):
        xs = np.random.uniform(int(-20e6), int(20e6), count)
        ys = np.random.uniform(int(-20e6), int(20e6), count)
        points = np.dstack((xs, ys))[0, :]
        return cKDTree(points, leafsize=leafsize)    
    
    
    def query_ball_point(self, **kwargs):
        nearest = [delayed(DaskKDTree._run_query_ball_point)(d, kwargs) for d in self.trees]
        b = db.from_delayed(nearest)
        return b.compute()
    
    
    @staticmethod
    def _run_query_ball_point(tree, query_info):
        indices = tree.query_ball_point(**query_info)
        return tree.data[indices]    

Facing this error, while replicating your example:
File "C:\Anaconda3\lib\site-packages\dask\bag\core.py", line 1460, in reify
if seq and isinstance(seq[0], Iterator):

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

For small number of kdtrees, this is an interesting approach.

For truly distributed KDTrees, I recently read a paper "Highly Parallel Fast KD‐tree Construction for Interactive Ray Tracing of Dynamic Scenes" [1] that looks like it might be a more optimal way to do this in parallel.

[1]: https://onlinelibrary.wiley.com/doi/pdf/10.1111/j.1467-8659.2007.01062.x