add skip_self for inrange and knn

Author: KristofferC

README.md

@@ -78,6 +78,7 @@ knn!(idxs, dists, tree, point, k [, skip=Returns(false)])
 * `point[s]`: A vector or matrix of points to find the `k` nearest neighbors for. A vector of numbers represents a single point; a matrix means the `k` nearest neighbors for each point (column) will be computed. `points` can also be a vector of vectors.
 * `k`: Number of nearest neighbors to find.
 * `skip` (optional): A predicate function to skip certain points, e.g., points already visited.
+* `skip_self` (optional, batched queries only): Skip the point with the same index as the current query when the query set is identical to the tree data, e.g. `knn(tree, data, 1; skip_self=true)`.
 
 
 For the single closest neighbor, you can use `nn`:
@@ -145,6 +146,12 @@ dists
 #  0.04556078331418939
 #  0.049967238112417205
 
+# Self-query the same dataset without returning each point as its own neighbor
+idxs, dists = knn(kdtree, data, 1; skip_self=true)
+
+# Retrieve just the nearest neighbor per point
+nn_idx, nn_dist = nn(kdtree, data; skip_self=true)
+
 # Preallocating input results
 idxs, dists = zeros(Int32, k), zeros(Float32, k)
 knn!(idxs, dists, kdtree, v, k)
@@ -162,6 +169,8 @@ inrange!(idxs, tree, point, radius)
 * `tree`: The tree instance.
 * `point[s]`: A vector or matrix of points to find neighbors for.
 * `radius`: Search radius.
+* `skip` (optional): Predicate to skip certain points.
+* `skip_self` (optional, batched queries only): When querying the same dataset, skip the point whose index matches the query.
 
 Note: Distances are not returned, only indices.
 
@@ -188,6 +197,10 @@ inrange!(idxs, balltree, point, r)
 
 # counts points without allocating index arrays
 neighborscount = inrangecount(balltree, point, r)
