make ClusterTree more generic

- A `ClusterTree` can now be constructed from `AbstractVector`

Author: maltezfaria

src/Trees/abstracttree.jl

@@ -86,6 +86,45 @@ function filter_tree!(f,nodes,tree,isterminal=true)
     return nodes
 end
 
+"""
+    partition_by_depth(tree)
+
+Given a `tree`, return a `partition` vector whose `i`-th entry stores all the nodes in
+`tree` with `depth=i-1`. Empty nodes are not added to the partition.
+"""
+function partition_by_depth(tree)
+    T         = eltype(tree)
+    partition = Vector{Vector{T}}()
+    depth = 0
+    _partition_by_depth!(partition,tree,depth)
+end
+
+function _partition_by_depth!(partition,tree,depth)
+    T = eltype(tree)
+    if length(partition) < depth+1
+         push!(partition,T[])
+    end
+    length(tree) > 0 && push!(partition[depth+1],tree)
+    for chd in children(tree)
+        _partition_by_depth!(partition,chd,depth+1)
+    end
+    return partition
+end
+
+"""
+    partition_by_height(tree)
+
+Given a `tree`, return a `partition` vector whose `i`-th entry stores all the nodes in
+`tree` with `height=i-1`. The `height` of the tree is thus `lenth(partition)`,
+with `partition(end)==tree`.
+"""
+function partition_by_height(tree)
+    # TODO: how to do this more or less efficiently? One idea is to start at the
+    # leaves, push them, then push their parents and recurse, making sure call
+    # `unique!` as you go up in order to avoid duplicate parents.
+end
+
+
 # interface to AbstractTrees. No children is determined by an empty tuple for
 # AbstractTrees.
 AbstractTrees.children(t::AbstractTree) = isleaf(t) ? () : t.children

src/Trees/clustertree.jl

@@ -1,13 +1,13 @@
 """
     mutable struct ClusterTree{T,S,D}
 
-Tree structure used to cluster elements of type `T` into containers of type `S`.
-The method `center(::T)::SVector` is required for the clustering algorithms. An
+Tree structure used to cluster elements of type `V = eltype(T)` into containers of type `S`.
+The method `center(::V)::SVector` is required for the clustering algorithms. An
 additional `data` field of type `D` can be associated with each node to store
 node-specific information (it defaults to `D=Nothing`).
 
 # Fields:
-- `_elements::Vector{T}` : vector containing the sorted elements.
+- `_elements::T` : vector containing the sorted elements.
 - `container::S` : container for the elements in the current node.
 - `index_range::UnitRange{Int}` : indices of elements contained in the current node.
 - `loc2glob::Vector{Int}` : permutation from the local indexing system to the
@@ -17,15 +17,15 @@ original (global) indexing system used as input in the construction of the tree.
 - `data::D` : generic data field of type `D`.
 """
 mutable struct ClusterTree{T,S,D} <: AbstractTree
-    _elements::Vector{T}
+    _elements::T
     container::S
     index_range::UnitRange{Int}
     loc2glob::Vector{Int}
     children::Vector{ClusterTree{T,S,D}}
     parent::ClusterTree{T,S,D}
     data::D
     # inner constructors handling missing fields.
-    function ClusterTree{D}(els::Vector{T},container::S,loc_idxs,loc2glob,children,parent,data=nothing) where {T,S,D}
+    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)
         clt.children = isnothing(children) ? Vector{typeof(clt)}() : children
         clt.parent   = isnothing(parent)   ? clt : parent
@@ -87,7 +87,7 @@ container_type(::ClusterTree{T,S}) where {T,S} = S
 
 Type of elements sorted in `clt`.
 """
-element_type(::ClusterTree{T}) where {T} = T
+element_type(::ClusterTree{T}) where {T} = eltype(T)
 
 isleaf(clt::ClusterTree)  = isempty(clt.children)
 isroot(clt::ClusterTree)  = clt.parent == clt

src/Trees/splitter.jl

@@ -17,7 +17,7 @@ abstract type AbstractSplitter end
 
 Determine whether or not a `ClusterTree` should be further divided.
 """
