improve performance of ClusterTree constructor

- use a buffer to store the intermediate permutations
- only permute the elements at the end of the tree construction

Author: maltezfaria

benchmarks/clustertree_bench.jl

@@ -0,0 +1,9 @@
+import WavePropBase as WPB
+using StaticArrays
+using BenchmarkTools
+
+n = 1_000_000
+
+pts      = rand(WPB.Point3D,n)
+splitter = WPB.GeometricSplitter(nmax=100)
+@btime WPB.ClusterTree($pts,$splitter;threads=false)

src/Trees/clustertree.jl

@@ -22,12 +22,13 @@ mutable struct ClusterTree{T,S,D} <: AbstractTree
     container::S
     index_range::UnitRange{Int}
     loc2glob::Vector{Int}
+    buffer::Vector{Int}
     children::Vector{ClusterTree{T,S,D}}
     parent::ClusterTree{T,S,D}
     data::D
     # inner constructors handling missing fields.
-    function ClusterTree{D}(els::T,container::S,loc_idxs,loc2glob,children,parent,data=nothing) where {T,S,D}
-        clt = new{T,S,D}(els,container,loc_idxs,loc2glob)
+    function ClusterTree{D}(els::T,container::S,loc_idxs,loc2glob,buffer,children,parent,data=nothing) where {T,S,D}
+        clt = new{T,S,D}(els,container,loc_idxs,loc2glob,buffer)
         clt.children = isnothing(children) ? Vector{typeof(clt)}() : children
         clt.parent   = isnothing(parent)   ? clt : parent
         clt.data     = isnothing(data)     ? D() : data
@@ -76,6 +77,8 @@ the (global) indexes used upon the construction of the tree.
 """
 loc2glob(clt::ClusterTree)      = clt.loc2glob
 
+buffer(clt::ClusterTree)      = clt.buffer
+
 """
     container_type(clt::ClusterTree)
 
