Merge pull request #173 from rikhuijzer/rh/tree-digits

Round digits in `print_tree`

Author: ablaom

src/DecisionTree.jl

@@ -80,14 +80,18 @@ length(ensemble::Ensemble) = length(ensemble.trees)
 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)")
+end
 function print_tree(leaf::Leaf, depth=-1, indent=0; feature_names=nothing)
-    matches = findall(leaf.values .== leaf.majority)
-    ratio = string(length(matches)) * "/" * string(length(leaf.values))
-    println("$(leaf.majority) : $(ratio)")
+    return print_tree(stdout, leaf, depth, indent; feature_names=feature_names)
 end
 
+
 """
-       print_tree(tree::Node, depth=-1, indent=0; feature_names=nothing)
+    print_tree([io::IO,] tree::Node, depth=-1, indent=0; sigdigits=4, feature_names=nothing)
 
 Print a textual visualization of the specified `tree`. For example, if
 for some input pattern the value of "Feature 3" is "-30" and the value
@@ -100,7 +104,7 @@ non-majority classes are not shown.
 
 # Example output:
 ```
-Feature 3 < -28.15
+Feature 3 < -28.15 ?
 ├─ Feature 2 < -161.0 ?
    ├─ 5 : 842/3650
    └─ 7 : 2493/10555
@@ -113,20 +117,24 @@ To facilitate visualisation of trees using third party packages, a `DecisionTree
 `DecisionTree.Node` object can be wrapped to obtain a tree structure implementing the
 AbstractTrees.jl interface. See  [`wrap`](@ref)` for details.
 """
-function print_tree(tree::Node, depth=-1, indent=0; feature_names=nothing)
+function print_tree(io::IO, tree::Node, depth=-1, indent=0; sigdigits=2, feature_names=nothing)
     if depth == indent
-        println()
+        println(io)
         return
     end
+    featval = round(tree.featval; sigdigits=sigdigits)
     if feature_names === nothing
-        println("Feature $(tree.featid) < $(tree.featval)")
+        println(io, "Feature $(tree.featid) < $featval ?")
     else
-        println("Feature $(tree.featid): \"$(feature_names[tree.featid])\" < $(tree.featval)")
+        println(io, "Feature $(tree.featid): \"$(feature_names[tree.featid])\" < $featval ?")
     end
-    print("    " ^ indent * "├─ ")
-    print_tree(tree.left, depth, indent + 1; feature_names = feature_names)
-    print("    " ^ indent * "└─ ")
-    print_tree(tree.right, depth, indent + 1; feature_names = feature_names)
+    print(io, "    " ^ indent * "├─ ")
+    print_tree(io, tree.left, depth, indent + 1; feature_names=feature_names)
+    print(io, "    " ^ indent * "└─ ")
+    print_tree(io, tree.right, depth, indent + 1; feature_names=feature_names)
+end
+function print_tree(tree::Node, depth=-1, indent=0; sigdigits=2, feature_names=nothing)
+    return print_tree(stdout, tree, depth, indent; sigdigits=sigdigits, feature_names=feature_names)
 end
 
 function show(io::IO, leaf::Leaf)

src/scikitlearnAPI.jl

@@ -386,5 +386,7 @@ length(dt::DecisionTreeClassifier)  = length(dt.root)
 length(dt::DecisionTreeRegressor)   = length(dt.root)
 
 print_tree(dt::DecisionTreeClassifier, depth=-1; kwargs...) = print_tree(dt.root, depth; kwargs...)
-print_tree(dt::DecisionTreeRegressor,  depth=-1; kwargs...) = print_tree(dt.root, depth; kwargs...)
-print_tree(n::Nothing, depth=-1; kwargs...)                 = show(n)
+print_tree(io::IO, dt::DecisionTreeClassifier, depth=-1; kwargs...) = print_tree(io, dt.root, depth; kwargs...)
+print_tree(dt::DecisionTreeRegressor, depth=-1; kwargs...) = print_tree(dt.root, depth; kwargs...)
+print_tree(io::IO, dt::DecisionTreeRegressor,  depth=-1; kwargs...) = print_tree(io, dt.root, depth; kwargs...)
+print_tree(n::Nothing, depth=-1; kwargs...) = show(n)

test/classification/random.jl

@@ -3,7 +3,7 @@
 
 Random.seed!(16)
 
-n,m = 10^3, 5;
+n, m = 10^3, 5;
 features = rand(n,m);
 weights = rand(-1:1,m);
 labels = round.(Int, features * weights);
@@ -15,7 +15,19 @@ preds = apply_tree(model, round.(Int, features))
 max_depth = 3
 model = build_tree(labels, features, 0, max_depth)
 @test depth(model) == max_depth
-print_tree(model, 3)
+
+io = IOBuffer()
+print_tree(io, model, 3)
+text = String(take!(io))
+println()
+print(text)
+println()
+
+# Read the regex as: many not arrow left followed by an arrow left, a space, some numbers and
+# a dot and a space and question mark.
+rx = r"[^<]*< [0-9\.]* ?"
+matches = eachmatch(rx, text)
+@test !isempty(matches)
 
 model = build_tree(labels, features)
 preds = apply_tree(model, features)