-function should_split(clt,splitter::AbstractSplitter)
+function should_split(clt, splitter::AbstractSplitter)
     abstract_method(splitter)
 end
 
@@ -26,7 +26,7 @@ end
 
 Divide `clt` using the strategy implemented by `splitter`.
 """
-function split!(clt,splitter::AbstractSplitter)
+function split!(clt, splitter::AbstractSplitter)
     abstract_method(splitter)
 end
 
@@ -37,24 +37,24 @@ Used to split an `N` dimensional `ClusterTree` into `2^N` children until at most
 `nmax` points are contained in node.
 """
 Base.@kwdef struct DyadicSplitter <: AbstractSplitter
-    nmax::Int=typemax(Int)
+    nmax::Int = typemax(Int)
 end
 
-function should_split(node::ClusterTree,splitter::DyadicSplitter)
+function should_split(node::ClusterTree, splitter::DyadicSplitter)
     return length(node) > splitter.nmax
 end
 
-function split!(parentcluster::ClusterTree,::DyadicSplitter)
-    d        = ambient_dimension(parentcluster)
+function split!(parentcluster::ClusterTree, ::DyadicSplitter)
+    d = ambient_dimension(parentcluster)
     clusters = [parentcluster]
     rec = container(parentcluster)
     rec_center = center(rec)
-    for i in 1:d
+    for i = 1:d
         pos = rec_center[i]
         nel = length(clusters) #2^(i-1)
-        for _ in 1:nel
+        for _ = 1:nel
             clt = popfirst!(clusters)
-            append!(clusters,_binary_split!(clt,i,pos;parentcluster))
+            append!(clusters, _binary_split!(clt, i, pos; parentcluster))
         end
     end
     return clusters
@@ -65,16 +65,16 @@ end
 
 Used to split a `ClusterTree` in half along the largest axis.
 """
-@Base.kwdef struct GeometricSplitter <: AbstractSplitter
-    nmax::Int=50
+Base.@kwdef struct GeometricSplitter <: AbstractSplitter
+    nmax::Int = 50
 end
 
-should_split(node::ClusterTree,splitter::GeometricSplitter) = length(node) > splitter.nmax
+should_split(node::ClusterTree, splitter::GeometricSplitter) = length(node) > splitter.nmax
 
-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)
+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]
 end
 
@@ -83,54 +83,54 @@ end
 
 Like [`GeometricSplitter`](@ref), but shrinks the children's containters.
 """
-@Base.kwdef struct GeometricMinimalSplitter <: AbstractSplitter
-    nmax::Int=50
+Base.@kwdef struct GeometricMinimalSplitter <: AbstractSplitter
+    nmax::Int = 50
 end
 
-should_split(node::ClusterTree,splitter::GeometricMinimalSplitter) = length(node) > splitter.nmax
+should_split(node::ClusterTree, splitter::GeometricMinimalSplitter) = length(node) > splitter.nmax
 
-function split!(cluster::ClusterTree,::GeometricMinimalSplitter)
-    rec  = cluster.container
-    wmax, imax  = findmax(high_corner(rec) - low_corner(rec))
-    mid = low_corner(rec)[imax]+wmax/2
+function split!(cluster::ClusterTree, ::GeometricMinimalSplitter)
+    rec = cluster.container
+    wmax, imax = findmax(high_corner(rec) - low_corner(rec))
+    mid = low_corner(rec)[imax] + wmax / 2
     predicate = (x) -> x[imax] < mid
-    left_node,right_node =  _binary_split!(cluster,predicate)
+    left_node, right_node = _binary_split!(cluster, predicate)
     return [left_node, right_node]
 end
 
 """
     struct PrincipalComponentSplitter <: AbstractSplitter
 """
