add v3 that uses own walkers with fat node

Author: KristofferC

benchmark/benchmark_tree_walk.jl

@@ -2,7 +2,7 @@ using BenchmarkTools
 using NearestNeighbors
 using StaticArrays
 using AbstractTrees
-using NearestNeighbors: children, _treeindex, children2, preorder2, leaves2, _isleaf2
+using NearestNeighbors: children, _treeindex, children2, preorder2, leaves2, _isleaf2, children3, preorder3, leaves3, _isleaf3
 
 # Create test trees of various sizes
 function make_trees(n_points, ndim=3)
@@ -30,6 +30,14 @@ function count_nodes_v2(tree)
     return count
 end
 
+function count_nodes_v3(tree)
+    count = 0
+    for _ in preorder3(tree)
+        count += 1
+    end
+    return count
+end
+
 # Benchmark 3: Sum of all leaf point norms
 function sum_leaf_norms_v1(tree)
     root = treeroot(tree)
@@ -52,6 +60,16 @@ function sum_leaf_norms_v2(tree)
     return total
 end
 
+function sum_leaf_norms_v3(tree)
+    total = 0.0
+    for node in leaves3(tree)
+        for pt in NearestNeighbors.leafpoints3(node)
+            total += sum(abs2, pt)
+        end
+    end
+    return total
+end
+
 # Benchmark 4: Collect all node indices
 function collect_indices_v1(tree)
     root = treeroot(tree)
@@ -70,26 +88,12 @@ function collect_indices_v2(tree)
     return indices
 end
 
-# Benchmark 5: Filter internal nodes
-function count_internal_nodes_v1(tree)
-    root = treeroot(tree)
-    count = 0
-    for node in PreOrderDFS(root)
-        if !isempty(children(node))
-            count += 1
-        end
-    end
-    return count
-end
-
-function count_internal_nodes_v2(tree)
-    count = 0
-    for node in preorder2(tree)
-        if !_isleaf2(tree, node)
-            count += 1
-        end
+function collect_indices_v3(tree)
+    indices = Int[]
+    for node in preorder3(tree)
+        push!(indices, node.index)
     end
-    return count
+    return indices
 end
 
 # Run benchmarks
@@ -98,7 +102,7 @@ function run_benchmarks()
 
     println("="^80)
     println("Tree Walking Benchmark Comparison")
-    println("AbstractTrees (v1) vs isbits nodes (v2)")
+    println("AbstractTrees (v1) vs isbits nodes (v2) vs wrapped nodes (v3)")
     println("="^80)
 
     for n in sizes
@@ -108,55 +112,56 @@ function run_benchmarks()
 
         kdtree, balltree = make_trees(n)
 
-        # Verify both approaches give same results
-        @assert count_nodes_v1(kdtree) == count_nodes_v2(kdtree)
-        @assert sum_leaf_norms_v1(kdtree) ≈ sum_leaf_norms_v2(kdtree)
-        @assert collect_indices_v1(kdtree) == collect_indices_v2(kdtree)
-        @assert count_internal_nodes_v1(kdtree) == count_internal_nodes_v2(kdtree)
+        # Verify all approaches give same results
+        @assert count_nodes_v1(kdtree) == count_nodes_v2(kdtree) == count_nodes_v3(kdtree)
+        @assert sum_leaf_norms_v1(kdtree) ≈ sum_leaf_norms_v2(kdtree) ≈ sum_leaf_norms_v3(kdtree)
+        @assert collect_indices_v1(kdtree) == collect_indices_v2(kdtree) == collect_indices_v3(kdtree)
 
         println("\n--- KDTree ---")
 
         print("Count nodes (v1): ")
         @btime count_nodes_v1($kdtree)
         print("Count nodes (v2): ")
         @btime count_nodes_v2($kdtree)
+        print("Count nodes (v3): ")
+        @btime count_nodes_v3($kdtree)
 
         print("\nSum leaf norms (v1): ")
         @btime sum_leaf_norms_v1($kdtree)
         print("Sum leaf norms (v2): ")
         @btime sum_leaf_norms_v2($kdtree)
+        print("Sum leaf norms (v3): ")
+        @btime sum_leaf_norms_v3($kdtree)
 
         print("\nCollect indices (v1): ")
         @btime collect_indices_v1($kdtree)
         print("Collect indices (v2): ")
         @btime collect_indices_v2($kdtree)
-
-        print("\nCount internal (v1): ")
-        @btime count_internal_nodes_v1($kdtree)
-        print("Count internal (v2): ")
-        @btime count_internal_nodes_v2($kdtree)
+        print("Collect indices (v3): ")
+        @btime collect_indices_v3($kdtree)
 
         println("\n--- BallTree ---")
 
         print("Count nodes (v1): ")
         @btime count_nodes_v1($balltree)
         print("Count nodes (v2): ")
         @btime count_nodes_v2($balltree)
+        print("Count nodes (v3): ")
+        @btime count_nodes_v3($balltree)
 
         print("\nSum leaf norms (v1): ")
         @btime sum_leaf_norms_v1($balltree)
         print("Sum leaf norms (v2): ")
         @btime sum_leaf_norms_v2($balltree)
+        print("Sum leaf norms (v3): ")
+        @btime sum_leaf_norms_v3($balltree)
 
         print("\nCollect indices (v1): ")
         @btime collect_indices_v1($balltree)
         print("Collect indices (v2): ")
         @btime collect_indices_v2($balltree)
-
-        print("\nCount internal (v1): ")
-        @btime count_internal_nodes_v1($balltree)
-        print("Count internal (v2): ")
-        @btime count_internal_nodes_v2($balltree)
+        print("Collect indices (v3): ")
+        @btime collect_indices_v3($balltree)
     end
 end
 

src/NearestNeighbors.jl

@@ -65,6 +65,7 @@ include("periodic_tree.jl")
 include("datafreetree.jl")
 include("tree_walk.jl")
 include("tree_walk_2.jl")
+include("tree_walk_3.jl")
 include("knn.jl")
 include("inrange.jl")