From 33ccb17e7cc896ff22917ffa68af061c48b2c187 Mon Sep 17 00:00:00 2001
From: Sebastian Ament <sebastianeament@gmail.com>
Date: Fri, 7 Jan 2022 13:32:20 +0100
Subject: [PATCH] parallelizing knn and inrange searches

---
 Project.toml            | 2 +-
 src/NearestNeighbors.jl | 1 +
 src/inrange.jl          | 2 +-
 src/knn.jl              | 8 ++++----
 test/test_inrange.jl    | 4 ++--
 test/test_knn.jl        | 4 ++--
 test/test_monkey.jl     | 6 +++---
 7 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/Project.toml b/Project.toml
index 5b5d135..0380f83 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,6 +1,6 @@
 name = "NearestNeighbors"
 uuid = "b8a86587-4115-5ab1-83bc-aa920d37bbce"
-version = "0.4.9"
+version = "0.4.10"
 
 [deps]
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
diff --git a/src/NearestNeighbors.jl b/src/NearestNeighbors.jl
index 6aa4796..7e8c5b1 100644
--- 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
index be5c7a8..d678cad 100644
--- 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
index 6fb4bb1..34bd2ce 100644
--- 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
index 652faf2..1988cf5 100644
--- 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
index 3648ae5..7c94776 100644
--- 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
index a3fdb0b..12f41f2 100644
--- 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