Turn Leaf struct into a frequency map

Author: tecosaur

src/DecisionTree.jl

@@ -28,16 +28,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}}
@@ -46,11 +48,15 @@ struct Ensemble{S, T}
     trees :: Vector{LeafOrNode{S, T}}
 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(String) = ""
-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))
 promote_rule(::Type{Node{S, T}}, ::Type{Leaf{T}}) where {S, T} = Node{S, T}
 promote_rule(::Type{Leaf{T}}, ::Type{Node{S, T}}) where {S, T} = Node{S, T}
 
@@ -81,9 +87,8 @@ depth(leaf::Leaf) = 0
 depth(tree::Node) = 1 + max(depth(tree.left), depth(tree.right))
 
 function print_tree(io::IO, leaf::Leaf, depth=-1, indent=0; 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; feature_names=nothing)
     return print_tree(stdout, leaf, depth, indent; feature_names=feature_names)
@@ -139,8 +144,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,7 +41,10 @@ function _convert(
     ) where {S, T}
 
     if node.is_leaf
-        return Leaf{T}(list[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, list, labels)
         right = _convert(node.r, list, labels)
@@ -120,7 +123,10 @@ function prune_tree(tree::LeafOrNode{S, T}, purity_thresh=1.0) where {S, T}
             matches = findall(all_labels .== majority)
             purity = length(matches) / length(all_labels)
             if purity >= purity_thresh
-                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
@@ -139,7 +145,8 @@ function prune_tree(tree::LeafOrNode{S, T}, purity_thresh=1.0) where {S, T}
 end
 
 
-apply_tree(leaf::Leaf, feature::AbstractVector) = leaf.majority
+apply_tree(leaf::Leaf{T}, feature::AbstractVector{S}) where {S, T} =
+    leaf.features[leaf.majority]
 
 function apply_tree(tree::Node{S, T}, features::AbstractVector{S}) where {S, T}
     if tree.featid == 0
@@ -173,7 +180,7 @@ n_labels` matrix of probabilities, each row summing up to 1.
 (eg. ["versicolor", "virginica", "setosa"]). It specifies the column ordering
 of the output matrix. """
 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}, features::AbstractVector{S}, labels) where {S, T}
     if tree.featval === nothing

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)