Certifiably optimal sparse decision trees, discovered by autoresearch.
FastSmallTreeClassifier returns the tree minimising misclassification rate plus
a penalty per leaf, and certifies that no other tree on the same binarized
features scores better. The search behind it lives in solver.py.
Expand source code
"""Certifiably optimal sparse decision trees, discovered by autoresearch.
`FastSmallTreeClassifier` returns the tree minimising misclassification rate plus
a penalty per leaf, and certifies that no other tree on the same binarized
features scores better. The search behind it lives in `solver.py`.
"""
import warnings
import numpy as np
import pandas as pd
import scipy.sparse
from sklearn.base import BaseEstimator, ClassifierMixin
from sklearn.tree import DecisionTreeClassifier
from sklearn.tree._tree import Tree
from sklearn.utils.validation import check_array
from imodels.tree.optimal_tree.solver import (HAVE_NUMBA, ST_LB, ST_UB, BinaryEncoder,
BitDataset, CompiledOptimizer, TargetEncoder,
TreeClassifier, cluster_rows)
from imodels.util.arguments import (_finite_check_kwarg, check_fit_arguments,
check_predict_X)
from imodels.util.introspection import RuleInspectionMixin
from imodels.util.numba_compile import notify_first_compile
from imodels.tree.optimal_tree import solver as _solver_module
NUMBA_HINT = (
"FastSmallTreeClassifier needs numba, which is not installed. Install it with "
"`pip install numba` (or `pip install imodels[optional]`). The search itself is "
"compiled, and interpreting it is orders of magnitude slower, so there is no "
"pure-Python fallback."
)
def _class_weights(costs: np.ndarray):
"""Per-class weights w when the objective is a reweighting of the counts.
That is the case when a mistake on a row of class j costs w[j] whatever was
predicted (zero diagonal, each column constant off it): the default objective
and ``balance``. Returns None for any other cost matrix.
"""
C = np.asarray(costs, dtype=np.float64)
K = C.shape[0]
if K < 2 or np.any(np.diag(C) != 0):
return None
off = ~np.eye(K, dtype=bool)
w = np.array([C[off[:, j], j][0] for j in range(K)])
if not all(np.all(C[off[:, j], j] == w[j]) for j in range(K)):
return None
return w
def _leaf_value(dist: np.ndarray, costs: np.ndarray, weights, prediction=None) -> np.ndarray:
"""A node's sklearn ``value``: normalised, with ``argmax`` the solver's choice.
``prediction`` is the class the solver certified for a leaf. Two classes can
tie exactly (class counts (1, 2, 11) under ``balance`` give 2/6 and 11/33),
and the solver's ``argmin(costs @ dist)`` and the ``argmax`` here then round
differently, so a near tie is broken toward ``prediction``. A gap larger than
rounding can explain is a real disagreement and raises.
"""
if weights is not None:
v = weights * dist
else:
cost = np.asarray(costs, dtype=np.float64) @ dist
v = cost.max() - cost
if v.sum() <= 0: # an empty node, or one where every label costs the same
v = np.zeros_like(dist)
if prediction is None:
prediction = int(np.argmin(np.asarray(costs, dtype=np.float64) @ dist))
v[int(prediction)] = 1.0
v = v / v.sum()
if prediction is not None and int(np.argmax(v)) != prediction:
top = v.max()
if v[prediction] < top * (1.0 - 1e-9):
raise AssertionError("sklearn leaf value does not pick the certified class")
# a relative margin far above one ulp, so normalising again (as sklearn's
# predict_proba does) cannot merge the two values back into a tie
v[prediction] = top * (1.0 + 1e-12)
v = v / v.sum()
return v
class FastSmallTreeClassifier(RuleInspectionMixin, ClassifierMixin, BaseEstimator):
"""A decision tree that is certifiably optimal for its objective.
Fits the tree minimising
misclassification rate + ``regularization`` * (number of leaves)
over *every* decision tree on the binarized features, and proves that no
other such tree scores better. This is the objective of `GOSDT
<https://arxiv.org/abs/2006.08690>`_ (Lin et al., ICML 2020); the
branch-and-bound search is described in `solver.py` and was developed by
autoresearch (see `Agentic-imodels <https://github.com/csinva/agentic-imodels>`_).
Greedy trees (CART, C4.5) pick each split without regard to the splits below
it, so they can miss the best tree of a given size. This model does not:
when ``optimal_`` is True, the objective it reports is a proven minimum.
Depth is not a parameter. ``regularization`` sets the price of a leaf, and
the optimal tree at that price is as deep as it needs to be. Since a leaf
that does not improve the misclassification rate by ``regularization`` is
never worth adding, the trees stay small.
Parameters
----------
regularization : float, default=0.05
Penalty added to the objective for each leaf, in units of
misclassification rate. Larger values give smaller trees. Values below
``1 / n_samples`` are not useful, since a leaf fixing a single training
point already pays for itself.
time_limit : float or None, default=60.0
Seconds after which the search returns the best tree found so far rather
than continuing. ``optimal_`` is then False and a warning is raised.
``0`` or None means no limit, which can run for a very long time on wide
or continuous data.
balance : bool, default=False
Weigh the classes equally, optimising balanced accuracy rather than
accuracy.
costs : array-like of shape (n_classes, n_classes), default=None
``costs[i, j]`` is the cost of predicting class ``i`` when the truth is
class ``j``. Overrides ``balance``.
memory_limit : int, default=0
Bytes held by the search's node store (the memo of solved subproblems)
above which the search stops, as the time limit does. This is not the
resident memory of the process, which also holds the data, the
workspace and Python itself. ``0`` means no limit. The store grows with
the number of subproblems kept, which is what makes hard instances
expensive.
verbose : bool, default=False
Print search statistics after fitting.
Attributes
----------
optimal_ : bool
Whether optimality was certified. False means a limit was reached first,
so the tree is the best one found rather than a proven optimum.
objective_ : float
Objective value of the returned tree: its misclassification rate plus
``regularization`` times its number of leaves. Under ``balance`` the
rate is the class-weighted (balanced) error, and under ``costs`` it is
the average cost per training row.
lowerbound_, upperbound_ : float
Interval the search closed on the optimal objective; equal when
``optimal_`` is True.
n_leaves_ : int
Number of leaves in the returned tree.
complexity_ : int
Number of splits in the returned tree, the imodels measure of size.
n_binary_features_ : int
Number of binary features the encoder produced from X. A continuous
column becomes one feature per distinct value, which is what makes
continuous data expensive to solve exactly.
estimator_ : sklearn.tree.DecisionTreeClassifier
The fitted tree as an sklearn decision tree, so sklearn's tree tooling
(``plot_tree``, ``export_text``, ``feature_importances_``, dtreeviz)
works on it directly. Each split is placed midway between the training
values on either side of the rule, as sklearn places its own. A new
value strictly between two consecutive training values of a column is
routed by that midpoint, sklearn's convention; the objective is defined
on the training rows, so this changes nothing the certificate covers.
`predict` and `predict_proba` walk these splits in float64 rather than
calling ``estimator_.predict``: sklearn compares X as float32, so on a
column whose training values are closer together than float32 resolves
(unix timestamps, say) ``estimator_.predict`` can route training rows
differently from the certified tree, while `predict` cannot.
tree_ : dict
The same tree as nested dicts of named rules, which printing the model
shows.
time_ : float
Seconds spent in the search.
classes_ : ndarray
Class labels seen during fit.
Notes
-----
The compiled search takes about 20 seconds to build the first time it runs
on a machine. The result is cached on disk, so later processes load it in
about a second. Set ``OPTTREE_NUMBA_CACHE=0`` to skip the cache if its
directory is unwritable.
Examples
--------
>>> import numpy as np
>>> from imodels import FastSmallTreeClassifier
>>> X = np.array([[0, 0], [0, 1], [1, 0], [1, 1]] * 8)
>>> y = (X[:, 0] == X[:, 1]).astype(int)
>>> model = FastSmallTreeClassifier(regularization=0.01).fit(X, y)
>>> model.optimal_
True
"""
def __init__(self, regularization: float = 0.05, time_limit: float = 60.0,
balance: bool = False, costs=None, memory_limit: int = 0,
verbose: bool = False):
self.regularization = regularization
self.time_limit = time_limit
self.balance = balance
self.costs = costs
self.memory_limit = memory_limit
self.verbose = verbose
def fit(self, X, y, feature_names=None):
"""Fit the optimal tree.
Parameters
----------
X : array-like of shape (n_samples, n_features)
y : array-like of shape (n_samples,)
feature_names : list of str, optional
Names used when printing the model; taken from X's columns when it
is a DataFrame.
Returns
-------
self
"""
if not HAVE_NUMBA:
raise ImportError(NUMBA_HINT)
notify_first_compile(_solver_module.__file__, "FastSmallTreeClassifier", "about 20 seconds",
_solver_module.NUMBA_CACHE, "OPTTREE_NUMBA_CACHE")
time_limit, memory_limit = self._check_limits()
X, y, feature_names = check_fit_arguments(self, X, y, feature_names)
X = self._as_float(X)
# the certificate rests on bounds of the form "a tree with a leaves costs at least
# a * regularization" and "a tree's loss cannot fall below zero", which need a
# nonnegative penalty and nonnegative, finite costs
lam = float(self.regularization)
if not np.isfinite(lam) or lam < 0.0:
raise ValueError(f"regularization must be a finite number >= 0, got {self.regularization!r}")
if self.costs is not None:
C = np.asarray(self.costs, dtype=np.float64)
if not np.all(np.isfinite(C)) or np.any(C < 0.0):
raise ValueError("costs must be finite and nonnegative")
frame = pd.DataFrame(X, columns=list(feature_names))
self.encoder_ = BinaryEncoder().fit(frame)
Xb = self.encoder_.transform(frame)
self.target_encoder_ = TargetEncoder().fit(y)
y_idx = self.target_encoder_.transform(y)
# the objective does not depend on the row order, but the kernels skip the
# empty words of a capture mask, so rows are clustered to keep the capture
# sets of deep nodes contiguous
perm = cluster_rows(Xb, y_idx, len(self.classes_), self.encoder_.groups)
data = BitDataset(Xb[perm], y_idx[perm], len(self.classes_),
costs=self.costs, balance=self.balance)
self.n_binary_features_ = data.m
# n_jobs=1: the search has a parallel phase, but a library model should not
# take every core of the caller's machine without being asked
opt = CompiledOptimizer(data, self.regularization, groups=self.encoder_.groups,
time_limit=time_limit, memory_limit=memory_limit,
verbose=self.verbose, n_jobs=1)
root = opt.run() # a row index into the compiled engine's node store
self.optimal_ = opt.optimal
self.stop_reason_ = opt.stop_reason
self.time_ = opt.elapsed
self.iterations_ = opt.iterations
self.lowerbound_ = float(opt.st[ST_LB][root])
self.upperbound_ = float(opt.st[ST_UB][root])
if not self.optimal_:
warnings.warn(
f"{self.stop_reason_} limit reached before optimality was certified; "
"returning the best tree found. Raise time_limit, or raise "
"regularization to search over smaller trees.", RuntimeWarning)
self.tree_ = self._decode(opt.extract(root), data)
opt.release() # the node store is large and is not needed past extraction
certified = TreeClassifier(self.tree_)
self.objective_ = certified.risk()
self.n_leaves_ = certified.leaves()
self.complexity_ = max(self.n_leaves_ - 1, 0)
self.estimator_, self._node_proba_ = self._to_sklearn(self.tree_, frame, data.costs)
self._node_pred_ = self._node_predictions(self.tree_)
self._check_sklearn_tree(certified, frame)
if self.verbose:
print(f"objective {self.objective_:.6g} in {self.time_:.3g}s "
f"({self.iterations_} subproblems, optimal={self.optimal_})")
return self
def _check_limits(self):
"""``time_limit`` and ``memory_limit`` as the search takes them (0 = none)."""
time_limit = self.time_limit
if time_limit is None:
time_limit = 0.0
try:
time_limit = float(time_limit)
except (TypeError, ValueError):
raise ValueError(f"time_limit must be a number >= 0 or None, got {self.time_limit!r}")
if not np.isfinite(time_limit) or time_limit < 0.0:
raise ValueError(f"time_limit must be a finite number >= 0 (0 or None for no "
f"limit), got {self.time_limit!r}")
memory_limit = self.memory_limit
if memory_limit is None:
memory_limit = 0
try:
memory_limit = float(memory_limit)
except (TypeError, ValueError):
raise ValueError(f"memory_limit must be a number of bytes >= 0, got {self.memory_limit!r}")
if not np.isfinite(memory_limit) or memory_limit < 0:
raise ValueError(f"memory_limit must be a finite number of bytes >= 0 (0 for no "
f"limit), got {self.memory_limit!r}")
return time_limit, int(memory_limit)
@staticmethod
def _as_float(X):
"""X as a float64 array, or a clear error before the search starts.
``check_fit_arguments`` leaves the object array that a DataFrame with a
pandas ``Categorical`` column becomes; numeric categories are kept as
their values, and anything else cannot be thresholded.
"""
if scipy.sparse.issparse(X):
X = X.toarray()
try:
return np.asarray(X, dtype=np.float64)
except (TypeError, ValueError):
raise ValueError(
"FastSmallTreeClassifier needs numeric features, but X has a column "
"that cannot be converted to float (for example a pandas Categorical "
"of strings). Encode such columns as numbers first, e.g. with "
"pd.get_dummies or OrdinalEncoder.") from None
def _decode(self, node: dict, data: BitDataset) -> dict:
"""Turn the solver's tree of capture sets into one of named rules.
Each leaf keeps the class counts of the training rows it captures, so
that `predict_proba` has something to report.
"""
if "prediction" in node:
count, dist, max_loss, _, _, prediction = data.leaf_stats(node["key"])
return {
"prediction": self.target_encoder_.inverse(prediction),
"name": "class",
"loss": float(max_loss),
"complexity": float(self.regularization),
"count": int(count),
"dist": np.asarray(dist, dtype=float).tolist(),
}
rule = self.encoder_.rules[node["feature"]]
return {
"feature": int(rule["feature"]),
"name": rule["name"],
"relation": rule["relation"],
"reference": rule["reference"],
"type": rule["type"],
"true": self._decode(node["true"], data),
"false": self._decode(node["false"], data),
}
def _to_sklearn(self, tree: dict, frame: pd.DataFrame, costs: np.ndarray):
"""The certified tree as a fitted `DecisionTreeClassifier`.
sklearn sends a row left when ``x[feature] <= threshold``, so a rule's
false side becomes the left child and its true side the right one. Every
rule is a threshold on a numeric column (``check_fit_arguments`` admits
nothing else), so every rule maps onto one sklearn split.
sklearn predicts ``argmax`` of a leaf's ``value``, so ``value`` must pick
the class the solver certified. The solver picks ``argmin(costs @ dist)``
with the first class winning ties, which for the default objective and
for ``balance`` is ``argmax`` of the class-weighted counts, what sklearn
itself stores under ``class_weight``. A general ``costs`` matrix has no
such weighting, so there ``value`` holds each class's saving over the
costliest one, which ``argmax`` ranks the same way.
Returns the estimator and each node's empirical class frequencies, which
`predict_proba` reports.
"""
K = len(self.classes_)
weights = _class_weights(costs)
values_of = {} # column -> its sorted training values
def threshold(node):
j = node["feature"]
if node["type"] == "categorical":
raise ValueError("categorical rules cannot be expressed as sklearn splits")
if j not in values_of:
values_of[j] = np.unique(frame.iloc[:, j].to_numpy(dtype=np.float64))
vals, ref = values_of[j], node["reference"]
# on the training values, ">= ref" and "== ref" (a two-valued column,
# ref its larger value) both hold exactly for the values >= ref
lo, hi = vals[vals < ref].max(), vals[vals >= ref].min()
mid = 0.5 * (lo + hi)
# two adjacent floats have no float between them, and the midpoint then
# rounds onto one of them; ``x <= lo`` still splits them as the rule does
return mid if lo <= mid < hi else lo
rows = []
def add(node, depth):
i = len(rows)
rows.append(None)
if "prediction" in node:
dist = np.asarray(node["dist"], dtype=np.float64)
rows[i] = (-1, -1, -2, -2.0, dist, int(node["prediction"]))
return dist, depth
left = len(rows)
left_dist, left_depth = add(node["false"], depth + 1)
right = len(rows)
right_dist, right_depth = add(node["true"], depth + 1)
dist = left_dist + right_dist
rows[i] = (left, right, int(node["feature"]), threshold(node), dist, None)
return dist, max(left_depth, right_depth)
_, max_depth = add(tree, 0)
est = Tree(self.n_features_in_, np.array([K], dtype=np.intp), 1)
nodes = np.zeros(len(rows), dtype=est.__getstate__()["nodes"].dtype)
values = np.zeros((len(rows), 1, K))
node_proba = np.zeros((len(rows), K))
for i, (left, right, feature, thr, dist, prediction) in enumerate(rows):
n = dist.sum()
freq = dist / n if n > 0 else np.full(K, 1.0 / K)
nodes[i]["left_child"], nodes[i]["right_child"] = left, right
nodes[i]["feature"], nodes[i]["threshold"] = feature, thr
nodes[i]["impurity"] = 1.0 - float(np.sum(freq ** 2))
nodes[i]["n_node_samples"] = nodes[i]["weighted_n_node_samples"] = n
if "missing_go_to_left" in nodes.dtype.names:
nodes[i]["missing_go_to_left"] = 1 # a missing value fails a rule: its false side
values[i, 0] = _leaf_value(dist, costs, weights, prediction)
node_proba[i] = freq
est.__setstate__({"max_depth": max_depth, "node_count": len(rows),
"nodes": nodes, "values": values})
clf = DecisionTreeClassifier()
clf.tree_ = est
clf.classes_ = self.classes_
clf.n_classes_ = K
clf.n_outputs_ = 1
clf.n_features_in_ = self.n_features_in_
clf.max_features_ = self.n_features_in_
if hasattr(self, "feature_names_in_"):
clf.feature_names_in_ = self.feature_names_in_
return clf, node_proba
def _node_predictions(self, tree: dict) -> np.ndarray:
"""Each node's certified class code, in `estimator_`'s node order (-1 inside)."""
out = []
def add(node):
if "prediction" in node:
out.append(int(node["prediction"]))
return
out.append(-1)
add(node["false"])
add(node["true"])
add(tree)
return np.asarray(out, dtype=np.intp)
def _check_sklearn_tree(self, certified, frame: pd.DataFrame):
"""Check `predict` against the certified tree on the training rows.
`predict` walks `estimator_`'s splits in float64, which route every
training row as the certified rules do, so a difference is a bug. sklearn's
own ``estimator_.predict`` compares X as float32, so two training values of
one column closer than float32 resolves can land on the same side of a
split; that only affects ``estimator_`` used directly, and is warned about.
"""
X = frame.to_numpy(dtype=np.float64)
ours = certified.predict_fast(frame).astype(int)
if not np.array_equal(self._node_pred_[self._apply(X)], ours):
raise AssertionError("predict routes training rows differently from the certified tree")
theirs = self.estimator_.predict(self._sklearn_X(X))
n_diff = int(np.sum(np.asarray(self.classes_[ours]) != np.asarray(theirs)))
if n_diff:
warnings.warn(
f"estimator_.predict (sklearn, float32) routes {n_diff} training rows "
"differently from the certified tree: some column has training values "
"closer together than float32 can separate. predict and predict_proba "
"are not affected; use estimator_ only for plotting and export here.",
RuntimeWarning)
def _apply(self, X: np.ndarray) -> np.ndarray:
"""Leaf index of each row of float64 X in `estimator_`, compared in float64.
A row goes left (a rule's false side) when ``x <= threshold``; a missing
value goes left too, as in `estimator_`.
"""
tree = self.estimator_.tree_
left, right = tree.children_left, tree.children_right
feature, threshold = tree.feature, tree.threshold
node = np.zeros(X.shape[0], dtype=np.intp)
rows = np.arange(X.shape[0])
for _ in range(tree.max_depth):
inner = left[node] >= 0
if not inner.any():
break
x = X[rows, np.where(inner, feature[node], 0)]
go_right = x > threshold[node] # NaN compares False: left
node = np.where(inner, np.where(go_right, right[node], left[node]), node)
return node
def _predict_X(self, X) -> np.ndarray:
"""X at predict time as float64, sparse input densified as fit does."""
X = check_predict_X(self, X) # checks fitted-ness before any attribute of ours
if scipy.sparse.issparse(X):
X = X.toarray()
return check_array(X, dtype=np.float64, **_finite_check_kwarg(True))
def _sklearn_X(self, X):
"""X as `estimator_` expects it: named when the model was fitted on names."""
X = np.asarray(X, dtype=float)
if hasattr(self, "feature_names_in_"):
return pd.DataFrame(X, columns=self.feature_names_in_)
return X
def predict(self, X):
"""Predict the class of each row of X with the certified tree, in float64."""
X = self._predict_X(X)
return self.classes_[self._node_pred_[self._apply(X)]]
def predict_proba(self, X):
"""Class probabilities, read off the training rows in each leaf.
A leaf of an optimal tree predicts one class, so these are the empirical
class frequencies of the training rows that reached the leaf rather than
a calibrated probability.
"""
X = self._predict_X(X)
return self._node_proba_[self._apply(X)]
def _display_tree(self, node: dict) -> dict:
"""A copy of the tree whose leaves carry the caller's class labels.
The tree itself predicts the integer codes `check_fit_arguments` made of
y, which is what routes rows; only printing wants the labels back.
"""
if "prediction" in node:
label = self.classes_[int(node["prediction"])]
# unwrap numpy scalars so the printed label reads as the caller wrote it
return {**node, "prediction": getattr(label, "item", lambda: label)()}
return {**node,
"true": self._display_tree(node["true"]),
"false": self._display_tree(node["false"])}Classes
class FastSmallTreeClassifier (regularization: float = 0.05, time_limit: float = 60.0, balance: bool = False, costs=None, memory_limit: int = 0, verbose: bool = False)-
A decision tree that is certifiably optimal for its objective.
Fits the tree minimising
misclassification rate + <code>regularization</code> * (number of leaves)over every decision tree on the binarized features, and proves that no other such tree scores better. This is the objective of
GOSDT <https://arxiv.org/abs/2006.08690>(Lin et al., ICML 2020); the branch-and-bound search is described insolver.pyand was developed by autoresearch (seeAgentic-imodels <https://github.com/csinva/agentic-imodels>).Greedy trees (CART, C4.5) pick each split without regard to the splits below it, so they can miss the best tree of a given size. This model does not: when
optimal_is True, the objective it reports is a proven minimum.Depth is not a parameter.
regularizationsets the price of a leaf, and the optimal tree at that price is as deep as it needs to be. Since a leaf that does not improve the misclassification rate byregularizationis never worth adding, the trees stay small.Parameters
regularization:float, default=0.05- Penalty added to the objective for each leaf, in units of
misclassification rate. Larger values give smaller trees. Values below
1 / n_samplesare not useful, since a leaf fixing a single training point already pays for itself. time_limit:floatorNone, default=60.0- Seconds after which the search returns the best tree found so far rather
than continuing.
optimal_is then False and a warning is raised.0or None means no limit, which can run for a very long time on wide or continuous data. balance:bool, default=False- Weigh the classes equally, optimising balanced accuracy rather than accuracy.
costs:array-likeofshape (n_classes, n_classes), default=Nonecosts[i, j]is the cost of predicting classiwhen the truth is classj. Overridesbalance.memory_limit:int, default=0- Bytes held by the search's node store (the memo of solved subproblems)
above which the search stops, as the time limit does. This is not the
resident memory of the process, which also holds the data, the
workspace and Python itself.
0means no limit. The store grows with the number of subproblems kept, which is what makes hard instances expensive. verbose:bool, default=False- Print search statistics after fitting.
Attributes
optimal_:bool- Whether optimality was certified. False means a limit was reached first, so the tree is the best one found rather than a proven optimum.
objective_:float- Objective value of the returned tree: its misclassification rate plus
regularizationtimes its number of leaves. Underbalancethe rate is the class-weighted (balanced) error, and undercostsit is the average cost per training row. lowerbound_,upperbound_:float- Interval the search closed on the optimal objective; equal when
optimal_is True. n_leaves_:int- Number of leaves in the returned tree.
complexity_:int- Number of splits in the returned tree, the imodels measure of size.
n_binary_features_:int- Number of binary features the encoder produced from X. A continuous column becomes one feature per distinct value, which is what makes continuous data expensive to solve exactly.
estimator_:sklearn.tree.DecisionTreeClassifier- The fitted tree as an sklearn decision tree, so sklearn's tree tooling
(
plot_tree,export_text,feature_importances_, dtreeviz) works on it directly. Each split is placed midway between the training values on either side of the rule, as sklearn places its own. A new value strictly between two consecutive training values of a column is routed by that midpoint, sklearn's convention; the objective is defined on the training rows, so this changes nothing the certificate covers.predictandpredict_probawalk these splits in float64 rather than callingestimator_.predict: sklearn compares X as float32, so on a column whose training values are closer together than float32 resolves (unix timestamps, say)estimator_.predictcan route training rows differently from the certified tree, whilepredictcannot. tree_:dict- The same tree as nested dicts of named rules, which printing the model shows.
time_:float- Seconds spent in the search.
classes_:ndarray- Class labels seen during fit.
Notes
The compiled search takes about 20 seconds to build the first time it runs on a machine. The result is cached on disk, so later processes load it in about a second. Set
OPTTREE_NUMBA_CACHE=0to skip the cache if its directory is unwritable.Examples
>>> import numpy as np >>> from imodels import FastSmallTreeClassifier >>> X = np.array([[0, 0], [0, 1], [1, 0], [1, 1]] * 8) >>> y = (X[:, 0] == X[:, 1]).astype(int) >>> model = FastSmallTreeClassifier(regularization=0.01).fit(X, y) >>> model.optimal_ TrueExpand source code
class FastSmallTreeClassifier(RuleInspectionMixin, ClassifierMixin, BaseEstimator): """A decision tree that is certifiably optimal for its objective. Fits the tree minimising misclassification rate + ``regularization`` * (number of leaves) over *every* decision tree on the binarized features, and proves that no other such tree scores better. This is the objective of `GOSDT <https://arxiv.org/abs/2006.08690>`_ (Lin et al., ICML 2020); the branch-and-bound search is described in `solver.py` and was developed by autoresearch (see `Agentic-imodels <https://github.com/csinva/agentic-imodels>`_). Greedy trees (CART, C4.5) pick each split without regard to the splits below it, so they can miss the best tree of a given size. This model does not: when ``optimal_`` is True, the objective it reports is a proven minimum. Depth is not a parameter. ``regularization`` sets the price of a leaf, and the optimal tree at that price is as deep as it needs to be. Since a leaf that does not improve the misclassification rate by ``regularization`` is never worth adding, the trees stay small. Parameters ---------- regularization : float, default=0.05 Penalty added to the objective for each leaf, in units of misclassification rate. Larger values give smaller trees. Values below ``1 / n_samples`` are not useful, since a leaf fixing a single training point already pays for itself. time_limit : float or None, default=60.0 Seconds after which the search returns the best tree found so far rather than continuing. ``optimal_`` is then False and a warning is raised. ``0`` or None means no limit, which can run for a very long time on wide or continuous data. balance : bool, default=False Weigh the classes equally, optimising balanced accuracy rather than accuracy. costs : array-like of shape (n_classes, n_classes), default=None ``costs[i, j]`` is the cost of predicting class ``i`` when the truth is class ``j``. Overrides ``balance``. memory_limit : int, default=0 Bytes held by the search's node store (the memo of solved subproblems) above which the search stops, as the time limit does. This is not the resident memory of the process, which also holds the data, the workspace and Python itself. ``0`` means no limit. The store grows with the number of subproblems kept, which is what makes hard instances expensive. verbose : bool, default=False Print search statistics after fitting. Attributes ---------- optimal_ : bool Whether optimality was certified. False means a limit was reached first, so the tree is the best one found rather than a proven optimum. objective_ : float Objective value of the returned tree: its misclassification rate plus ``regularization`` times its number of leaves. Under ``balance`` the rate is the class-weighted (balanced) error, and under ``costs`` it is the average cost per training row. lowerbound_, upperbound_ : float Interval the search closed on the optimal objective; equal when ``optimal_`` is True. n_leaves_ : int Number of leaves in the returned tree. complexity_ : int Number of splits in the returned tree, the imodels measure of size. n_binary_features_ : int Number of binary features the encoder produced from X. A continuous column becomes one feature per distinct value, which is what makes continuous data expensive to solve exactly. estimator_ : sklearn.tree.DecisionTreeClassifier The fitted tree as an sklearn decision tree, so sklearn's tree tooling (``plot_tree``, ``export_text``, ``feature_importances_``, dtreeviz) works on it directly. Each split is placed midway between the training values on either side of the rule, as sklearn places its own. A new value strictly between two consecutive training values of a column is routed by that midpoint, sklearn's convention; the objective is defined on the training rows, so this changes nothing the certificate covers. `predict` and `predict_proba` walk these splits in float64 rather than calling ``estimator_.predict``: sklearn compares X as float32, so on a column whose training values are closer together than float32 resolves (unix timestamps, say) ``estimator_.predict`` can route training rows differently from the certified tree, while `predict` cannot. tree_ : dict The same tree as nested dicts of named rules, which printing the model shows. time_ : float Seconds spent in the search. classes_ : ndarray Class labels seen during fit. Notes ----- The compiled search takes about 20 seconds to build the first time it runs on a machine. The result is cached on disk, so later processes load it in about a second. Set ``OPTTREE_NUMBA_CACHE=0`` to skip the cache if its directory is unwritable. Examples -------- >>> import numpy as np >>> from imodels import FastSmallTreeClassifier >>> X = np.array([[0, 0], [0, 1], [1, 0], [1, 1]] * 8) >>> y = (X[:, 0] == X[:, 1]).astype(int) >>> model = FastSmallTreeClassifier(regularization=0.01).fit(X, y) >>> model.optimal_ True """ def __init__(self, regularization: float = 0.05, time_limit: float = 60.0, balance: bool = False, costs=None, memory_limit: int = 0, verbose: bool = False): self.regularization = regularization self.time_limit = time_limit self.balance = balance self.costs = costs self.memory_limit = memory_limit self.verbose = verbose def fit(self, X, y, feature_names=None): """Fit the optimal tree. Parameters ---------- X : array-like of shape (n_samples, n_features) y : array-like of shape (n_samples,) feature_names : list of str, optional Names used when printing the model; taken from X's columns when it is a DataFrame. Returns ------- self """ if not HAVE_NUMBA: raise ImportError(NUMBA_HINT) notify_first_compile(_solver_module.__file__, "FastSmallTreeClassifier", "about 20 seconds", _solver_module.NUMBA_CACHE, "OPTTREE_NUMBA_CACHE") time_limit, memory_limit = self._check_limits() X, y, feature_names = check_fit_arguments(self, X, y, feature_names) X = self._as_float(X) # the certificate rests on bounds of the form "a tree with a leaves costs at least # a * regularization" and "a tree's loss cannot fall below zero", which need a # nonnegative penalty and nonnegative, finite costs lam = float(self.regularization) if not np.isfinite(lam) or lam < 0.0: raise ValueError(f"regularization must be a finite number >= 0, got {self.regularization!r}") if self.costs is not None: C = np.asarray(self.costs, dtype=np.float64) if not np.all(np.isfinite(C)) or np.any(C < 0.0): raise ValueError("costs must be finite and nonnegative") frame = pd.DataFrame(X, columns=list(feature_names)) self.encoder_ = BinaryEncoder().fit(frame) Xb = self.encoder_.transform(frame) self.target_encoder_ = TargetEncoder().fit(y) y_idx = self.target_encoder_.transform(y) # the objective does not depend on the row order, but the kernels skip the # empty words of a capture mask, so rows are clustered to keep the capture # sets of deep nodes contiguous perm = cluster_rows(Xb, y_idx, len(self.classes_), self.encoder_.groups) data = BitDataset(Xb[perm], y_idx[perm], len(self.classes_), costs=self.costs, balance=self.balance) self.n_binary_features_ = data.m # n_jobs=1: the search has a parallel phase, but a library model should not # take every core of the caller's machine without being asked opt = CompiledOptimizer(data, self.regularization, groups=self.encoder_.groups, time_limit=time_limit, memory_limit=memory_limit, verbose=self.verbose, n_jobs=1) root = opt.run() # a row index into the compiled engine's node store self.optimal_ = opt.optimal self.stop_reason_ = opt.stop_reason self.time_ = opt.elapsed self.iterations_ = opt.iterations self.lowerbound_ = float(opt.st[ST_LB][root]) self.upperbound_ = float(opt.st[ST_UB][root]) if not self.optimal_: warnings.warn( f"{self.stop_reason_} limit reached before optimality was certified; " "returning the best tree found. Raise time_limit, or raise " "regularization to search over smaller trees.", RuntimeWarning) self.tree_ = self._decode(opt.extract(root), data) opt.release() # the node store is large and is not needed past extraction certified = TreeClassifier(self.tree_) self.objective_ = certified.risk() self.n_leaves_ = certified.leaves() self.complexity_ = max(self.n_leaves_ - 1, 0) self.estimator_, self._node_proba_ = self._to_sklearn(self.tree_, frame, data.costs) self._node_pred_ = self._node_predictions(self.tree_) self._check_sklearn_tree(certified, frame) if self.verbose: print(f"objective {self.objective_:.6g} in {self.time_:.3g}s " f"({self.iterations_} subproblems, optimal={self.optimal_})") return self def _check_limits(self): """``time_limit`` and ``memory_limit`` as the search takes them (0 = none).""" time_limit = self.time_limit if time_limit is None: time_limit = 0.0 try: time_limit = float(time_limit) except (TypeError, ValueError): raise ValueError(f"time_limit must be a number >= 0 or None, got {self.time_limit!r}") if not np.isfinite(time_limit) or time_limit < 0.0: raise ValueError(f"time_limit must be a finite number >= 0 (0 or None for no " f"limit), got {self.time_limit!r}") memory_limit = self.memory_limit if memory_limit is None: memory_limit = 0 try: memory_limit = float(memory_limit) except (TypeError, ValueError): raise ValueError(f"memory_limit must be a number of bytes >= 0, got {self.memory_limit!r}") if not np.isfinite(memory_limit) or memory_limit < 0: raise ValueError(f"memory_limit must be a finite number of bytes >= 0 (0 for no " f"limit), got {self.memory_limit!r}") return time_limit, int(memory_limit) @staticmethod def _as_float(X): """X as a float64 array, or a clear error before the search starts. ``check_fit_arguments`` leaves the object array that a DataFrame with a pandas ``Categorical`` column becomes; numeric categories are kept as their values, and anything else cannot be thresholded. """ if scipy.sparse.issparse(X): X = X.toarray() try: return np.asarray(X, dtype=np.float64) except (TypeError, ValueError): raise ValueError( "FastSmallTreeClassifier needs numeric features, but X has a column " "that cannot be converted to float (for example a pandas Categorical " "of strings). Encode such columns as numbers first, e.g. with " "pd.get_dummies or OrdinalEncoder.") from None def _decode(self, node: dict, data: BitDataset) -> dict: """Turn the solver's tree of capture sets into one of named rules. Each leaf keeps the class counts of the training rows it captures, so that `predict_proba` has something to report. """ if "prediction" in node: count, dist, max_loss, _, _, prediction = data.leaf_stats(node["key"]) return { "prediction": self.target_encoder_.inverse(prediction), "name": "class", "loss": float(max_loss), "complexity": float(self.regularization), "count": int(count), "dist": np.asarray(dist, dtype=float).tolist(), } rule = self.encoder_.rules[node["feature"]] return { "feature": int(rule["feature"]), "name": rule["name"], "relation": rule["relation"], "reference": rule["reference"], "type": rule["type"], "true": self._decode(node["true"], data), "false": self._decode(node["false"], data), } def _to_sklearn(self, tree: dict, frame: pd.DataFrame, costs: np.ndarray): """The certified tree as a fitted `DecisionTreeClassifier`. sklearn sends a row left when ``x[feature] <= threshold``, so a rule's false side becomes the left child and its true side the right one. Every rule is a threshold on a numeric column (``check_fit_arguments`` admits nothing else), so every rule maps onto one sklearn split. sklearn predicts ``argmax`` of a leaf's ``value``, so ``value`` must pick the class the solver certified. The solver picks ``argmin(costs @ dist)`` with the first class winning ties, which for the default objective and for ``balance`` is ``argmax`` of the class-weighted counts, what sklearn itself stores under ``class_weight``. A general ``costs`` matrix has no such weighting, so there ``value`` holds each class's saving over the costliest one, which ``argmax`` ranks the same way. Returns the estimator and each node's empirical class frequencies, which `predict_proba` reports. """ K = len(self.classes_) weights = _class_weights(costs) values_of = {} # column -> its sorted training values def threshold(node): j = node["feature"] if node["type"] == "categorical": raise ValueError("categorical rules cannot be expressed as sklearn splits") if j not in values_of: values_of[j] = np.unique(frame.iloc[:, j].to_numpy(dtype=np.float64)) vals, ref = values_of[j], node["reference"] # on the training values, ">= ref" and "== ref" (a two-valued column, # ref its larger value) both hold exactly for the values >= ref lo, hi = vals[vals < ref].max(), vals[vals >= ref].min() mid = 0.5 * (lo + hi) # two adjacent floats have no float between them, and the midpoint then # rounds onto one of them; ``x <= lo`` still splits them as the rule does return mid if lo <= mid < hi else lo rows = [] def add(node, depth): i = len(rows) rows.append(None) if "prediction" in node: dist = np.asarray(node["dist"], dtype=np.float64) rows[i] = (-1, -1, -2, -2.0, dist, int(node["prediction"])) return dist, depth left = len(rows) left_dist, left_depth = add(node["false"], depth + 1) right = len(rows) right_dist, right_depth = add(node["true"], depth + 1) dist = left_dist + right_dist rows[i] = (left, right, int(node["feature"]), threshold(node), dist, None) return dist, max(left_depth, right_depth) _, max_depth = add(tree, 0) est = Tree(self.n_features_in_, np.array([K], dtype=np.intp), 1) nodes = np.zeros(len(rows), dtype=est.__getstate__()["nodes"].dtype) values = np.zeros((len(rows), 1, K)) node_proba = np.zeros((len(rows), K)) for i, (left, right, feature, thr, dist, prediction) in enumerate(rows): n = dist.sum() freq = dist / n if n > 0 else np.full(K, 1.0 / K) nodes[i]["left_child"], nodes[i]["right_child"] = left, right nodes[i]["feature"], nodes[i]["threshold"] = feature, thr nodes[i]["impurity"] = 1.0 - float(np.sum(freq ** 2)) nodes[i]["n_node_samples"] = nodes[i]["weighted_n_node_samples"] = n if "missing_go_to_left" in nodes.dtype.names: nodes[i]["missing_go_to_left"] = 1 # a missing value fails a rule: its false side values[i, 0] = _leaf_value(dist, costs, weights, prediction) node_proba[i] = freq est.__setstate__({"max_depth": max_depth, "node_count": len(rows), "nodes": nodes, "values": values}) clf = DecisionTreeClassifier() clf.tree_ = est clf.classes_ = self.classes_ clf.n_classes_ = K clf.n_outputs_ = 1 clf.n_features_in_ = self.n_features_in_ clf.max_features_ = self.n_features_in_ if hasattr(self, "feature_names_in_"): clf.feature_names_in_ = self.feature_names_in_ return clf, node_proba def _node_predictions(self, tree: dict) -> np.ndarray: """Each node's certified class code, in `estimator_`'s node order (-1 inside).""" out = [] def add(node): if "prediction" in node: out.append(int(node["prediction"])) return out.append(-1) add(node["false"]) add(node["true"]) add(tree) return np.asarray(out, dtype=np.intp) def _check_sklearn_tree(self, certified, frame: pd.DataFrame): """Check `predict` against the certified tree on the training rows. `predict` walks `estimator_`'s splits in float64, which route every training row as the certified rules do, so a difference is a bug. sklearn's own ``estimator_.predict`` compares X as float32, so two training values of one column closer than float32 resolves can land on the same side of a split; that only affects ``estimator_`` used directly, and is warned about. """ X = frame.to_numpy(dtype=np.float64) ours = certified.predict_fast(frame).astype(int) if not np.array_equal(self._node_pred_[self._apply(X)], ours): raise AssertionError("predict routes training rows differently from the certified tree") theirs = self.estimator_.predict(self._sklearn_X(X)) n_diff = int(np.sum(np.asarray(self.classes_[ours]) != np.asarray(theirs))) if n_diff: warnings.warn( f"estimator_.predict (sklearn, float32) routes {n_diff} training rows " "differently from the certified tree: some column has training values " "closer together than float32 can separate. predict and predict_proba " "are not affected; use estimator_ only for plotting and export here.", RuntimeWarning) def _apply(self, X: np.ndarray) -> np.ndarray: """Leaf index of each row of float64 X in `estimator_`, compared in float64. A row goes left (a rule's false side) when ``x <= threshold``; a missing value goes left too, as in `estimator_`. """ tree = self.estimator_.tree_ left, right = tree.children_left, tree.children_right feature, threshold = tree.feature, tree.threshold node = np.zeros(X.shape[0], dtype=np.intp) rows = np.arange(X.shape[0]) for _ in range(tree.max_depth): inner = left[node] >= 0 if not inner.any(): break x = X[rows, np.where(inner, feature[node], 0)] go_right = x > threshold[node] # NaN compares False: left node = np.where(inner, np.where(go_right, right[node], left[node]), node) return node def _predict_X(self, X) -> np.ndarray: """X at predict time as float64, sparse input densified as fit does.""" X = check_predict_X(self, X) # checks fitted-ness before any attribute of ours if scipy.sparse.issparse(X): X = X.toarray() return check_array(X, dtype=np.float64, **_finite_check_kwarg(True)) def _sklearn_X(self, X): """X as `estimator_` expects it: named when the model was fitted on names.""" X = np.asarray(X, dtype=float) if hasattr(self, "feature_names_in_"): return pd.DataFrame(X, columns=self.feature_names_in_) return X def predict(self, X): """Predict the class of each row of X with the certified tree, in float64.""" X = self._predict_X(X) return self.classes_[self._node_pred_[self._apply(X)]] def predict_proba(self, X): """Class probabilities, read off the training rows in each leaf. A leaf of an optimal tree predicts one class, so these are the empirical class frequencies of the training rows that reached the leaf rather than a calibrated probability. """ X = self._predict_X(X) return self._node_proba_[self._apply(X)] def _display_tree(self, node: dict) -> dict: """A copy of the tree whose leaves carry the caller's class labels. The tree itself predicts the integer codes `check_fit_arguments` made of y, which is what routes rows; only printing wants the labels back. """ if "prediction" in node: label = self.classes_[int(node["prediction"])] # unwrap numpy scalars so the printed label reads as the caller wrote it return {**node, "prediction": getattr(label, "item", lambda: label)()} return {**node, "true": self._display_tree(node["true"]), "false": self._display_tree(node["false"])}Ancestors
- RuleInspectionMixin
- RulesMixin
- TextMixin
- LeavesMixin
- sklearn.base.ClassifierMixin
- sklearn.base.BaseEstimator
- sklearn.utils._estimator_html_repr._HTMLDocumentationLinkMixin
- sklearn.utils._metadata_requests._MetadataRequester
Methods
def fit(self, X, y, feature_names=None)-
Fit the optimal tree.
Parameters
X:array-likeofshape (n_samples, n_features)y:array-likeofshape (n_samples,)feature_names:listofstr, optional- Names used when printing the model; taken from X's columns when it is a DataFrame.
Returns
self
Expand source code
def fit(self, X, y, feature_names=None): """Fit the optimal tree. Parameters ---------- X : array-like of shape (n_samples, n_features) y : array-like of shape (n_samples,) feature_names : list of str, optional Names used when printing the model; taken from X's columns when it is a DataFrame. Returns ------- self """ if not HAVE_NUMBA: raise ImportError(NUMBA_HINT) notify_first_compile(_solver_module.__file__, "FastSmallTreeClassifier", "about 20 seconds", _solver_module.NUMBA_CACHE, "OPTTREE_NUMBA_CACHE") time_limit, memory_limit = self._check_limits() X, y, feature_names = check_fit_arguments(self, X, y, feature_names) X = self._as_float(X) # the certificate rests on bounds of the form "a tree with a leaves costs at least # a * regularization" and "a tree's loss cannot fall below zero", which need a # nonnegative penalty and nonnegative, finite costs lam = float(self.regularization) if not np.isfinite(lam) or lam < 0.0: raise ValueError(f"regularization must be a finite number >= 0, got {self.regularization!r}") if self.costs is not None: C = np.asarray(self.costs, dtype=np.float64) if not np.all(np.isfinite(C)) or np.any(C < 0.0): raise ValueError("costs must be finite and nonnegative") frame = pd.DataFrame(X, columns=list(feature_names)) self.encoder_ = BinaryEncoder().fit(frame) Xb = self.encoder_.transform(frame) self.target_encoder_ = TargetEncoder().fit(y) y_idx = self.target_encoder_.transform(y) # the objective does not depend on the row order, but the kernels skip the # empty words of a capture mask, so rows are clustered to keep the capture # sets of deep nodes contiguous perm = cluster_rows(Xb, y_idx, len(self.classes_), self.encoder_.groups) data = BitDataset(Xb[perm], y_idx[perm], len(self.classes_), costs=self.costs, balance=self.balance) self.n_binary_features_ = data.m # n_jobs=1: the search has a parallel phase, but a library model should not # take every core of the caller's machine without being asked opt = CompiledOptimizer(data, self.regularization, groups=self.encoder_.groups, time_limit=time_limit, memory_limit=memory_limit, verbose=self.verbose, n_jobs=1) root = opt.run() # a row index into the compiled engine's node store self.optimal_ = opt.optimal self.stop_reason_ = opt.stop_reason self.time_ = opt.elapsed self.iterations_ = opt.iterations self.lowerbound_ = float(opt.st[ST_LB][root]) self.upperbound_ = float(opt.st[ST_UB][root]) if not self.optimal_: warnings.warn( f"{self.stop_reason_} limit reached before optimality was certified; " "returning the best tree found. Raise time_limit, or raise " "regularization to search over smaller trees.", RuntimeWarning) self.tree_ = self._decode(opt.extract(root), data) opt.release() # the node store is large and is not needed past extraction certified = TreeClassifier(self.tree_) self.objective_ = certified.risk() self.n_leaves_ = certified.leaves() self.complexity_ = max(self.n_leaves_ - 1, 0) self.estimator_, self._node_proba_ = self._to_sklearn(self.tree_, frame, data.costs) self._node_pred_ = self._node_predictions(self.tree_) self._check_sklearn_tree(certified, frame) if self.verbose: print(f"objective {self.objective_:.6g} in {self.time_:.3g}s " f"({self.iterations_} subproblems, optimal={self.optimal_})") return self def predict(self, X)-
Predict the class of each row of X with the certified tree, in float64.
Expand source code
def predict(self, X): """Predict the class of each row of X with the certified tree, in float64.""" X = self._predict_X(X) return self.classes_[self._node_pred_[self._apply(X)]] def predict_proba(self, X)-
Class probabilities, read off the training rows in each leaf.
A leaf of an optimal tree predicts one class, so these are the empirical class frequencies of the training rows that reached the leaf rather than a calibrated probability.
Expand source code
def predict_proba(self, X): """Class probabilities, read off the training rows in each leaf. A leaf of an optimal tree predicts one class, so these are the empirical class frequencies of the training rows that reached the leaf rather than a calibrated probability. """ X = self._predict_X(X) return self._node_proba_[self._apply(X)]
Inherited members