draw-scikit-tree 0.1.3

A simple interface to modify scikit-learn's generated DOT representation of a decision tree.

DrawScikitTree

A simple interface to modify scikit-learn's generated DOT string representation of a trained decision tree. Some basic function include changing the shape and color of each node, and tracing the decision paths taken for a test sample.

Installation

pip install draw-scikit-tree

Basic usage

Using the iris dataset as the classical example.

from sklearn.datasets import load_iris
from sklearn import tree
iris = load_iris()
X, y = iris.data, iris.target
clf = tree.DecisionTreeClassifier()
clf = clf.fit(X, y)

Next, use the trained classifier to initialize the TreeGraph object.

from DrawScikitTree import TreeGraph
treeGraph = TreeGraph(clf, impurity=False, label="none", fontname="Arial")

To trace the decisions paths taken for some test samples, use the .trace_paths(X_sample) function.

import numpy as np
import graphviz

# Get some random samples
random_indices = np.random.randint(X.shape[0], size=5)
X_sample = X[random_indices, :]

# Setting verbose=True will print out the decision paths for each sample
treeGraph.trace_paths(X_sample, color="red", verbose=True)

# Displaying the newly modified tree
new_dot_data = treeGraph.export()
graph = graphviz.Source(new_dot_data)
display(graph)

For more examples check out the examples.