@@ -126,27 +129,29 @@ function ClusterTree{D}(elements,splitter=CardinalitySplitter();copy_elements=tr
         bbox         = HyperRectangle(elements)
     end
     n            = length(elements)
-    irange     = 1:n
+    irange       = 1:n
     loc2glob     = collect(irange)
+    buffer       = collect(irange)
     children     = nothing
     parent       = nothing
     #build the root, then recurse
-    root         = ClusterTree{D}(elements,bbox,irange,loc2glob,children,parent)
+    root         = ClusterTree{D}(elements,bbox,irange,loc2glob,buffer,children,parent)
     _build_cluster_tree!(root,splitter,threads)
+    # finally, permute the elements so as to use the local indexing
+    copy!(elements,elements[loc2glob]) # faster than permute!
     return root
 end
 ClusterTree(args...;kwargs...) = ClusterTree{Nothing}(args...;kwargs...)
 
 function _build_cluster_tree!(current_node,splitter,threads)
     if should_split(current_node,splitter)
-        children = split!(current_node,splitter)
-        current_node.children = children
+        split!(current_node,splitter)
         if threads
-            Threads.@threads for child in children
+            Threads.@threads for child in children(current_node)
                 _build_cluster_tree!(child,splitter,threads)
             end
         else
-            for child in children
+            for child in children(current_node)
                 _build_cluster_tree!(child,splitter,threads)
             end
         end
@@ -180,18 +185,18 @@ function _common_binary_split!(cluster::ClusterTree{T,S,D},conditions;
                                parentcluster) where {T,S,D}
     els    = root_elements(cluster)
     l2g    = loc2glob(cluster)
+    buf    = buffer(cluster)
     irange = index_range(cluster)
     # get split data
     npts_left,npts_right,left_rec,right_rec,buff = binary_split_data(cluster,conditions)
     @assert npts_left + npts_right == length(irange) "elements lost during split"
-    l2g[irange] = l2g[buff]
-    els[irange] = els[buff] # reorders the global index set
+    copy!(view(l2g,irange),buff)
     # new ranges for cluster
     left_indices  = irange.start:(irange.start)+npts_left-1
     right_indices = (irange.start+npts_left):irange.stop
     # create children
-    clt1 = ClusterTree{D}(els,left_rec,left_indices,l2g,nothing,parentcluster)
-    clt2 = ClusterTree{D}(els,right_rec,right_indices,l2g,nothing,parentcluster)
+    clt1 = ClusterTree{D}(els,left_rec,left_indices,l2g,buf,nothing,parentcluster)
+    clt2 = ClusterTree{D}(els,right_rec,right_indices,l2g,buf,nothing,parentcluster)
     return clt1, clt2
 end
 function binary_split_data(cluster::ClusterTree{T,S},conditions::Function) where {T,S}
@@ -200,23 +205,24 @@ function binary_split_data(cluster::ClusterTree{T,S},conditions::Function) where
     els        = root_elements(cluster)
     irange     = index_range(cluster)
     n          = length(irange)
+    buff       = view(cluster.buffer,irange)
+    l2g        = loc2glob(cluster)
     npts_left  = 0
     npts_right = 0
-    buff = Vector{Int}(undef,length(cluster))
     xl_left = xl_right = high_corner(rec)
     xu_left = xu_right = low_corner(rec)
     #sort the points into left and right rectangle
     for i in irange
-        pt = els[i] |> center
+        pt = els[l2g[i]] |> center
         if f(pt)
             xl_left = min.(xl_left,pt)
             xu_left = max.(xu_left,pt)
             npts_left += 1
-            buff[npts_left] = i
+            buff[npts_left] = l2g[i]
         else
             xl_right = min.(xl_right,pt)
             xu_right = max.(xu_right,pt)
-            buff[n-npts_right] = i
+            buff[n-npts_right] = l2g[i]
             npts_right += 1
         end
     end
@@ -231,19 +237,21 @@ function binary_split_data(cluster::ClusterTree,conditions::Tuple{Integer,Real})
     els        = root_elements(cluster)
     irange     = index_range(cluster)
     n          = length(irange)
+    l2g        = loc2glob(cluster)
     npts_left  = 0
     npts_right = 0
-    buff = Vector{Int}(undef,length(cluster))
+    buff       = view(cluster.buffer,irange)
     # bounding boxes
     left_rec, right_rec = split(rec,dir,pos)
     #sort the points into left and right rectangle
     for i in irange
-        pt = els[i] |> center
+        pt = els[l2g[i]] |> center
         if pt in left_rec
             npts_left += 1
-            buff[npts_left] = i
+            buff[npts_left] = l2g[i]
         else # pt in right_rec
-            buff[n-npts_right] = i
+            # @assert pt in right_rec
+            buff[n-npts_right] = l2g[i]
             npts_right += 1
         end
     end

src/Trees/splitter.jl

@@ -24,7 +24,9 @@ end
 """
     split!(clt::ClusterTree,splitter::AbstractSplitter)
 
-Divide `clt` using the strategy implemented by `splitter`.
+Divide `clt` using the strategy implemented by `splitter`. This function is
+reponsible of assigning the `children` and `parent` fields, as well as of
+permuting the data of `clt`.
 """
 function split!(clt, splitter::AbstractSplitter)
     abstractmethod(splitter)
@@ -59,7 +61,8 @@ function split!(parentcluster::ClusterTree, ::DyadicSplitter)
             append!(clusters, _binary_split!(clt, i, pos; parentcluster))
         end
     end
-    return clusters
+    parentcluster.children = clusters
+    return parentcluster
 end
 
 """
@@ -77,7 +80,8 @@ function split!(cluster::ClusterTree, ::GeometricSplitter)
     rec = cluster.container
     wmax, imax = findmax(high_corner(rec) - low_corner(rec))
     left_node, right_node = _binary_split!(cluster, imax, low_corner(rec)[imax] + wmax / 2)
-    return [left_node, right_node]
+    cluster.children = [left_node,right_node]
+    return cluster
 end
 
 """
@@ -97,7 +101,8 @@ function split!(cluster::ClusterTree, ::GeometricMinimalSplitter)
     mid = low_corner(rec)[imax] + wmax / 2
     predicate = (x) -> x[imax] < mid
     left_node, right_node = _binary_split!(cluster, predicate)
-    return [left_node, right_node]
+    cluster.children = [left_node,right_node]
+    return cluster
 end
 
 """
@@ -114,25 +119,28 @@ function split!(cluster::ClusterTree, ::PrincipalComponentSplitter)
     irange = cluster.index_range
     xc = center_of_mass(cluster)
     # compute covariance matrix for principal direction
+    l2g = loc2glob(cluster)
     cov = sum(irange) do i
-        x = coords(pts[i])
+        x = coords(pts[l2g[i]])
         (x - xc) * transpose(x - xc)
     end
     v = eigvecs(cov)[:, end]
     predicate = (x) -> dot(x - xc, v) < 0
     left_node, right_node = _binary_split!(cluster, predicate)
-    return [left_node, right_node]
+    cluster.children = [left_node,right_node]
+    return cluster
 end
 
 function center_of_mass(clt::ClusterTree)
     pts = clt._elements
     loc_idxs = clt.index_range
+    l2g = loc2glob(clt)
     # w    = clt.weights
     n = length(loc_idxs)
     # M    = isempty(w) ? n : sum(i->w[i],glob_idxs)
     # xc   = isempty(w) ? sum(i->pts[i]/M,glob_idxs) : sum(i->w[i]*pts[i]/M,glob_idxs)
     M = n
-    xc = sum(i -> coords(pts[i]) / M, loc_idxs)
+    xc = sum(i -> coords(pts[l2g[i]]) / M, loc_idxs)
     return xc
 end
 
@@ -156,8 +164,10 @@ function split!(cluster::ClusterTree, ::CardinalitySplitter)
     irange = cluster.index_range
     rec = container(cluster)
     _, imax = findmax(high_corner(rec) - low_corner(rec))
-    med = median(coords(points[i])[imax] for i in irange) # the median along largest axis `imax`
+    l2g = loc2glob(cluster)
+    med = median(coords(points[l2g[i]])[imax] for i in irange) # the median along largest axis `imax`
     predicate = (x) -> x[imax] < med
     left_node, right_node = _binary_split!(cluster, predicate)
-    return [left_node, right_node]
+    cluster.children = [left_node,right_node]
+    return cluster
 end