How to plot a decision tree (using a graphics package)

[roland-KA]

Is there a possibility to plot a decision tree using Plots.jl (or some other graphics package)?

I'm using MLJ and the only means to visualize a decision tree seems to be report(mach).print_tree(n) where mach is the trained machine.

If there is no such possibility: How can I access the tree (data structure) directly in MLJ?

[roland-KA]

I think I've solved most of the questions above: Using AbstractTrees.jl and GraphRecipes.jl it's relatively easy to implement.

• The decision tree can be accessed via fitted_params(mach).tree.
• Then the AbstractTrees functions children and print node have to be implemented:

using AbstractTrees

function AbstractTrees.children(node::DecisionTree.Node)
	return(node.left, node.right)
end

function AbstractTrees.printnode(io::IO, node::DecisionTree.Node)
	print(io, "ID: ", node.featid, " - ", node.featval)
end

function AbstractTrees.printnode(io::IO, leaf::DecisionTree.Leaf)
	print(io, "maj: ", leaf.majority, " - vals: ", length(leaf.values))
end

• Finally GraphRecipes can be used together with Plots (dtree is the decision tree to be plotted):

using GraphicRecipes

plot(TreePlot(dtree), method = :tree, nodeshape = :ellipse)

Unfortunately DecisionTree.Node stores only the id of the feature which is used for a split (featid). It would be nice, if the feature name could also be shown. Within the DecisionTree package exists an array feature_names with these names. But I didn't find a way to access it. How can this be done (in printnode)?

[Rahulub3r]

Hi @roland-KA, I have created a function to plot decision tree using CairoMakie (https://github.com/Rahulub3r/Julia-ML-Utils/blob/main/decisionTreeUtils.jl). Although this is a minimum working example, I think it can be changed to be production material for integrating into the package. I am interested in contributing. Let me know if this looks good.

[roland-KA]

Hi @Rahulub3r , thank you for your efforts!
Could you perhaps give a working example on how to call drawTree? I.e. which initial values should be used to draw a tree from its root down to the leafs?

[Rahulub3r]

Sure. An MWE is as follows. Suggestions are appreciated.

using MLJ
using DecisionTree
using CairoMakie
using Random

X, y = make_blobs(300;rng=MersenneTwister(1234))

dtc = @load DecisionTreeClassifier pkg=DecisionTree verbosity=0
dtc_model = dtc(min_purity_increase=0.005, min_samples_leaf=1, min_samples_split=2, max_depth=3)
dtc_mach = machine(dtc_model, X, y)
MLJ.fit!(dtc_mach)
x = fitted_params(dtc_mach)
#print_tree(x.tree)

f = Figure(;resolution=(1000, 800))
ax1 = Axis(f[1,1])
drawTree(x.tree, x.encoding, ax1; feature_names=["X1", "X2"], 
        nodetextsize=20, nodetextcolor=:black, nodewth=12,
        linetextsize=13, leaftextsize=13, leafwth=4)
hidespines!(ax1)
hidedecorations!(ax1)
f

[roland-KA]

Thank's, that looks really good! 👍

[ablaom]

Very cool! I wonder what the best way to integrate this contribution might be. Be great for MLJ users to be able to do this (without adding Makie as dependency to existing MLJ packages).

@Rahulub3r Any interest in working helping out with visualisation in MLJ more generally?

Probably missing something but why is it necessary to specify feature names? Are they not part of x.encoding?

[Rahulub3r]

Yes, I am interested. However, I am not sure how we would implement them without Makie.
RE feature names, if you do not specify them, they will be shown as Feature 1, Feature 2, .. in the plot, but if you specify the names, you will see them as shown in the plot above. x.encoding contains the class label information and not the features.

[ablaom]

x.encoding contains the class label information and not the features.

Yes, of course you right 👍🏾 . Perhaps we could get the MLJ interface to expose the feature names (as we do in, eg, the MLJGLMInterface).

[roland-KA]

As far as I can see, we have now three options to visualize a decision tree:

1. Using GraphRecipes (with Plots in the background and the decision tree 'extended' to an AbstractTree) as described in my post above.
2. The code developed by @Rahulub3r, which implements its own tree layout algorithm using Makie.jl.
3. Using Graphs (with GraphMakie the background).

No. 2 is without doubt visually the most beautiful solution. But I think we should also consider some software engineering aspects:

• Maintenance: I think it is better to build on existing efforts to keep the own efforts for development and maintenance low (our resources are scarce). No. 1 & 3 only need some 'translation' of the decision tree to another data structure and use then ready made layout and plotting algorithms.
• Loose coupling: As @ablaom mentioned, there shouldn't be a direct dependency between MLJ and a graphics package. With option 1 every package capable of plotting an AbstractTree will do the job and with option 3 the same holds for every package capable of plotting a Graph. So we have a decoupling using these abstract data types.
  And of course it would be helpful, if MLJ could provide the class labels as well as the feature names through a clear interface. That's a good idea @ablaom!

I've been digging a bit deeper into options 1 & 3 recently, so some more comments on that:

Option 1:
Using class labels and attribute names (from 'somewhere', ideally from an MLJ interface 😊), trees with all that information can be plotted using the means stated above.

BTW: I've started a three-part tutorial on 'Towards Data Science' on how Julia and its ecosystem can be used for ML (for an audience that doesn't know Julia). The first part has been published last week. The third part uses exactly no. 1 as an example to show how easy it is in Julia to 'glue' several packages together and create a new solution with only little coding. So there will be an extensive description on how to do it.

