[PYTHON] Visualize scikit-learn decision trees with Plotly's Treemap

this

Do this: angel:

manner

Visualize with plot_tree

First, create a decision tree normally and visualize it with plot_tree.

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.datasets import load_breast_cancer
from sklearn.tree import DecisionTreeClassifier, plot_tree

data = load_breast_cancer()
X, y = data['data'], data['target']
feature_names = data['feature_names']

model = DecisionTreeClassifier(criterion='entropy').fit(X,y)

plt.figure(figsize=(12, 4), dpi=200)
plot_tree(model, feature_names=feature_names, filled=True)
plt.show()

it is normal.

Visualize with Treemap

Next, let's disassemble the decision tree we made and visualize it with Treemap.

import plotly.graph_objects as go

labels = [''] * model.tree_.node_count
parents = [''] * model.tree_.node_count
labels[0] = 'root'
for i, (f, t, l, r) in enumerate(zip(
    model.tree_.feature,
    model.tree_.threshold,
    model.tree_.children_left,
    model.tree_.children_right,
)):
    if l != r:
        labels[l] = f'{feature_names[f]} <= {t:g}'
        labels[r] = f'{feature_names[f]} > {t:g}'
        parents[l] = parents[r] = labels[i]

fig = go.Figure(go.Treemap(
    branchvalues='total',
    labels=labels,
    parents=parents,
    values=model.tree_.n_node_samples,
    textinfo='label+value+percent root',
    marker=dict(colors=model.tree_.impurity),
    customdata=list(map(str, model.tree_.value)),
    hovertemplate='''
<b>%{label}</b><br>
impurity: %{color}<br>
samples: %{value} (%{percentRoot:%.2f})<br>
value: %{customdata}'''
))
fig.show()

Nodes that are crushed and invisible can be seen by clicking on the sector.

the end

Unlike plot_tree, you can't color each class, so it can be difficult to use without binary classification or regression: sweat_smile:

References

• https://scikit-learn.org/stable/auto_examples/tree/plot_unveil_tree_structure.html