Interface for tree traversal / walking of BallTree/KDTree

[dgleich]

It'd be great to have an interface to walk the tree structures. This could be used internally by some of the search options and externally as in a few cases I have below.

see https://github.com/dgleich/GraphPlayground.jl/blob/main/src/manybodyforce.jl for usage (which just broke given one of your changes :) ... which is what I should expect for using undocumented internals!) and https://discourse.julialang.org/t/packages-for-quadtrees-in-julia/113078/3

So I'm trying to get a more useful walk interface...

I have a few ideas, but I'm trying to solicit use cases at the moment.

• my use cases above
• inrange_kernel in NearestNeighbors.

Are there others? The plan is to submit a pull request at some point with an initial take on how to make this generally useful.

[KristofferC]

Yes, it's a good idea. It's in #180

[KristofferC]

It would be good here to list the set of methods required and what their return values should be.

[dgleich]

The initial idea was something like this...

SpaceTree::Union{KDTree,BallTree}

struct KDTreeNode{T,R}
  index::Int
  tree::T
  region::R
end 

index(node::KDTreeNode) = node.index 
points(node::KDTreeNode) = [ generator for all points contained within the region described by the node ]
points_indices(node::KDTreeNode) = [ generator for all indices of points in the range ] # handles cases where you have metadata
region(node::KDTreeNode) = node.region 
isleaf(node::KDTreeNode) = isleaf(node.tree.tree_data.n_internal_nodes, node.index)

children(node::KDTreeNode) = union{Tuple{KDTreeNode,KDTreeNode},Nothing}

treeroot(tree::SpaceTree) = KDTreeNode(tree, 1, T.hyper_rec)
# or just root?
root(tree::SpaceTree) = KDTreeNode(tree, 1, T.hyper_rec)

I think this is a super easy to use interface.

Things I don't like about this:

• region computations are fast, but not free. I have situations where they aren't required so it would be good to elide them.
• "tree" is copied to each node.

This is where I'd start and then optimize from there to see if the compiler is smart enough to elide the region info if it isn't used in a function. (Maybe with aggressive use of @inline? )