Making sklearn's decision trees easier to traverse

Question

scikit-learn's decision tree structure is difficult for me to navigate. I would prefer to have functionality like tree.left, tree.right, tree.parent, etc.

I wrote the following:

class SkNode:
    """Rooted binary tree class with attributes to match sklearn's decision tree object."""
    def __init__(self, label = None):
        self.label = label 
        self.left  = None #left child node
        self.right = None #right child node
        self.feature = None # the feature we are splitting on
        self.threshold = None # the threshold value for the feature we are splitting on
        self.parent = None # the parent node
        self.leaf = False # whether or not this is a leaf node
        self.i_am_left = False # if this is the left child of the parent
        self.i_am_right = False # if this is the right child of the parent
        self.pred = None # For a leaf node, this will be the index of the majority class
        self.n_classes = None 
        self.n_features = None 

    def find_constraints(self):
        """Recursively traverses the tree from child to parent and collects the constraints on each feature"""
        mins = [float('-inf') for i in range(self.n_features)] # list of mins by feature index
        maxs = [float('inf') for i in range(self.n_features)] # list of maxs by feature index
        if self.parent:
            (mins, maxs) = self.parent.find_constraints() # we start with the mins and maxs of our parent
            i = self.parent.feature
            if self.i_am_left:
                maxs[i] = self.parent.threshold 
            if self.i_am_right:
                mins[i] = self.parent.threshold
        return (mins, maxs)


def traversable_nodes(sk_tree_fit_model_object):
    """ Converts a fit sklearn model object into a dictionary of SkNodes"""
    sk_tree = sk_tree_fit_model_object.tree_
    nodes = {i: SkNode(label = i) for i in range(sk_tree.node_count)}
    n_classes = sk_tree.n_classes[0]
    n_features = sk_tree.n_features
    for i in range(sk_tree.node_count):
        nodes[i].n_classes = n_classes
        nodes[i].n_features = n_features
        if sk_tree.feature[i] >= 0:
            nodes[i].left = nodes[sk_tree.children_left[i]]
            nodes[i].left.parent = nodes[i]
            nodes[i].left.i_am_left = True

            nodes[i].right = nodes[sk_tree.children_right[i]]
            nodes[i].right.parent = nodes[i]
            nodes[i].right.i_am_right = True

            nodes[i].threshold  = sk_tree.threshold[i]

            nodes[i].feature = sk_tree.feature[i]
            
        if sk_tree.feature[i] < 0:
            nodes[i].leaf = True
            nodes[i].pred = sk_tree.value.reshape((-1,n_classes)).argmax(axis = 1)[i]
            
    return nodes

The code works. Here is an example:

import matplotlib.pyplot as plt
import numpy as np
import matplotlib.patches as patches

n_samples = 300
X = np.random.uniform(low = -2, high = 2, size=(n_samples,2))

y = np.zeros(n_samples)
y[((X[:,0]-1)**2 +X[:,1]**2 > 1) & (X[:,0]**2 + X[:,1]**2 < 4)] = 1
y[((X[:,0]-1)**2 +X[:,1]**2 < 1)] = 2

tree_clf = DecisionTreeClassifier(max_depth  = 8)

tree_clf.fit(X, y)

plt.figure(figsize=(8,8))

plt.scatter(X[y == 0,0], X[y == 0,1], c = 'red', marker = 'x', label = "0")
plt.scatter(X[y == 1,0], X[y == 1,1], c = 'orange', marker = 'v',label = "1")
plt.scatter(X[y == 2,0], X[y == 2,1], c = 'blue', label = "2")
plt.gca().add_patch(plt.Circle((1, 0), 1, color = 'k', fill = False))
plt.gca().add_patch(plt.Circle((0, 0), 2, color = 'k', fill = False))

nodes = traversable_nodes(tree_clf)

for node in nodes.values():
    if node.leaf:
        (xmin,ymin),(xmax,ymax) = node.find_constraints()
        xmin = np.max([xmin,-2])
        xmax = np.min([xmax,2])
        ymin = np.max([ymin,-2])
        ymax = np.min([ymax,2])
        if node.pred == 0:
            plt.gca().add_patch(patches.Rectangle((xmin, ymin), xmax-xmin, ymax-ymin, linewidth=2, edgecolor='k', facecolor = 'red', zorder = -2, alpha = 0.4))
        if node.pred == 1:
            plt.gca().add_patch(patches.Rectangle((xmin, ymin), xmax-xmin, ymax-ymin, linewidth=2, edgecolor='k', facecolor = 'orange', zorder = -2, alpha = 0.4))
        if node.pred == 2:
            plt.gca().add_patch(patches.Rectangle((xmin, ymin), xmax-xmin, ymax-ymin, linewidth=2, edgecolor='k', facecolor = 'blue', zorder = -2, alpha = 0.4))

plt.title(f"Maximum Depth of 8")
plt.xlabel("$x_1$", fontsize=12)
plt.ylabel("$x_2$", fontsize=12)
plt.legend(fontsize=12)

plt.show()

produces this picture:

I am looking for advice on how to improve my SkNode class and transversable_nodes function. Any advice you can give will be appreciated.

Answer 1

In traversable_nodes, the 2 if statements should be combined into an if/else. Change:

    if sk_tree.feature[i] >= 0:
        ...
    if sk_tree.feature[i] < 0:

to:

    if sk_tree.feature[i] >= 0:
        ...
    else:

Not only is it simpler, but it is more efficient because there is no need for the second check if the first is true, and this is inside a loop.

In the example code, you can similarly replace the 3 if statements with an if/elif:

    if node.pred == 0:
        ...
    if node.pred == 1:
        ...
    if node.pred == 2:

to:

    if node.pred == 0:
        ...
    elif node.pred == 1:
        ...
    elif node.pred == 2:

DRY

This long line is repeated 3 times, except for facecolor:

plt.gca().add_patch(patches.Rectangle((xmin, ymin), xmax-xmin, ymax-ymin, linewidth=2, edgecolor='k', facecolor = 'red', zorder = -2, alpha = 0.4))

You could push that into a function:

def add_patch(face_color):
    plt.gca().add_patch(patches.Rectangle(
        (xmin, ymin), xmax-xmin, ymax-ymin, linewidth=2, edgecolor='k', facecolor = face_color, zorder = -2, alpha = 0.4)
    )