add knninrange to uniformly sample k points in range within a given range

Project.toml

@@ -4,6 +4,7 @@ version = "0.4.24"
 
 [deps]
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [compat]

README.md

@@ -157,11 +157,15 @@ A range search finds all neighbors within the range `r` of given point(s). This
 ```julia
 inrange(tree, point[s], radius) -> idxs
 inrange!(idxs, tree, point, radius)
+knninrange(tree, point[s], radius, k) -> idxs
+knninrange!(idxs, tree, point, radius, k)
+inrangecount(tree, point[s], radius) -> count
 ```
 
 * `tree`: The tree instance.
 * `point[s]`: A vector or matrix of points to find neighbors for.
 * `radius`: Search radius.
+* `k`: Maximum number of neighbors to return when sampling with `knninrange`. Every point inside the radius has equal probability of being selected.
 
 Note: Distances are not returned, only indices.
 
@@ -186,6 +190,11 @@ idxs = inrange(balltree, point, r)
 idxs = Int32[]
 inrange!(idxs, balltree, point, r)
 
+# Sample up to `k` neighbors uniformly at random without allocating new buffers
+buf = zeros(Int, 5)
+nsampled = knninrange!(buf, balltree, point, r, 5)
+random_subset = buf[1:nsampled]
+
 # counts points without allocating index arrays
 neighborscount = inrangecount(balltree, point, r)
 ```

src/NearestNeighbors.jl

@@ -3,10 +3,11 @@ module NearestNeighbors
 using Distances: Distances, PreMetric, Metric, UnionMinkowskiMetric, eval_reduce, eval_end, eval_op, eval_start, evaluate, parameters, Euclidean, Cityblock, Minkowski, Chebyshev, Hamming, Mahalanobis, WeightedEuclidean, WeightedCityblock, WeightedMinkowski
 
 using StaticArrays: StaticArrays, MVector, SVector
+using Random
 using Base: setindex
 
 export NNTree, BruteTree, KDTree, BallTree, DataFreeTree, PeriodicTree
-export knn, knn!, nn, inrange, inrange!, inrangecount, inrange_pairs # TODOs?, npairs
+export knn, knn!, nn, inrange, inrange!, inrangecount, inrange_pairs, knninrange, knninrange! # TODOs?, npairs
 export injectdata
 
 export Euclidean,

src/ball_tree.jl

@@ -210,7 +210,7 @@ end
 function _inrange(tree::BallTree{V},
                   point::AbstractVector,
                   radius::Number,
-                  idx_in_ball::Union{Nothing, Vector{<:Integer}},
+                  idx_in_ball::Union{Nothing, AbstractVector{<:Integer}},
                   skip::F) where {V, F}
     ball = HyperSphere(convert(V, point), convert(eltype(V), radius)) # The "query ball"
     return inrange_kernel!(tree, 1, point, ball, idx_in_ball, skip, nothing) # Call the recursive range finder
@@ -220,7 +220,7 @@ function inrange_kernel!(tree::BallTree,
                          index::Int,
                          point::AbstractVector,
                          query_ball::HyperSphere,
-                         idx_in_ball::Union{Nothing, Vector{<:Integer}},
+                         idx_in_ball::Union{Nothing, AbstractVector{<:Integer}},
                          skip::F,
                          dedup::MaybeBitSet) where {F}
 

src/brute_tree.jl

@@ -78,7 +78,7 @@ end
 function _inrange(tree::BruteTree,
                   point::AbstractVector,
                   radius::Number,
-                  idx_in_ball::Union{Nothing, Vector{<:Integer}},
+                  idx_in_ball::Union{Nothing, AbstractVector{<:Integer}},
                   skip::F,) where {F}
     return inrange_kernel!(tree, point, radius, idx_in_ball, skip, nothing)
 end
@@ -87,7 +87,7 @@ end
 function inrange_kernel!(tree::BruteTree,
                          point::AbstractVector,
                          r::Number,
-                         idx_in_ball::Union{Nothing, Vector{<:Integer}},
+                         idx_in_ball::Union{Nothing, AbstractVector{<:Integer}},
                          skip::Function,
                          dedup::MaybeBitSet)
     count = 0

src/inrange.jl

@@ -1,5 +1,46 @@
 check_radius(r) = r < 0 && throw(ArgumentError("the query radius r must be ≧ 0"))
 
+mutable struct ReservoirSampler{T<:Integer,RNG<:AbstractRNG,V<:AbstractVector{T}} <: AbstractVector{T}
+    storage::V
+    capacity::Int
+    len::Int
+    seen::Int
+    rng::RNG
+end
+
+function ReservoirSampler(storage::AbstractVector{T}, capacity::Integer, rng::AbstractRNG) where {T<:Integer}
+    capacity < 0 && throw(ArgumentError("k must be ≥ 0"))
+    capacity <= length(storage) || throw(ArgumentError("storage length must be ≥ k"))
+    return ReservoirSampler{T, typeof(rng), typeof(storage)}(storage, capacity, 0, 0, rng)
+end
+
+Base.IndexStyle(::Type{<:ReservoirSampler}) = IndexLinear()
+Base.size(rs::ReservoirSampler) = (rs.len,)
+Base.length(rs::ReservoirSampler) = rs.len
+Base.getindex(rs::ReservoirSampler, i::Int) = rs.storage[i]
+Base.setindex!(rs::ReservoirSampler, value, i::Int) = setindex!(rs.storage, value, i)
+
+function Base.push!(rs::ReservoirSampler{T}, value) where {T}
+    rs.seen += 1
+    rs.capacity == 0 && return rs
+    if rs.seen <= rs.capacity
+        rs.len = rs.seen
+        rs.storage[rs.len] = value
+    else
+        j = rand(rs.rng, 1:rs.seen)
+        if j <= rs.capacity
+            rs.storage[j] = value
+        end
+    end
+    return rs
+end
+
+function Base.sort!(rs::ReservoirSampler; kwargs...)
+    rs.len <= 1 && return rs
+    sort!(view(rs.storage, 1:rs.len); kwargs...)
+    return rs
+end
+
 """
     inrange(tree::NNTree, points, radius) -> indices
 
@@ -147,3 +188,69 @@ function inrange_pairs(tree::NNTree, radius::Number, sortres=false, skip::F=Retu
     check_radius(radius)
     return _inrange_pairs(tree, radius, sortres, skip)
 end
+
+"""
+    knninrange(tree::NNTree, point, radius, k; rng=Random.default_rng(), sortres=false, skip=Returns(false))
+
+Return up to `k` indices drawn uniformly at random (without replacement) from the points
+that lie within `radius` of `point`.
+
+This behaves similarly to `inrange`, but it avoids returning more than `k` neighbors.
+
+See also: `knninrange!`.
+"""
+function knninrange(tree::NNTree{V}, point::AbstractVector{T}, radius::Number, k::Integer;
+                    rng::AbstractRNG=Random.default_rng(), sortres=false, skip::F=Returns(false)) where {V, T <: Number, F}
+    check_input(tree, point)
+    check_radius(radius)
+    k < 0 && throw(ArgumentError("k must be ≥ 0"))
+    k == 0 && return Int[]
+    buf = Vector{Int}(undef, k)
+    nsampled = knninrange!(buf, tree, point, radius, k; rng=rng, sortres=sortres, skip=skip)
+    resize!(buf, nsampled)
+    return buf
+end
+
+function knninrange(tree::NNTree{V}, points::AbstractVector{T}, radius::Number, k::Integer;
+                    rng::AbstractRNG=Random.default_rng(), sortres=false, skip::F=Returns(false)) where {V, T <: AbstractVector, F}
+    check_input(tree, points)
+    check_radius(radius)
+    return [knninrange(tree, points[i], radius, k; rng=rng, sortres=sortres, skip=skip) for i in 1:length(points)]
+end
+
+function knninrange(tree::NNTree{V}, points::AbstractMatrix{T}, radius::Number, k::Integer;
+                    rng::AbstractRNG=Random.default_rng(), sortres=false, skip::F=Returns(false)) where {V, T <: Number, F}
+    check_input(tree, points)
+    check_radius(radius)
+    dim = size(points, 1)
+    n_points = size(points, 2)
+    idxs = Vector{Vector{Int}}(undef, n_points)
+    for i in 1:n_points
+        point = SVector{dim,T}(ntuple(j -> points[j, i], Val(dim)))
+        idxs[i] = knninrange(tree, point, radius, k; rng=rng, sortres=sortres, skip=skip)
+    end
+    return idxs
+end
+
+"""
+    knninrange!(idxs, tree, point, radius, k; rng=Random.default_rng(), sortres=false, skip=Returns(false))
+
+Mutating version of `knninrange`. The first `k` entries of `idxs` are used as storage for the
+reservoir sampler and will contain the sampled indices after the call returns. The function returns
+the number of valid samples that were written (i.e. `min(k, number_in_range)`).
+
+The length of `idxs` must be at least `k`. The contents beyond the returned sample length are left
+untouched.
+"""
+function knninrange!(idxs::AbstractVector{<:Integer}, tree::NNTree{V}, point::AbstractVector{T},
+                     radius::Number, k::Integer=length(idxs); rng::AbstractRNG=Random.default_rng(),
+                     sortres=false, skip::F=Returns(false)) where {V, T <: Number, F}
+    check_input(tree, point)
+    check_radius(radius)
+    k < 0 && throw(ArgumentError("k must be ≥ 0"))
+    k == 0 && return 0
+    k <= length(idxs) || throw(ArgumentError("idxs must have length ≥ k"))
+    sampler = ReservoirSampler(idxs, k, rng)
+    _inrange_point!(tree, point, radius, sortres, sampler, skip)
+    return length(sampler)
+end

src/kd_tree.jl

@@ -222,7 +222,7 @@ function _inrange(
         tree::KDTree,
         point::AbstractVector,
         radius::Number,
-        idx_in_ball::Union{Nothing, Vector{<:Integer}},
+        idx_in_ball::Union{Nothing, AbstractVector{<:Integer}},
         skip::F) where {F}
     init_min = get_min_distance_no_end(tree.metric, tree.hyper_rec, point)
     init_max_contribs = get_max_distance_contributions(tree.metric, tree.hyper_rec, point)
@@ -239,7 +239,7 @@ function inrange_kernel!(
         index::Int,
         point::AbstractVector,
         r::Number,
-        idx_in_ball::Union{Nothing, Vector{<:Integer}},
+        idx_in_ball::Union{Nothing, AbstractVector{<:Integer}},
         hyper_rec::HyperRectangle,
         min_dist,
         max_dist_contribs::SVector,

src/periodic_tree.jl

@@ -233,7 +233,7 @@ end
 function _inrange(tree::PeriodicTree{V},
     point::AbstractVector,
     radius::Number,
-    idx_in_ball::Union{Nothing, Vector{<:Integer}},
+    idx_in_ball::Union{Nothing, AbstractVector{<:Integer}},
     skip::F) where {V, F}
 
     dedup_state = empty!(tree.dedup_set)

src/tree_ops.jl

@@ -180,7 +180,7 @@ end
 
 # Add all points in this subtree since we have determined
 # they are all within the desired range
-function addall(tree::NNTree, index::Int, idx_in_ball::Union{Nothing, Vector{<:Integer}}, skip::Function,
+function addall(tree::NNTree, index::Int, idx_in_ball::Union{Nothing, AbstractVector{<:Integer}}, skip::Function,
                 dedup::MaybeBitSet)
     tree_data = tree.tree_data
     if isleaf(tree_data.n_internal_nodes, index)

test/test_inrange.jl

@@ -1,3 +1,5 @@
+using StableRNGs
+
 # Does not test leafsize
 @testset "inrange" begin
     @testset "metric" for metric in [Euclidean()]
@@ -84,6 +86,54 @@ end
     @test idxs isa Vector{Vector{Int}}
 end
 
+@testset "knninrange" begin
+    points = [
+        0.0 0.0 0.0 0.0 1.0 1.0 1.0 1.0 0.5 0.5;
+        0.0 0.0 1.0 1.0 0.0 0.0 1.0 1.0 0.5 0.5;
+        0.0 1.0 0.0 1.0 0.0 1.0 0.0 1.0 0.5 0.2
+    ]
+    target = SVector{3,Float64}(0.5, 0.5, 0.5)
+    target_far = SVector{3,Float64}(2.0, 2.0, 2.0)
+    radius = 2.0
+    small_radius = 0.6
+    trees = (BruteTree, KDTree, BallTree)
+    idx_range = 1:size(points, 2)
+
+    for T in trees
+        reorder_vals = T === BruteTree ? (false,) : (false, true)
+        for reorder in reorder_vals
+            tree = T(points; leafsize=2, reorder=reorder)
+
+            sample = knninrange(tree, target, radius, 3; rng=StableRNG(42))
+            @test length(sample) == 3
+            @test all(in(idx_range), sample)
+            @test length(unique(sample)) == 3
+
+            buf = fill(-1, 5)
+            len = knninrange!(buf, tree, target, radius, 4; rng=StableRNG(7))
+            @test len == min(4, length(inrange(tree, target, radius)))
+            @test all(in(idx_range), buf[1:len])
+            @test length(unique(buf[1:len])) == len
+
+            tight = knninrange(tree, target_far, 0.05, 5)
+            @test tight == Int[]  # no points within the radius
+
+            near = knninrange(tree, target, small_radius, 5)
+            expected = sort(inrange(tree, target, small_radius))
+            @test sort(near) == expected
+
+            @test knninrange(tree, target, radius, 0) == Int[]
+            @test knninrange!(Int[], tree, target, radius, 0) == 0
+        end
+    end
+
+    tree = KDTree(points)
+    multi = knninrange(tree, [target, target], radius, 2; rng=StableRNG(3))
+    @test multi isa Vector{Vector{Int}}
+    @test all(length(idx) <= 2 for idx in multi)
+    @test all(all(in(idx_range), idx) for idx in multi)
+end
+
 @testset "mutating" begin
     for T in (KDTree, BallTree, BruteTree)
         data = T(rand(3, 100))