Turn Leaf struct into a frequency map

Author: tecosaur

src/DecisionTree.jl

@@ -26,16 +26,18 @@ export InfoNode, InfoLeaf, wrap
 ###########################
 ########## Types ##########
 
-struct Leaf{T}
-    majority :: T
-    values   :: Vector{T}
+struct Leaf{T, N}
+    features :: NTuple{N, T}
+    majority :: Int
+    values   :: NTuple{N, Int}
+    total    :: Int
 end
 
-struct Node{S, T}
+struct Node{S, T, N}
     featid  :: Int
     featval :: S
-    left    :: Union{Leaf{T}, Node{S, T}}
-    right   :: Union{Leaf{T}, Node{S, T}}
+    left    :: Union{Leaf{T, N}, Node{S, T, N}}
+    right   :: Union{Leaf{T, N}, Node{S, T, N}}
 end
 
 const LeafOrNode{S, T} = Union{Leaf{T}, Node{S, T}}
@@ -52,13 +54,15 @@ struct Ensemble{S, T}
     featim  :: Vector{Float64}
 end
 
+Leaf(features::NTuple{T, N}) where {T, N} =
+    Leaf(features, 0, Tuple(zeros(T, N)), 0)
 
 is_leaf(l::Leaf) = true
 is_leaf(n::Node) = false
 
 _zero(::Type{String}) = ""
 _zero(x::Any) = zero(x)
-convert(::Type{Node{S, T}}, lf::Leaf{T}) where {S, T} = Node(0, _zero(S), lf, Leaf(_zero(T), [_zero(T)]))
+convert(::Type{Node{S, T}}, lf::Leaf{T}) where {S, T} = Node(0, _zero(S), lf, Leaf(lf.features))
 convert(::Type{Root{S, T}}, node::LeafOrNode{S, T}) where {S, T} = Root{S, T}(node, 0, Float64[])
 convert(::Type{LeafOrNode{S, T}}, tree::Root{S, T}) where {S, T} = tree.node
 promote_rule(::Type{Node{S, T}}, ::Type{Leaf{T}}) where {S, T} = Node{S, T}
@@ -97,9 +101,8 @@ depth(tree::Node) = 1 + max(depth(tree.left), depth(tree.right))
 depth(tree::Root) = depth(tree.node)
 
 function print_tree(io::IO, leaf::Leaf, depth=-1, indent=0; sigdigits=4, feature_names=nothing)
-    n_matches = count(leaf.values .== leaf.majority)
-    ratio = string(n_matches, "/", length(leaf.values))
-    println(io, "$(leaf.majority) : $(ratio)")
+    println(io, leaf.features[leaf.majority], " : ",
+            leaf.values[leaf.majority], '/', leaf.total)
 end
 function print_tree(leaf::Leaf, depth=-1, indent=0; sigdigits=4, feature_names=nothing)
     return print_tree(stdout, leaf, depth, indent; sigdigits, feature_names)
@@ -162,8 +165,8 @@ end
 
 function show(io::IO, leaf::Leaf)
     println(io, "Decision Leaf")
-    println(io, "Majority: $(leaf.majority)")
-    print(io,   "Samples:  $(length(leaf.values))")
+    println(io, "Majority: ", leaf.features[leaf.majority])
+    print(io,   "Samples:  ", leaf.total)
 end
 
 function show(io::IO, tree::Node)

src/classification/main.jl

@@ -41,11 +41,14 @@ function _convert(
     ) where {S, T}
 
     if node.is_leaf
-        return Leaf{T}(list[node.label], labels[node.region])
+        featfreq = Tuple(sum(labels[node.region] .== l) for l in list)
+        return Leaf{T, length(list)}(
+            Tuple(list), argmax(featfreq), featfreq, length(node.region))
     else
         left = _convert(node.l, list, labels)
         right = _convert(node.r, list, labels)
-        return Node{S, T}(node.feature, node.threshold, left, right)
+        return Node{S, T, length(list)}(
+            node.feature, node.threshold, left, right)
     end
 end
 
@@ -233,7 +236,10 @@ function prune_tree(
             if !isempty(fi)
                 update_pruned_impurity!(tree, fi, ntt, loss)
             end
-            return Leaf{T}(majority, all_labels)
+            features = Tuple(unique(all_labels))
+            featfreq = Tuple(sum(all_labels .== f) for f in features)
+            return Leaf{T}(features, argmax(featfreq),
+                           featfreq, length(all_labels))
         else
             return tree
         end
@@ -268,7 +274,7 @@ function prune_tree(
 end
 
 
-apply_tree(leaf::Leaf, feature::AbstractVector) = leaf.majority
+apply_tree(leaf::Leaf, feature::AbstractVector) = leaf.features[leaf.majority]
 apply_tree(
     tree::Root{S, T},
     features::AbstractVector{S}
@@ -314,7 +320,7 @@ of the output matrix.
 apply_tree_proba(tree::Root{S, T}, features::AbstractVector{S}, labels) where {S, T} =
     apply_tree_proba(tree.node, features, labels)
 apply_tree_proba(leaf::Leaf{T}, features::AbstractVector{S}, labels) where {S, T} =
-    compute_probabilities(labels, leaf.values)
+    collect(leaf.values ./ leaf.total)
 
 function apply_tree_proba(
     tree::Node{S, T},

src/regression/main.jl

@@ -2,7 +2,10 @@ include("tree.jl")
 
 function _convert(node::treeregressor.NodeMeta{S}, labels::Array{T}) where {S, T <: Float64}
     if node.is_leaf
-        return Leaf{T}(node.label, labels[node.region])
+        features = Tuple(unique(labels))
+        featfreq = Tuple(sum(labels[node.region] .== f) for f in features)
+        return Leaf{T, length(features)}(
+            features, argmax(featfreq), featfreq, length(node.region))
     else
         left = _convert(node.l, labels)
         right = _convert(node.r, labels)