API Reference / algorithms / causal /

meta_learners.py

Meta-learners for uplift / heterogeneous treatment effect estimation.

The S-, T- and X-learners wrap an arbitrary TuiML regressor (for example DecisionTreeRegressor) and re-arrange the (X, treatment, y) data so that an ordinary supervised learner can estimate the conditional average treatment effect

\tau(x) = E[Y(1) - Y(0) \mid X = x],

where Y(1) and Y(0) are the potential outcomes under treatment and control.

Classes

SLearner

class algorithms.causal.meta_learners.SLearner(UpliftModel)

S-learner: a single model on [X, treatment].

Constructor 
__init__(
    self,
    estimator: Optional[object] = None,
)

Summary

The S-learner stacks the treatment indicator onto the covariates and fits one model f(X, t). The uplift is the difference between the two counterfactual predictions:

\hat{\tau}(x) = f(x, 1) - f(x, 0).

Overview

  1. Append the treatment column to X.
  2. Fit a single regressor on the augmented [X, treatment].
  3. Predict the uplift as f(X, 1) - f(X, 0).

Theory

The S-learner lets a single model learn both response surfaces jointly, regularizing them toward each other. This is efficient when the two surfaces are similar, but the treatment indicator can be ignored by flexible learners (its signal diluted among the other features), which shrinks the estimated uplift toward zero.

Parameters

estimator
object) = DecisionTreeRegressor(
A TuiML regressor (instance or class) with fit/predict.

Attributes

model_
object
The fitted regressor on [X, treatment].
n_features_in_
int
Number of features in X (without the treatment column).
n_treated_, n_control_
int
Number of samples in each treatment group.

Notes

Complexity: one supervised fit plus 2 predictions per sample.

When to use: a strong default; the T- and X-learners beat it when the treatment groups are imbalanced or their response surfaces differ a lot.

References

Kunzel2019
Kunzel, S.R., Sekhon, J.S., Bickel, P.J. and Yu, B. (2019). Metalearners for estimating heterogeneous treatment effects using machine learning. Proceedings of the National Academy of Sciences, 116(10), 4156-4165. DOI: 10.1073/pnas.1804597116

See Also

TLearner

Two separate group models.
XLearner

T-learner plus imputed effects.
python 
>>> from tuiml.algorithms.causal import SLearner
>>> from tuiml.algorithms.trees import DecisionTreeRegressor
>>> import numpy as np
>>> rng = np.random.RandomState(0)
>>> X = rng.uniform(-1, 1, size=(300, 2))
>>> t = rng.randint(0, 2, size=300)
>>> y = 1.0 + X[:, 1] + t * (2.0 * X[:, 0]) + rng.normal(0, 0.1, size=300)
>>> model = SLearner(DecisionTreeRegressor(max_depth=4)).fit(X, t, y)
>>> model.predict_uplift(X).shape
(300,)

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) -> 'SLearner'

Fit a single model on the augmented [X, treatment].

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
SLearner
Fitted estimator.
predict_uplift (self, X: np.ndarray) -> np.ndarray

Return the predicted uplift f(X, 1) - f(X, 0).

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

TLearner

class algorithms.causal.meta_learners.TLearner(UpliftModel)

T-learner: two models, one per treatment group.

Constructor 
__init__(
    self,
    estimator: Optional[object] = None,
)

Summary

The T-learner fits two regressors f_0 and f_1 on the control and treated samples separately. The uplift is their difference:

\hat{\tau}(x) = f_1(x) - f_0(x).

Overview

  1. Split (X, y) by the treatment indicator.
  2. Fit f_0 on the control group and f_1 on the treated.
  3. Predict the uplift as f_1(X) - f_0(X).

Theory

Each group gets its own response surface, so a strong treatment signal in one group cannot be diluted by the other. The trade-off is data efficiency: each model sees only its own group, which can hurt when one group is small or the surfaces share a lot of structure.

Parameters

