parallelizing knn and inrange searches

KristofferC · Jan 7, 2022 · 5ae189d · 5ae189d
1 parent 1de39f5
commit 5ae189d
Show file tree

Hide file tree

Showing 6 changed files with 13 additions and 12 deletions.
diff --git a/src/NearestNeighbors.jl b/src/NearestNeighbors.jl
@@ -5,6 +5,7 @@ import Distances: Metric, result_type, eval_reduce, eval_end, eval_op, eval_star
 
 using StaticArrays
 import Base.show
+using Base.Threads: @threads
 
 export NNTree, BruteTree, KDTree, BallTree, DataFreeTree
 export knn, nn, inrange # TODOs? , allpairs, distmat, npairs

diff --git a/src/inrange.jl b/src/inrange.jl
@@ -15,7 +15,7 @@ function inrange(tree::NNTree,
 
     idxs = [Vector{Int}() for _ in 1:length(points)]
 
-    for i in 1:length(points)
+    @threads for i in 1:length(points)
         inrange_point!(tree, points[i], radius, sortres, idxs[i])
     end
     return idxs

diff --git a/src/knn.jl b/src/knn.jl
@@ -11,16 +11,16 @@ end
 Performs a lookup of the `k` nearest neigbours to the `points` from the data
 in the `tree`. If `sortres = true` the result is sorted such that the results are
 in the order of increasing distance to the point. `skip` is an optional predicate
-to determine if a point that would be returned should be skipped based on its 
+to determine if a point that would be returned should be skipped based on its
 index.
 """
 function knn(tree::NNTree{V}, points::Vector{T}, k::Int, sortres=false, skip::F=always_false) where {V, T <: AbstractVector, F<:Function}
     check_input(tree, points)
     check_k(tree, k)
     n_points = length(points)
-     dists = [Vector{get_T(eltype(V))}(undef, k) for _ in 1:n_points]
-     idxs = [Vector{Int}(undef, k) for _ in 1:n_points]
-    for i in 1:n_points
+    dists = [Vector{get_T(eltype(V))}(undef, k) for _ in 1:n_points]
+    idxs = [Vector{Int}(undef, k) for _ in 1:n_points]
+    @threads for i in 1:n_points
         knn_point!(tree, points[i], sortres, dists[i], idxs[i], skip)
     end
     return idxs, dists

diff --git a/test/test_inrange.jl b/test/test_inrange.jl
@@ -1,7 +1,7 @@
 # Does not test leafsize
 @testset "inrange" begin
-    @testset "metric" for metric in [Euclidean()]
-        @testset "tree type" for TreeType in trees_with_brute
+    @testset "metric $Metric" for metric in [Euclidean()]
+        @testset "tree type $TreeType" for TreeType in trees_with_brute
             function test(data)
                 tree = TreeType(data, metric; leafsize=2)
                 dosort = true

diff --git a/test/test_knn.jl b/test/test_knn.jl
@@ -3,8 +3,8 @@
 import Distances.evaluate
 
 @testset "knn" begin
-    @testset "metric" for metric in [metrics; WeightedEuclidean(ones(2))]
-        @testset "tree type" for TreeType in trees_with_brute
+    @testset "metric $metric" for metric in [metrics; WeightedEuclidean(ones(2))]
+        @testset "tree type $TreeType" for TreeType in trees_with_brute
             function test(data)
                 tree = TreeType(data, metric; leafsize=2)
 

diff --git a/test/test_monkey.jl b/test/test_monkey.jl
@@ -3,9 +3,9 @@ import NearestNeighbors.MinkowskiMetric
 # some edge case has been missed in the real tests
 
 
-@testset "metric" for metric in fullmetrics
-    @testset "tree type" for TreeType in trees_with_brute
-        @testset "type" for T in (Float32, Float64)
+@testset "metric $metric" for metric in fullmetrics
+    @testset "tree type $TreeType" for TreeType in trees_with_brute
+        @testset "element type $T" for T in (Float32, Float64)
             @testset "knn monkey" begin
                 # Checks that we find existing point in the tree
                 # and that it is the closest