On that way, I identified an issue with the TreePlot recipe (JuliaPlots/GraphRecipes.jl#172), which has to be resolved, before we can use it to produce correct tree plots.

Option 3:
In order to get a 'beautiful' tree plot, some extra work is needed as the standard output from GraphMakie is a bit basic and doesn't offer more out-of-the-box as I a learned here: MakieOrg/GraphMakie.jl#57

I can provide a code example on how to apply this option next week (I'm a bit busy this week).

[ablaom]

@roland-KA Very nice summary of the issues and best options. I agree that while @Rahulub3r code is a lovely bit of work, a more generic solution is preferred (1 or 3).

It probably makes sense to have a separate "MLJPlotRecipes" package for plotting things that are very specific to MLJ (we have, for example, recipe in MLJTuning to plot outcomes of hyper-parameter optimization). But this wouldn't fall into this category.

I've opened JuliaAI/MLJDecisionTreeInterface.jl#13. This should be pretty easy (PR welcome 😄 ).

Makie.jl or Plots.jl? That's a very difficult question. I'm waiting for someone to convince me one way or the other 😉 . Of course one could do both, starting with Plots.jl which seems more ready to roll.

See also the MLJ visualization project in this GSoC list.

[roland-KA]

I think a separate "MLJPlotRecipes" package is a good idea, as there are surely more models that could be visualized (and it would be another advantage of MLJ over similar ML-packages in general).

As a first step, an implementation for Plots.jl would be preferable in my opinion too, because:

• Plots.jl is more mature in comparison. In Makie.jl there are still quite a few loose ends and a lot of changes going on.
• The maintainer of GraphMakie.jl informed me, that he is planning to extend that package with the functionality we need (see: Text boxes instead of text labels MakieOrg/GraphMakie.jl#57). So it hasn't to implemented for decision trees in MLJPlotRecipes.

... and if time permits, I will have a look at JuliaAI/MLJDecisionTreeInterface.jl#13 😊

[roland-KA]

And as promised, here the code to plot a decision tree using Graphs and GraphMakie. This variant needs a bit more coding than GraphRecipies & Plots since Graphs expects the tree in a breadth-first order whereas the structures of the decision tree are better suited for a depth-first traversal.

So I'm doing basically a depth-first traversal of the decision tree but I collect the nodes (as well as the information for the labels) in a breadth-first order. Sounds a bit strange (and is a bit strange 😀), but can be done with relatively little effort.

The trick is the use of the counter dictionary below, which has a counter for each level of the (binary) tree and delivers a numbering of the nodes in breadth-first order (1st level starts with 1, 2nd level with 2, 3rd level with 4 etc.).

The function traverse_tree visits each node and adds on this visit label information to label and adds the edge to its left and right child to the graph g (in the call to add_node). We assume that the decision tree is in tree and has been created beforehand.

using Graph, GraphMakie, CairoMakie

depth = DecisionTree.depth(tree)
g = SimpleDiGraph(2^(depth+1)-1)  # graph (for building the tree) of size given as an argument
label = Dict{Int, String}()  # for collecting label information
counter = Dict([i => 2^i for i in 0:depth])

function traverse_tree(tree::DecisionTree.Node, level::Int, cnt::Dict{Int, Int}, g::SimpleDiGraph, label::Dict{Int, String})
	label[cnt[level]] = "ID = " * string(tree.featid) * " - " * string(tree.featval)
	add_node!(g, tree.left, level, cnt, label) 	# left child
	add_node!(g, tree.right, level, cnt, label) # right child
end

traverse_tree(tree::DecisionTree.Leaf, level::Int, cnt::Dict{Int, Int}, g::SimpleDiGraph, label::Dict{Int, String}) = 
	label[cnt[level]] = "maj = " * string(tree.majority)

function add_node!(g::SimpleDiGraph, node::Union{DecisionTree.Node, DecisionTree.Leaf}, level::Int, cnt::Dict{Int, Int}, label::Dict{Int, String})
	add_edge!(g, cnt[level], cnt[level+1])
	traverse_tree(node, level + 1, cnt, g, label)
	cnt[level+1] += 1
end

The function traverse_tree has to be applied to the initial data structures defined above and a level of 0. Afterwards we have a tree structure in g similar to the decision tree and all label information in label. The keys of label give us the (breadth-first) order we need for plotting. Therefore we have to extract it in that order (call to sort) and then everything can be plotted using graphplot.

traverse_tree(tree, 0, counter, g, label)
tree_labels = collect(values(sort(label)))
f, ax, p = graphplot(g, layout = layout, nlabels = tree_labels, nlabels_attr=(;justification = :center, align = (:center, :top), color = :blue))

The resulting tree plot is the one depicted in MakieOrg/GraphMakie.jl#57.