+
+# Self-query without returning each point itself
+idxs_self = inrange(balltree, data, r; skip_self=true)
+counts_self = inrangecount(balltree, data, r; skip_self=true)
 ```
 
 ### Self-Pair Searches

src/ball_tree.jl

@@ -167,8 +167,9 @@ function _knn(tree::BallTree,
               best_idxs::Union{Integer, AbstractVector{<:Integer}},
               best_dists::Union{Number, AbstractVector},
               ::Union{Nothing, AbstractVector},
-              skip::F) where {F}
-    return knn_kernel!(tree, 1, point, best_idxs, best_dists, skip, nothing)
+              skip::F,
+              self_idx::Int) where {F}
+    return knn_kernel!(tree, 1, point, best_idxs, best_dists, skip, nothing, self_idx)
 end
 
 
@@ -178,9 +179,10 @@ function knn_kernel!(tree::BallTree{V},
                      best_idxs::Union{Integer, AbstractVector{<:Integer}},
                      best_dists::Union{Number, AbstractVector},
                      skip::F,
-                     dedup::MaybeBitSet) where {V, F}
+                     dedup::MaybeBitSet,
+                     self_idx::Int) where {V, F}
     if isleaf(tree.tree_data.n_internal_nodes, index)
-        return add_points_knn!(best_dists, best_idxs, tree, index, point, true, skip, dedup)
+        return add_points_knn!(best_dists, best_idxs, tree, index, point, true, skip, dedup, self_idx)
     end
 
     left_sphere = tree.hyper_spheres[getleft(index)]
@@ -192,16 +194,16 @@ function knn_kernel!(tree::BallTree{V},
     best_dist_1 = first(best_dists)
     if left_dist <= best_dist_1 || right_dist <= best_dist_1
         if left_dist < right_dist
-            best_idxs, best_dists = knn_kernel!(tree, getleft(index), point, best_idxs, best_dists, skip, dedup)
+            best_idxs, best_dists = knn_kernel!(tree, getleft(index), point, best_idxs, best_dists, skip, dedup, self_idx)
             best_dist_1 = first(best_dists)
             if right_dist <= best_dist_1
-                best_idxs, best_dists = knn_kernel!(tree, getright(index), point, best_idxs, best_dists, skip, dedup)
+                best_idxs, best_dists = knn_kernel!(tree, getright(index), point, best_idxs, best_dists, skip, dedup, self_idx)
             end
         else
-            best_idxs, best_dists = knn_kernel!(tree, getright(index), point, best_idxs, best_dists, skip, dedup)
+            best_idxs, best_dists = knn_kernel!(tree, getright(index), point, best_idxs, best_dists, skip, dedup, self_idx)
             best_dist_1 = first(best_dists)
             if left_dist <= best_dist_1
-                best_idxs, best_dists = knn_kernel!(tree, getleft(index), point, best_idxs, best_dists, skip, dedup)
+                best_idxs, best_dists = knn_kernel!(tree, getleft(index), point, best_idxs, best_dists, skip, dedup, self_idx)
             end
         end
     end
@@ -212,9 +214,10 @@ function _inrange(tree::BallTree{V},
                   point::AbstractVector,
                   radius::Number,
                   idx_in_ball::Union{Nothing, Vector{<:Integer}},
-                  skip::F) where {V, F}
+                  skip::F,
+                  self_idx::Int) 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
+    return inrange_kernel!(tree, 1, point, ball, idx_in_ball, skip, nothing, self_idx) # Call the recursive range finder
 end
 
 function inrange_kernel!(tree::BallTree,
@@ -223,7 +226,8 @@ function inrange_kernel!(tree::BallTree,
                          query_ball::HyperSphere,
                          idx_in_ball::Union{Nothing, Vector{<:Integer}},
                          skip::F,
-                         dedup::MaybeBitSet) where {F}
+                         dedup::MaybeBitSet,
+                         self_idx::Int) where {F}
 
     if index > length(tree.hyper_spheres)
         return 0
@@ -241,17 +245,17 @@ 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, skip, dedup)
+        return add_points_inrange!(idx_in_ball, tree, index, point, r, skip, dedup, self_idx)
     end
 
     # The query ball encloses the sub tree bounding sphere. Add all points in the
     # sub tree without checking the distance function.
     if encloses_fast(dist, tree.metric, sphere, query_ball)
-        return addall(tree, index, idx_in_ball, skip, dedup)
+        return addall(tree, index, idx_in_ball, skip, dedup, self_idx)
     else
         # Recursively call the left and right sub tree.
-        return inrange_kernel!(tree,  getleft(index), point, query_ball, idx_in_ball, skip, dedup) +
-               inrange_kernel!(tree, getright(index), point, query_ball, idx_in_ball, skip, dedup)
+        return inrange_kernel!(tree,  getleft(index), point, query_ball, idx_in_ball, skip, dedup, self_idx) +
+               inrange_kernel!(tree, getright(index), point, query_ball, idx_in_ball, skip, dedup, self_idx)
     end
 end
 

src/brute_tree.jl

@@ -46,20 +46,22 @@ function _knn(tree::BruteTree{V},
                  best_idxs::Union{Integer, AbstractVector{<:Integer}},
                  best_dists::Union{Number, AbstractVector},
                  ::Union{Nothing, AbstractVector},
-                 skip::F) where {V, F}
+                 skip::F,
+                 self_idx::Int) where {V, F}
 
-    return knn_kernel!(tree, point, best_idxs, best_dists, skip, nothing)
+    return knn_kernel!(tree, point, best_idxs, best_dists, skip, nothing, self_idx)
 end
 
 function knn_kernel!(tree::BruteTree{V},
                      point::AbstractVector,
                      best_idxs::Union{Integer, AbstractVector{<:Integer}},
                      best_dists::Union{Number, AbstractVector},
                      skip::F,
-                     dedup::MaybeBitSet) where {V, F}
+                     dedup::MaybeBitSet,
+                     self_idx::Int) where {V, F}
     has_set = dedup !== nothing
     for i in 1:length(tree.data)
-        if skip(i)
+        if (i == self_idx) || skip(i)
             continue
         end
 
@@ -80,8 +82,9 @@ function _inrange(tree::BruteTree,
                   point::AbstractVector,
                   radius::Number,
                   idx_in_ball::Union{Nothing, Vector{<:Integer}},
-                  skip::F,) where {F}
-    return inrange_kernel!(tree, point, radius, idx_in_ball, skip, nothing)
+                  skip::F,
+                  self_idx::Int) where {F}
+    return inrange_kernel!(tree, point, radius, idx_in_ball, skip, nothing, self_idx)
 end
 
 
@@ -90,11 +93,12 @@ function inrange_kernel!(tree::BruteTree,
                          r::Number,
                          idx_in_ball::Union{Nothing, Vector{<:Integer}},
                          skip::Function,
-                         dedup::MaybeBitSet)
+                         dedup::MaybeBitSet,
+                         self_idx::Int)
     count = 0
     has_set = dedup !== nothing
     for i in 1:length(tree.data)
-        if skip(i)
+        if (i == self_idx) || skip(i)
             continue
         end
         d = evaluate(tree.metric, tree.data[i], point)

src/inrange.jl

@@ -9,6 +9,8 @@ Find all the points in the tree which are closer than `radius` to `points`.
 - `tree`: The tree instance
 - `points`: Query point(s) - can be a vector (single point), matrix (multiple points), or vector of vectors
 - `radius`: Search radius
+- `skip` (optional): Predicate to skip certain points.
+- `skip_self` (optional, batched queries only): When querying the same dataset, skip the point whose index matches the query index.
 
 # Returns
 - `indices`: Vector of indices of points within the radius
@@ -19,23 +21,25 @@ function inrange(tree::NNTree,
                  points::AbstractVector{T},
                  radius::Number,
                  sortres=false,
-                 skip::F = Returns(false)) where {T <: AbstractVector, F}
+                 skip::F = Returns(false);
+                 skip_self::Bool=false) where {T <: AbstractVector, F}
     check_input(tree, points)
     check_for_nan_in_points(points)
     check_radius(radius)
 
     idxs = [Vector{Int}() for _ in 1:length(points)]
 
     for i in 1:length(points)
-        inrange_point!(tree, points[i], radius, sortres, idxs[i], skip)
+        self_idx = skip_self ? i : 0
+        inrange_point!(tree, points[i], radius, sortres, idxs[i], skip, self_idx)
     end
     return idxs
 end
 
-inrange_point!(tree, point, radius, sortres, idx, skip::F) where {F} = _inrange_point!(tree, point, radius, sortres, idx, skip)
+inrange_point!(tree, point, radius, sortres, idx, skip::F, self_idx::Int) where {F} = _inrange_point!(tree, point, radius, sortres, idx, skip, self_idx)
 
-function _inrange_point!(tree, point, radius, sortres, idx, skip::F) where {F}
-    count = _inrange(tree, point, radius, idx, skip)
+function _inrange_point!(tree, point, radius, sortres, idx, skip::F, self_idx::Int) where {F}
+    count = _inrange(tree, point, radius, idx, skip, self_idx)
     if idx !== nothing
         inner_tree = get_tree(tree)
         if inner_tree.reordered
@@ -62,25 +66,38 @@ Useful to avoid allocations or specify the element type of the output vector.
 
 See also: `inrange`, `inrangecount`.
 """
-function inrange!(idxs::AbstractVector, tree::NNTree{V}, point::AbstractVector{T}, radius::Number, sortres=false, skip=Returns(false)) where {V, T <: Number}
+function inrange!(idxs::AbstractVector, tree::NNTree{V}, point::AbstractVector{T}, radius::Number, sortres=false, skip=Returns(false); skip_self::Bool=false) where {V, T <: Number}
+    skip_self && throw(ArgumentError("skip_self is only supported for batched queries; pass a skip predicate instead for single points"))
     check_input(tree, point)
     check_for_nan_in_points(point)
     check_radius(radius)
     length(idxs) == 0 || throw(ArgumentError("idxs must be empty"))
-    inrange_point!(tree, point, radius, sortres, idxs, skip)
+    inrange_point!(tree, point, radius, sortres, idxs, skip, 0)
     return idxs
 end
 
-function inrange(tree::NNTree{V}, point::AbstractVector{T}, radius::Number, sortres=false) where {V, T <: Number}
+function inrange(tree::NNTree{V}, point::AbstractVector{T}, radius::Number, sortres=false; skip_self::Bool=false) where {V, T <: Number}
+    skip_self && throw(ArgumentError("skip_self is only supported for batched queries; pass a skip predicate instead for single points"))
     return inrange!(Int[], tree, point, radius, sortres)
 end
 
-function inrange(tree::NNTree{V}, points::AbstractMatrix{T}, radius::Number, sortres=false) where {V, T <: Number}
+# Single-point variant with an explicit skip predicate
+function inrange(tree::NNTree{V}, point::AbstractVector{T}, radius::Number, sortres::Bool, skip::F=Returns(false); skip_self::Bool=false) where {V, T <: Number, F}
+    skip_self && throw(ArgumentError("skip_self is only supported for batched queries; pass a skip predicate instead for single points"))
+    return inrange!(Int[], tree, point, radius, sortres, skip)
+end
+
+function inrange(tree::NNTree{V}, points::AbstractMatrix{T}, radius::Number, sortres=false; skip_self::Bool=false) where {V, T <: Number}
+    dim = size(points, 1)
+    inrange_matrix(tree, points, radius, Val(dim), sortres; skip_self)
+end
+
+function inrange(tree::NNTree{V}, points::AbstractMatrix{T}, radius::Number, sortres::Bool, skip::F=Returns(false); skip_self::Bool=false) where {V, T <: Number, F}
     dim = size(points, 1)
-    inrange_matrix(tree, points, radius, Val(dim), sortres)
+    inrange_matrix(tree, points, radius, Val(dim), sortres, skip; skip_self)
 end
 
-function inrange_matrix(tree::NNTree{V}, points::AbstractMatrix{T}, radius::Number, ::Val{dim}, sortres, skip::F=Returns(false)) where {V, T <: Number, dim, F}
+function inrange_matrix(tree::NNTree{V}, points::AbstractMatrix{T}, radius::Number, ::Val{dim}, sortres, skip::F=Returns(false); skip_self::Bool=false) where {V, T <: Number, dim, F}
     # TODO: DRY with inrange for AbstractVector
     check_input(tree, points)
     check_for_nan_in_points(points)
@@ -90,7 +107,8 @@ function inrange_matrix(tree::NNTree{V}, points::AbstractMatrix{T}, radius::Numb
 
     for i in 1:n_points
         point = SVector{dim,T}(ntuple(j -> points[j, i], Val(dim)))
-        inrange_point!(tree, point, radius, sortres, idxs[i], skip)
+        self_idx = skip_self ? i : 0
+        inrange_point!(tree, point, radius, sortres, idxs[i], skip, self_idx)
     end
     return idxs
 end
@@ -104,32 +122,40 @@ Count all the points in the tree which are closer than `radius` to `points`.
 - `tree`: The tree instance
 - `points`: Query point(s) - can be a vector (single point), matrix (multiple points), or vector of vectors
 - `radius`: Search radius
+- `skip` (optional): Predicate to skip certain points.
+- `skip_self` (optional, batched queries only): When querying the same dataset, skip the point whose index matches the query index.
 
 # Returns
 - `count`: Number of points within the radius (integer for single point, vector for multiple points)
 """
-function inrangecount(tree::NNTree{V}, point::AbstractVector{T}, radius::Number, skip::F=Returns(false)) where {V, T <: Number, F}
+function inrangecount(tree::NNTree{V}, point::AbstractVector{T}, radius::Number, skip::F=Returns(false); skip_self::Bool=false) where {V, T <: Number, F}
+    skip_self && throw(ArgumentError("skip_self is only supported for batched queries; pass a skip predicate instead for single points"))
     check_input(tree, point)
     check_for_nan_in_points(point)
     check_radius(radius)
-    return inrange_point!(tree, point, radius, false, nothing, skip)
+    return inrange_point!(tree, point, radius, false, nothing, skip, 0)
 end
 
 function inrangecount(tree::NNTree,
         points::AbstractVector{T},
-        radius::Number, skip::F=Returns(false)) where {T <: AbstractVector, F}
+        radius::Number, skip::F=Returns(false); skip_self::Bool=false) where {T <: AbstractVector, F}
     check_input(tree, points)
     check_for_nan_in_points(points)
     check_radius(radius)
-    return inrange_point!.(Ref(tree), points, radius, false, nothing, skip)
+    counts = Vector{Int}(undef, length(points))
+    for i in 1:length(points)
+        self_idx = skip_self ? i : 0
+        counts[i] = inrange_point!(tree, points[i], radius, false, nothing, skip, self_idx)
+    end
+    return counts
 end
 
-function inrangecount(tree::NNTree{V}, points::AbstractMatrix{T}, radius::Number, skip::F=Returns(false)) where {V, T <: Number, F}
+function inrangecount(tree::NNTree{V}, points::AbstractMatrix{T}, radius::Number, skip::F=Returns(false); skip_self::Bool=false) where {V, T <: Number, F}
     dim = size(points, 1)
-    inrangecount_matrix(tree, points, radius, Val(dim), skip)
+    inrangecount_matrix(tree, points, radius, Val(dim), skip; skip_self)
 end
 
-function inrangecount_matrix(tree::NNTree{V}, points::AbstractMatrix{T}, radius::Number, ::Val{dim}, skip::F=Returns(false)) where {V, T <: Number, dim, F}
+function inrangecount_matrix(tree::NNTree{V}, points::AbstractMatrix{T}, radius::Number, ::Val{dim}, skip::F=Returns(false); skip_self::Bool=false) where {V, T <: Number, dim, F}
     check_input(tree, points)
     check_for_nan_in_points(points)
     check_radius(radius)
@@ -138,7 +164,8 @@ function inrangecount_matrix(tree::NNTree{V}, points::AbstractMatrix{T}, radius:
 
     for i in 1:n_points
         point = SVector{dim,T}(ntuple(j -> points[j, i], Val(dim)))
-        counts[i] = inrange_point!(tree, point, radius, false, nothing, skip)
+        self_idx = skip_self ? i : 0
+        counts[i] = inrange_point!(tree, point, radius, false, nothing, skip, self_idx)
     end
     return counts
 end