scipy kdtree with meta data

Question

I'm currently looking for a way to build a couple of kd trees for quickly querying some n-dimensional data. However, I'm having some issue with the scipy KD tree algorithm

My data consists of id -> {data: somedata, coordinate: x, y}

I want to be able to query base on the coordinate and k-nearest neighbour's ids as well as getting the fix radius neghbour's id. Judging from the scipy implementation of KDTree and cKDtree, this is not available.

My other options are writing my own KD tree, which will not be that great because I'm just me, or ...?

Answer 1

From playing around with the KDTree, it looks like it insists on having a 2D numpy array like object given to its constructor, but the API returns indices into that array. Further, you can't just monkey patch your data onto the individual coordinates.

But, you can still take advantage of the guts of the KDTree. You need to split your object into a parallel list of the attached data objects, and a numpy array of the coordinates that the KDTree sees. When it say, returns the 3rd point for some search, you know that corresponds to the 3rd index in your attached data object list.

Sure, it's more awkward than an API that would give you a key function that gives coordinates from an arbitrary object, akin to how you can pass key to the python sort function, but it's much better than rolling your own KDTree implementation.