How can I train a decision tree constrained to have number of decision nodes = tree depth?

Question

In order to make a classifier dead easy to understand/interpret, I want to classify tabular data (with n columns) according to a set of nested rules, with the constraint that the number of decision nodes is equal to the depth of the tree. Given a vector x with n components, the classification logic will thus look like:

def classify(x):
  if x[0] < t_0:
    if x[1] < t_1:
      if x[2] < t_2:
        if x[3] < t_3:
          ...
            if x[m-1] < t_m_minus_1:
                return a_m_minus_1
            else:
                return a_m_minus_2
          ...
        else:
          return a_3
      else:
        return a_2 
    else:
      return a_1
  else:
    return a_0

with m <= n, so that the none of else statements can contain an if statement nested in it. As a consequence the number m of decision nodes will be equal to the depth of the tree.

Graphically (in the case m = 3), this will look like:

A Sankey diagram can also help visualize this:

Incidentally, I would like to use this classifier as a multi-class classifier (so not necessarily binary).

Also, ideally I'd like "else" leaf nodes to have very low Gini impurity index. Each node should be split according to a condition on a single feature only.

Is there a way to train a decision tree in scikit-learn while enforcing this constraint? What could be another library/approach to optimize such a classifier? (I would avoid coding a greedy algorithm from scratch if that's reinventing the wheel).

Answer 1

The structure you want seems to be expressable with an ordered series of if, else if, ... statements. This is a common structure for interpretable models, often called a Rule List, or Decision List. It is discussed in chapter "Rules" in the book Interpretable Machine Learning by Christoph Molnar.

There are several Python libraries that implements learning of a Rule List.

The imodels library in the submodule imodels.rule_list implements many methods that can produce Rule List models, such as Optimal rule list (CORELS), Bayesian rule list, Greedy rule list and OneR rule list.

The GreedyRuleListClassifier is probably the closest to your intent, the authors call it "like a decision tree that only ever splits going left".

OneR only considers one feature in total, which is an additional restriction.

The Optimal Rule list and Bayesian rule lists requires discretizing continuous features. This can for example be done using quantile binning, or another model to find relevant/candidate breakpoints. So it is considerably more involved, but may lead to better decision lists, especially if using the probabilistic outputs.

Example code for a GreedyRuleListClassifier may go as follows:

import pandas
import sklearn
import sklearn.datasets
from sklearn.model_selection import train_test_split

from imodels import GreedyRuleListClassifier

X, Y = sklearn.datasets.load_breast_cancer(as_frame=True, return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, Y, test_size = 0.3, random_state=4)

# NOTE: fitting sometimes fails with an Exception, or gives model with very bad performance
# Should attempt multiple fits and keep the best one, as estimated per a validation set
model = GreedyRuleListClassifier(max_depth=10)
model.fit(X_train, y_train, feature_names=X_train.columns)

y_pred = model.predict(X_test)
from sklearn.metrics import accuracy_score
score = accuracy_score(y_test.values,y_pred)
print('Accuracy:\n', score)

print('Rule list:\n')
print(model)

Should output something like

Accuracy:
 0.631578947368421
Rule list:

mean 0.603 (398 pts)
if worst area >= 869.3 then 0.908 (262 pts)
mean 0.015 (136 pts)
if worst texture >= 16.67 then 1.0 (1 pts)
mean 0.007 (135 pts)
if area error >= 22.18 then 0.2 (5 pts)
mean 0 (130 pts)

Here is a quick attempt at visualizing this with a Sankey diagram.

def plot_decision_rules_sankey(ax, rules):
    # https://matplotlib.org/stable/api/sankey_api.html
    # TODO: read the arguments
    

    from matplotlib.sankey import Sankey

    def format_rule(r):
        op = '>='
        s = f"{r['col']}\n {op} {r['cutoff']}\np={r['val_right']:.2f}"
        return s
    
    df = pandas.DataFrame.from_records(model.rules_)
    print(df)
    df = df.dropna()
    df['label'] = df.apply(format_rule, axis=1)
    df['orientation'] = [1] * len(df)
    df['out'] = df['num_pts'] / df['num_pts'].sum()

    p = Sankey(ax=ax,
       margin=0.0,
       format='',
       flows=[0.0] + list(df['out'] * -1),
       labels=['Input'] + list(df['label']),
       orientations=[0] + list(df['orientation']),
    ).finish()

    ax.axis('off')

from matplotlib import pyplot as plt
fig, ax = plt.subplots(1, figsize=(8, 6))
plot_decision_rules_sankey(ax, model.rules_)
fig.tight_layout()
fig.savefig('decision-rules-sankey.png')