Runtime function

README.md

@@ -164,6 +164,41 @@ inrange!(idxs, balltree, point, r)
 neighborscount = inrangecount(balltree, point, r)
 ```
 
+### Passing a runtime function into the range search
+```julia
+inrange_runtime!(tree, points, radius, runtime_function)
+```
+
+Example:
+```julia
+using NearestNeighbors
+data = rand(3,10^4)
+data_values = rand(10^4)
+r = 0.05
+points = rand(3,10)
+results = zeros(10)
+
+# this function will sum the `data_values` corresponding to the `data` that is in range of `points`
+# `p_idx` is the index of `points` i.e. 1-10
+# `data_idx` is is the index of the data in the tree that is in range
+# `p` is `points[:,p_idx]`
+# `values` is data needed for the operation
+# `results` is a storage space for the results
+function sum_values!(p_idx, data_idx, p, values, results)
+    results[p_idx] += values[data_idx]
+end
+
+# `runtime_function` must be of the form f(p_idx, data_idx, p)
+runtime_function(p_idx, data_idx, p) = sum_values!(p_idx, data_idx, p, values, results)
+
+kdtree = KDTree(data)
+
+# runs the runtime_function with all tree data points in range of points. In this case sums the `data_values` corresponding to the `data` that is in range of `points`
+inrange_runtime!(tree, points, radius, runtime_function)
+```
+
+
+
 ## Using On-Disk Data Sets
 
 By default, trees store a copy of the `data` provided during construction. For data sets larger than available memory, `DataFreeTree` can be used to strip a tree of its data field and re-link it later.

src/NearestNeighbors.jl

@@ -7,7 +7,7 @@ using StaticArrays
 import Base.show
 
 export NNTree, BruteTree, KDTree, BallTree, DataFreeTree
-export knn, knn!, nn, inrange, inrange!,inrangecount # TODOs? , allpairs, distmat, npairs
+export knn, knn!, nn, inrange, inrange!,inrangecount, inrange_runtime! # TODOs? , allpairs, distmat, npairs
 export injectdata
 
 export Euclidean,

src/ball_tree.jl

@@ -177,16 +177,20 @@ end
 function _inrange(tree::BallTree{V},
                   point::AbstractVector,
                   radius::Number,
-                  idx_in_ball::Union{Nothing, Vector{<:Integer}}) where {V}
+                  idx_in_ball::Union{Nothing, Vector{<:Integer}},
+                  point_index::Int = 1,
+                  runtime_function::Union{Nothing, Function} = nothing) where {V}
     ball = HyperSphere(convert(V, point), convert(eltype(V), radius)) # The "query ball"
-    return inrange_kernel!(tree, 1, point, ball, idx_in_ball) # Call the recursive range finder
+    return inrange_kernel!(tree, 1, point, ball, idx_in_ball, runtime_function, point_index) # Call the recursive range finder
 end
 
 function inrange_kernel!(tree::BallTree,
                          index::Int,
                          point::AbstractVector,
                          query_ball::HyperSphere,
-                         idx_in_ball::Union{Nothing, Vector{<:Integer}})
+                         idx_in_ball::Union{Nothing, Vector{<:Integer}},
+                         runtime_function::Union{Nothing, Function},
+                         point_index::Int)
 
     if index > length(tree.hyper_spheres)
         return 0
@@ -204,7 +208,7 @@ function inrange_kernel!(tree::BallTree,
     # At a leaf node, check all points in the leaf node
     if isleaf(tree.tree_data.n_internal_nodes, index)
         r = tree.metric isa MinkowskiMetric ? eval_pow(tree.metric, query_ball.r) : query_ball.r
-        return add_points_inrange!(idx_in_ball, tree, index, point, r)
+        return add_points_inrange!(idx_in_ball, tree, index, point, r, runtime_function, point_index)
     end
 
     count = 0
@@ -215,8 +219,8 @@ function inrange_kernel!(tree::BallTree,
         count += addall(tree, index, idx_in_ball)
     else
         # Recursively call the left and right sub tree.
-        count += inrange_kernel!(tree,  getleft(index), point, query_ball, idx_in_ball)
-        count += inrange_kernel!(tree, getright(index), point, query_ball, idx_in_ball)
+        count += inrange_kernel!(tree,  getleft(index), point, query_ball, idx_in_ball, runtime_function, point_index)
+        count += inrange_kernel!(tree, getright(index), point, query_ball, idx_in_ball, runtime_function, point_index)
     end
     return count
 end

src/brute_tree.jl

@@ -61,21 +61,26 @@ end
 function _inrange(tree::BruteTree,
                   point::AbstractVector,
                   radius::Number,
-                  idx_in_ball::Union{Nothing, Vector{<:Integer}})
-    return inrange_kernel!(tree, point, radius, idx_in_ball)
+                  idx_in_ball::Union{Nothing, Vector{<:Integer}},
+                  point_index::Int = 1,
+                  runtime_function::Union{Nothing, Function} = nothing)
+    return inrange_kernel!(tree, point, radius, idx_in_ball, runtime_function, point_index)
 end
 
 
 function inrange_kernel!(tree::BruteTree,
                          point::AbstractVector,
                          r::Number,
-                         idx_in_ball::Union{Nothing, Vector{<:Integer}})
+                         idx_in_ball::Union{Nothing, Vector{<:Integer}},
+                         runtime_function::Union{Nothing, Function},
+                         point_index::Int)
     count = 0
     for i in 1:length(tree.data)
         d = evaluate(tree.metric, tree.data[i], point)
         if d <= r
             count += 1
             idx_in_ball !== nothing && push!(idx_in_ball, i)
+            !isnothing(runtime_function) && runtime_function(point_index, i, point)
         end
     end
     return count

src/inrange.jl

@@ -104,3 +104,60 @@ function inrangecount(tree::NNTree{V}, point::AbstractMatrix{T}, radius::Number)
     end
     return inrangecount(tree, new_data, radius)
 end
+
+"""
+    inrange_runtime!(tree::NNTree{V}, points::AbstractVector{T}, radius::Number, runtime_function::Function)
+
+Compute a runtime function for all in range queries.
+Instead of returning the indicies, the `runtime_function` is called for each point in points
+and is given the points, the index of the point, and the index of the neighbor.
+The `runtime_function` should return nothing.
+The `runtime_function` should be of the form:
+runtime_function(point_index::Int, neighbor_index::Int, point::AbstractVector{T})
+where `point_index` is the index of the point in `points`, `neighbor_index` is the index of the neighbor in the tree,
+and `point` is the point in points.
+
+The `runtime_function` should not modify the tree or the points.
+
+Ananymous functions can be used as well.
+For example:
+```julia
+function runtime_function(point_index, neighbor_index, point, random_storage_of_results, neightbors_data)
+    # do something with the points
+    return nothing
+end
+
+random_storage_of_results = rand(3, 100)
+neightbors_data = rand(3, 100)
+f(a, b, c) = runtime_function(a, b, c, random_storage_of_results, neightbors_data)
+```
+"""
+function inrange_runtime!(tree::NNTree{V}, points::AbstractVector{T}, radius::Number, runtime_function::F) where {V, T <: AbstractVector, F}
+    check_input(tree, points)
+    check_radius(radius)
+
+    for i in eachindex(points)
+        _inrange(tree, points[i], radius, nothing, i, runtime_function)
+    end
+    return nothing
+end
+
+function inrange_runtime!(tree::NNTree{V}, points::AbstractMatrix{T}, radius::Number, runtime_function::F) where {V, T <: Number, F}
+    return inrange_runtime!(tree, points, radius, runtime_function, Val(size(points, 1)))
+end
+
+function inrange_runtime!(tree::NNTree{V}, points::AbstractVector{T}, radius::Number, runtime_function::F) where {V, T <: Number, F}
+    points = reshape(points, size(points, 1), 1)
+    return inrange_runtime!(tree, points, radius, runtime_function, Val(size(points, 1)))
+end
+
+function inrange_runtime!(tree::NNTree{V}, points::AbstractMatrix{T}, radius::Number, runtime_function::F, ::Val{dim}) where {V, T <: Number, F, dim}
+    check_input(tree, points)
+    check_radius(radius)
+    n_points = size(points, 2)
+    for i in 1:n_points
+        point = SVector{dim,T}(ntuple(j -> points[j, i], Val(dim)))
+        _inrange(tree, point, radius, nothing, i, runtime_function)
+    end
+    return nothing
+end

src/kd_tree.jl

@@ -207,10 +207,12 @@ end
 function _inrange(tree::KDTree,
                   point::AbstractVector,
                   radius::Number,
-                  idx_in_ball::Union{Nothing, Vector{<:Integer}} = Int[])
+                  idx_in_ball::Union{Nothing, Vector{<:Integer}} = Int[],
+                  point_index::Int = 1,
+                  runtime_function::Union{Nothing, Function} = nothing)
     init_min = get_min_distance_no_end(tree.metric, tree.hyper_rec, point)
     return inrange_kernel!(tree, 1, point, eval_pow(tree.metric, radius), idx_in_ball,
-            tree.hyper_rec, init_min)
+            tree.hyper_rec, init_min, runtime_function, point_index)
 end
 
 # Explicitly check the distance between leaf node and point while traversing
@@ -220,15 +222,17 @@ function inrange_kernel!(tree::KDTree,
                          r::Number,
                          idx_in_ball::Union{Nothing, Vector{<:Integer}},
                          hyper_rec::HyperRectangle,
-                         min_dist)
+                         min_dist,
+                         runtime_function::Union{Nothing, Function},
+                         point_index::Int)
     # Point is outside hyper rectangle, skip the whole sub tree
     if min_dist > r
         return 0
     end
 
     # At a leaf node. Go through all points in node and add those in range
     if isleaf(tree.tree_data.n_internal_nodes, index)
-        return add_points_inrange!(idx_in_ball, tree, index, point, r)
+        return add_points_inrange!(idx_in_ball, tree, index, point, r, runtime_function, point_index)
     end
 
     split_val = tree.split_vals[index]
@@ -255,7 +259,7 @@ function inrange_kernel!(tree::KDTree,
         ddiff = max(zero(lo - p_dim), lo - p_dim)
     end
     # Call closer sub tree
-    count += inrange_kernel!(tree, close, point, r, idx_in_ball, hyper_rec_close, min_dist)
+    count += inrange_kernel!(tree, close, point, r, idx_in_ball, hyper_rec_close, min_dist, runtime_function, point_index)
 
     # TODO: We could potentially also keep track of the max distance
     # between the point and the hyper rectangle and add the whole sub tree
@@ -267,6 +271,6 @@ function inrange_kernel!(tree::KDTree,
     ddiff_pow = eval_pow(M, ddiff)
     diff_tot = eval_diff(M, split_diff_pow, ddiff_pow, split_dim)
     new_min = eval_reduce(M, min_dist, diff_tot)
-    count += inrange_kernel!(tree, far, point, r, idx_in_ball, hyper_rec_far, new_min)
+    count += inrange_kernel!(tree, far, point, r, idx_in_ball, hyper_rec_far, new_min, runtime_function, point_index)
     return count
 end

src/tree_ops.jl

@@ -115,13 +115,16 @@ end
 # This will probably prevent SIMD and other optimizations so some care is needed
 # to evaluate if it is worth it.
 @inline function add_points_inrange!(idx_in_ball::Union{Nothing, AbstractVector{<:Integer}}, tree::NNTree,
-                                     index::Int, point::AbstractVector, r::Number)
+                                     index::Int, point::AbstractVector, r::Number, 
+                                     runtime_function::Union{Nothing, Function},
+                                     point_index::Int)
     count = 0
-    for z in get_leaf_range(tree.tree_data, index)
+    @inbounds for z in get_leaf_range(tree.tree_data, index)
         idx = tree.reordered ? z : tree.indices[z]
         if check_in_range(tree.metric, tree.data[idx], point, r)
             count += 1
             idx_in_ball !== nothing && push!(idx_in_ball, idx)
+            @inbounds !isnothing(runtime_function) && runtime_function(point_index, tree.reordered ? tree.indices[idx] : idx, point)
         end
     end
     return count

test/test_inrange.jl

@@ -93,3 +93,26 @@ end
         @test idxs == idxs2
     end
 end
+
+@testset "inrange_runtime function" begin
+    function runtime_test(point_index, neighbor_index, point, sum_of_random_data, neightbor_points)
+        sum_of_random_data[1] += sum(neightbor_points[4:6,neighbor_index])
+        return nothing
+    end
+
+    for T in (KDTree, BallTree, BruteTree)
+        sum_runtime = fill(0.0, 1)
+        data = rand(6, 100) # first 3 rows are "locations", last 3 rows are random data
+        f(a, b, c) = runtime_test(a, b, c, sum_runtime, data)
+
+        tree = KDTree(data[1:3, :])
+        inrange_runtime!(tree, [0.5, 0.5, 0.5], 1.0, f)
+        idxs = inrange(tree, [0.5, 0.5, 0.5], 1.0)
+        sum_idxs = 0.0
+        for i in eachindex(idxs)
+            sum_idxs += sum(data[4:6, idxs[i]])
+        end
+
+        @test sum_idxs == sum_runtime[1]
+    end
+end