tests for periodic tree

Author: KristofferC

test/test_inrange.jl

@@ -93,9 +93,10 @@ end
 end
 
 @testset "skip_self keyword" begin
-    for T in (KDTree, BallTree, BruteTree)
+    for T in (KDTree, BallTree, BruteTree, PeriodicTree)
         data = rand(3, 20)
-        tree = T(data)
+        base_tree = T === PeriodicTree ? KDTree(data) : T(data)
+        tree = T === PeriodicTree ? PeriodicTree(base_tree, zeros(3), ones(3)) : base_tree
         # radius 0 includes only identical points
         idxs = inrange(tree, data, 0.0; skip_self=true)
         counts = inrangecount(tree, data, 0.0; skip_self=true)

test/test_knn.jl

@@ -98,29 +98,34 @@ end
     @test distance ≈ [0.02239688629947563, 0.13440059522389006]
 
     @testset "skip_self keyword" begin
-        data = rand(3, 25)
-        tree = KDTree(data)
-        idxs_skip, dists_skip = knn(tree, data, 1, true; skip_self=true)
-        @test all(i -> idxs_skip[i][1] != i, eachindex(idxs_skip))
-        expected_idxs = Vector{Int}(undef, size(data, 2))
-        expected_dists = Vector{Float64}(undef, size(data, 2))
-        for i in 1:size(data, 2)
-            single_idx, single_dist = knn(tree, data[:, i], 1, true, j -> j == i)
-            expected_idxs[i] = single_idx[1]
-            expected_dists[i] = single_dist[1]
+        for tree_builder in (data -> KDTree(data),
+                             data -> BallTree(data),
+                             data -> BruteTree(data),
+                             data -> PeriodicTree(KDTree(data), zeros(3), ones(3)))
+            data = rand(3, 25)
+            tree = tree_builder(data)
+            idxs_skip, dists_skip = knn(tree, data, 1, true; skip_self=true)
+            @test all(i -> idxs_skip[i][1] != i, eachindex(idxs_skip))
+            expected_idxs = Vector{Int}(undef, size(data, 2))
+            expected_dists = Vector{Float64}(undef, size(data, 2))
+            for i in 1:size(data, 2)
+                single_idx, single_dist = knn(tree, data[:, i], 1, true, j -> j == i)
+                expected_idxs[i] = single_idx[1]
+                expected_dists[i] = single_dist[1]
+            end
+            @test [idxs_skip[i][1] for i in eachindex(idxs_skip)] == expected_idxs
+            @test [dists_skip[i][1] for i in eachindex(dists_skip)] ≈ expected_dists
+
+            # Works for nn with batched queries too
+            nn_idx, nn_dist = nn(tree, data; skip_self=true)
+            @test nn_idx == expected_idxs
+            @test nn_dist ≈ expected_dists
+
+            # skip_self combines with a custom skip predicate
+            block_idx = 5
+            idxs_blocked, _ = knn(tree, data, 1, true, j -> j == block_idx; skip_self=true)
+            @test all(i -> idxs_blocked[i][1] != i && idxs_blocked[i][1] != block_idx, eachindex(idxs_blocked))
         end
-        @test [idxs_skip[i][1] for i in eachindex(idxs_skip)] == expected_idxs
-        @test [dists_skip[i][1] for i in eachindex(dists_skip)] ≈ expected_dists
-
-        # Works for nn with batched queries too
-        nn_idx, nn_dist = nn(tree, data; skip_self=true)
-        @test nn_idx == expected_idxs
-        @test nn_dist ≈ expected_dists
-
-        # skip_self combines with a custom skip predicate
-        block_idx = 5
-        idxs_blocked, _ = knn(tree, data, 1, true, j -> j == block_idx; skip_self=true)
-        @test all(i -> idxs_blocked[i][1] != i && idxs_blocked[i][1] != block_idx, eachindex(idxs_blocked))
     end
 end