estimator
object) = DecisionTreeRegressor(
A TuiML regressor (instance or class) with fit/predict.

Attributes

model_0_
object
Fitted regressor on the control group.
model_1_
object
Fitted regressor on the treated group.
n_features_in_
int
Number of features in X.
n_treated_, n_control_
int
Number of samples in each treatment group.

Notes

Complexity: two supervised fits plus two predictions per sample.

When to use: the two group models are genuinely independent, which makes the T-learner the cleanest baseline when treatment groups are balanced and large.

References

Kunzel2019
Kunzel, S.R., Sekhon, J.S., Bickel, P.J. and Yu, B. (2019). Metalearners for estimating heterogeneous treatment effects using machine learning. Proceedings of the National Academy of Sciences, 116(10), 4156-4165. DOI: 10.1073/pnas.1804597116

See Also

SLearner

A single shared model.
XLearner

T-learner plus imputed effects.
python 
>>> from tuiml.algorithms.causal import TLearner
>>> from tuiml.algorithms.trees import DecisionTreeRegressor
>>> import numpy as np
>>> rng = np.random.RandomState(0)
>>> X = rng.uniform(-1, 1, size=(300, 2))
>>> t = rng.randint(0, 2, size=300)
>>> y = 1.0 + X[:, 1] + t * (2.0 * X[:, 0]) + rng.normal(0, 0.1, size=300)
>>> model = TLearner(DecisionTreeRegressor(max_depth=4)).fit(X, t, y)
>>> model.predict_uplift(X).shape
(300,)

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) -> 'TLearner'

Fit separate models for the treated and control groups.

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
TLearner
Fitted estimator.
predict_uplift (self, X: np.ndarray) -> np.ndarray

Return the predicted uplift f_1(X) - f_0(X).

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

XLearner

class algorithms.causal.meta_learners.XLearner(UpliftModel)

X-learner: T-learner plus cross-group imputed-effect models.

Constructor 
__init__(
    self,
    estimator: Optional[object] = None,
    propensity_model: Optional[object] = None,
)

Summary

The X-learner starts from a T-learner, then fits two further models on the imputed treatment effects — the residuals each group leaves behind under the other group's model — and combines them with a propensity weight.

Overview

  1. Fit the T-learner response models f_0 (control) and
f_1 (treated).
  1. Impute the effect for each treated unit
D_1 = y_1 - f_0(X_1) and each control unit D_0 = f_1(X_0) - y_0.
  1. Fit \tau_1 on (X_1, D_1) and \tau_0 on
(X_0, D_0).
  1. Predict \hat{\tau}(x) = p(x)\,\tau_0(x) + (1 - p(x))\,\tau_1(x), where p(x) is the propensity score.

Theory

The imputed effect D_i is a noisy, per-unit estimate of the individual treatment effect: for a treated unit it is the outcome above what the control model would have predicted; for a control unit it is the outcome below what the treated model would have predicted. Modeling these imputed effects directly recovers \tau(x) even when one group is much smaller than the other, and the propensity-weighted combination regularizes the two estimates toward the model with more local data.

Parameters

estimator
object) = DecisionTreeRegressor(
A TuiML regressor (instance or class) used for all four sub-models.
propensity_model
object or None = None
A TuiML classifier (instance or class) used to estimate :math:`P(\text{treatment} = 1 \mid X)`. If None, a constant propensity equal to the overall treatment rate is used.

Attributes

model_0_
object
Fitted control response model :math:`f_0`.
model_1_
object
Fitted treated response model :math:`f_1`.
tau_0_
object
Fitted imputed-effect model on control units.
tau_1_
object
Fitted imputed-effect model on treated units.
propensity_model_
object or None
Fitted propensity model (None when a constant propensity is used).
propensity_
float
Constant propensity (overall treatment rate) when no model is given.
n_features_in_
int
Number of features in X.
n_treated_, n_control_
int
Number of samples in each treatment group.

Notes

Complexity: four supervised fits plus four predictions per sample.

When to use: imbalanced treatment groups, or when the base learners for the two groups are not equally accurate.

References

Kunzel2019
Kunzel, S.R., Sekhon, J.S., Bickel, P.J. and Yu, B. (2019). Metalearners for estimating heterogeneous treatment effects using machine learning. Proceedings of the National Academy of Sciences, 116(10), 4156-4165. DOI: 10.1073/pnas.1804597116

See Also

SLearner

A single shared model.
TLearner

Two separate group models.
python 
>>> from tuiml.algorithms.causal import XLearner
>>> from tuiml.algorithms.trees import DecisionTreeRegressor
>>> import numpy as np
>>> rng = np.random.RandomState(0)
>>> X = rng.uniform(-1, 1, size=(300, 2))
>>> t = rng.randint(0, 2, size=300)
>>> y = 1.0 + X[:, 1] + t * (2.0 * X[:, 0]) + rng.normal(0, 0.1, size=300)
>>> model = XLearner(DecisionTreeRegressor(max_depth=4)).fit(X, t, y)
>>> model.predict_uplift(X).shape
(300,)

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) -> 'XLearner'

Fit the T-learner, the imputed-effect models, and the propensity.

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
XLearner
Fitted estimator.
predict_uplift (self, X: np.ndarray) -> np.ndarray

Return the propensity-weighted imputed-effect prediction.

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