-@Base.kwdef struct PrincipalComponentSplitter <: AbstractSplitter
-    nmax::Int=50
+Base.@kwdef struct PrincipalComponentSplitter <: AbstractSplitter
+    nmax::Int = 50
 end
 
-should_split(node::ClusterTree,splitter::PrincipalComponentSplitter) = length(node) > splitter.nmax
+should_split(node::ClusterTree, splitter::PrincipalComponentSplitter) = length(node) > splitter.nmax
 
-function split!(cluster::ClusterTree,::PrincipalComponentSplitter)
-    pts       = cluster._elements
-    irange    = cluster.index_range
-    xc        = center_of_mass(cluster)
+function split!(cluster::ClusterTree, ::PrincipalComponentSplitter)
+    pts = cluster._elements
+    irange = cluster.index_range
+    xc = center_of_mass(cluster)
     # compute covariance matrix for principal direction
-    cov  = sum(irange) do i
+    cov = sum(irange) do i
         x = coords(pts[i])
-        (x - xc)*transpose(x - xc)
+        (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)
+    v = eigvecs(cov)[:, end]
+    predicate = (x) -> dot(x - xc, v) < 0
+    left_node, right_node = _binary_split!(cluster, predicate)
     return [left_node, right_node]
 end
 
 function center_of_mass(clt::ClusterTree)
-    pts       = clt._elements
-    loc_idxs  = clt.index_range
+    pts = clt._elements
+    loc_idxs = clt.index_range
     # w    = clt.weights
-    n    = length(loc_idxs)
+    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)
+    M = n
+    xc = sum(i -> coords(pts[i]) / M, loc_idxs)
     return xc
 end
 
@@ -141,19 +141,19 @@ Used to split a `ClusterTree` along the largest dimension if
 `length(tree)>nmax`. The split is performed so the `data` is evenly distributed
 amongst all children.
 """
-@Base.kwdef struct CardinalitySplitter <: AbstractSplitter
-    nmax::Int=50
+Base.@kwdef struct CardinalitySplitter <: AbstractSplitter
+    nmax::Int = 50
 end
 
-should_split(node::ClusterTree,splitter::CardinalitySplitter) = length(node) > splitter.nmax
+should_split(node::ClusterTree, splitter::CardinalitySplitter) = length(node) > splitter.nmax
 
-function split!(cluster::ClusterTree,::CardinalitySplitter)
-    points     = cluster._elements
-    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`
+function split!(cluster::ClusterTree, ::CardinalitySplitter)
+    points = cluster._elements
+    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`
     predicate = (x) -> x[imax] < med
-    left_node, right_node = _binary_split!(cluster,predicate)
+    left_node, right_node = _binary_split!(cluster, predicate)
     return [left_node, right_node]
 end

src/Utils/Utils.jl

@@ -244,8 +244,9 @@ that `0 ≤ θ ≤ π` and ` -π < φ ≤ π`.
 """
 function cart2sph(x,y,z)
     azimuth   = atan(y,x)
-    elevation = atan(sqrt(x^2 + y^2),z)
-    r = sqrt(x^2 + y^2 + z^2)
+    a = x^2 + y^2
+    elevation = atan(sqrt(a),z)
+    r = sqrt(a + z^2)
     return r, elevation, azimuth
 end
 

test/Mesh/cartesianmesh_test.jl

@@ -29,7 +29,7 @@ end
 @testset "Two dimensions" begin
     l    = HyperRectangle((0.0,0.0),(1.0,1.0))
     E    = typeof(l) # type of mesh element
-    mesh = UniformCartesianMesh(domain=l,sz=(10,20))
+    mesh = UniformCartesianMesh(l,(10,20))
     iter = ElementIterator(mesh,E)
     @test iter[1,1] ≈ HyperRectangle((0,0),(0.1,0.05))
     @test length(iter) == 200