API Reference / algorithms / causal /

uplift_tree.py

A decision tree that splits directly on uplift gain.

Ordinary regression trees split to reduce outcome variance; an uplift tree instead splits to separate high-treatment-effect regions from low-treatment-effect regions, so the leaves are groups of individuals for whom the treatment works differently.

Classes

UpliftTreeClassifier

class algorithms.causal.uplift_tree.UpliftTreeClassifier(UpliftModel)

Single uplift tree that splits on the difference in treatment effects.

Constructor
__init__(
    self,
    max_depth: Optional[int] = None,
    min_samples_split: int = 2,
    min_samples_leaf: int = 20,
    max_features: Optional[int] = None,
    random_state: Optional[int] = None,
)

Summary

A greedy binary tree whose split criterion is the uplift gain — the squared difference in treatment effect between the two child nodes, weighted by their size. Each leaf stores the observed treatment effect \bar{y}_{t=1} - \bar{y}_{t=0} of the samples that land there.

Overview

  1. For each candidate feature and threshold, split the node and estimate
the uplift of each child as its treated-mean minus control-mean.
  1. Choose the split that maximizes
\frac{n_L n_R}{n}\,(\hat{\tau}_L - \hat{\tau}_R)^2.
  1. Recurse until a stopping rule fires, then store the leaf uplift.

Theory

Let a node contain n samples, n_t treated and n_c control, with outcomes y. Its estimated uplift is

\hat{\tau} = \frac{1}{n_t}\sum_{i: t_i=1} y_i - \frac{1}{n_c}\sum_{i: t_i=0} y_i.

A split sends the node's samples to a left child L and a right child R. The split is chosen to maximize the between-child uplift variance,

\text{gain} = \frac{n_L n_R}{n} \left(\hat{\tau}_L - \hat{\tau}_R\right)^2,

which is large exactly when the two children have very different treatment effects. This targets heterogeneity directly rather than the outcome level.

Parameters

max_depth
int or None = None
Maximum tree depth (None for no limit).
min_samples_split
int = 2
Minimum samples required to split an internal node.
min_samples_leaf
int = 20
Minimum samples required in a child for a split to be accepted. Each child must also contain at least one treated and one control sample so its uplift is defined.
max_features
int or None = None
Number of features to consider at each split (None uses all).
random_state
int or None = None
Random seed for feature sub-sampling.

Attributes

tree_
dict
The root node of the fitted tree. Internal nodes have feature, threshold, left and right keys; leaves have uplift, n_treated and n_control.
n_features_in_
int
Number of features in X.
n_nodes_
int
Total number of nodes in the fitted tree.
max_depth_
int
Depth of the fitted tree.

Notes

Complexity: each split sorts the node by each candidate feature, so training is roughly O(d \, n \log n) per level and prediction is O(\text{depth}) per sample.

When to use: when you want a single, inspectable tree (rather than a black-box meta-learner) and the treatment effect is piecewise-constant in the features.

References

Rzepakowski2012
Rzepakowski, P. and Jaroszewicz, S. (2012). Decision trees for uplift modeling with single and multiple treatments. Knowledge and Information Systems, 32(2), 303-327. DOI: 10.1007/s10115-011-0434-0
Athey2016
Athey, S. and Imbens, G. (2016). Recursive partitioning for heterogeneous causal effects. Proceedings of the National Academy of Sciences, 113(27), 7353-7360. DOI: 10.1073/pnas.1510489113
python
>>> from tuiml.algorithms.causal import UpliftTreeClassifier
>>> import numpy as np
>>> rng = np.random.RandomState(0)
>>> X = rng.uniform(-1, 1, size=(500, 2))
>>> t = rng.randint(0, 2, size=500)
>>> y = 1.0 + X[:, 1] + t * (2.0 * X[:, 0]) + rng.normal(0, 0.1, size=500)
>>> model = UpliftTreeClassifier(max_depth=4, min_samples_leaf=20).fit(X, t, y)
>>> model.predict_uplift(X).shape
(500,)

Methods

get_parameter_schema (cls) -> Dict[str, Dict[str, Any]]

Return JSON Schema for constructor parameters.

get_capabilities (cls) -> List[str]

Return supported capabilities.

get_complexity (cls) -> str

Return complexity analysis.

get_references (cls) -> List[str]

Return academic citations.

fit (self, X, treatment, y) -> 'UpliftTreeClassifier'

Build the uplift tree.

Parameters
X
np.ndarray of shape (n_samples, n_features)
Covariates.
treatment
np.ndarray of shape (n_samples,)
Binary treatment indicator.
y
np.ndarray of shape (n_samples,)
Numeric outcome.
Returns
self
UpliftTreeClassifier
Fitted estimator.
predict_uplift (self, X: np.ndarray) -> np.ndarray

Return the leaf uplift for each sample.

Parameters
X
np.ndarray of shape (n_samples, n_features)
Covariates.
Returns
uplift
np.ndarray of shape (n_samples,)
Predicted individual treatment effect.