statspai.metalearners

metalearners

Meta-Learners for heterogeneous treatment effect estimation.

Provides S/T/X/R/DR-Learner implementations that decompose CATE estimation into standard supervised-learning sub-problems. All learners accept any scikit-learn compatible estimator.

References

Kunzel et al. (2019). Metalearners for estimating heterogeneous treatment effects using machine learning. PNAS, 116(10), 4156-4165. [@kunzel2019metalearners]

Nie & Wager (2021). Quasi-oracle estimation of heterogeneous treatment effects. Biometrika, 108(2), 299-319. [@nie2021quasi]

Kennedy (2023). Towards optimal doubly robust estimation of heterogeneous causal effects. Electronic Journal of Statistics, 17(2), 3008-3049. [@kennedy2023towards]

AutoCATEResult dataclass

Bases: ResultProtocolMixin

Leaderboard + winner from sp.auto_cate.

Attributes:

Name	Type	Description
leaderboard	DataFrame	One row per learner with ATE, SE, CI, R-loss, BLP columns and CATE dispersion, sorted by R-loss ascending. blp_beta1 is the BLP coefficient on D - e(X) (the ATE) and blp_beta1_pvalue tests it against zero; blp_beta2 / blp_beta2_pvalue are the heterogeneity slope and its test of zero, all on out-of-fold CATE predictions.
best_learner	str	Full name of the chosen winner (e.g. "DR-Learner").
best_result	CausalResult	Full fitted result for the winner — supports .summary(), .tidy(), .glance().
results	dict[str, CausalResult]	All fitted learners keyed by short code ('s', 't', ...).
agreement	DataFrame	Pearson-rho matrix of CATE vectors across learners (in-sample). High agreement suggests stable heterogeneity; low agreement suggests model dependence.
selection_rule	str	Human-readable description of the rule that picked the winner.
n_obs	int	Sample size (after dropping NA on modelled columns).

Examples:

>>> import statspai as sp
>>> import numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> age = rng.normal(size=n)
>>> edu = rng.normal(size=n)
>>> training = (rng.uniform(size=n) < 0.5).astype(int)
>>> wage = (1.0 + 0.5 * age + edu + training * (1.0 + age)
...         + rng.normal(size=n))
>>> df = pd.DataFrame({"wage": wage, "training": training,
...                    "age": age, "edu": edu})
>>> result = sp.auto_cate(
...     df, y="wage", treat="training", covariates=["age", "edu"],
...     learners=("t", "dr"), n_folds=2, n_bootstrap=50, random_state=0,
... )
>>> type(result).__name__
'AutoCATEResult'
>>> result.n_obs
200
>>> sorted(result.results)
['dr', 't']
>>> bool(result.best_learner in set(result.leaderboard["learner"]))
True

CATEGainCurveResult dataclass

Bases: ResultProtocolMixin

Output of :func:cate_gain_curve.

Attributes:

Name	Type	Description
auc	float	Sum of the (normalized) cumulative gain curve over its steps: zero in expectation for a ranking unrelated to the effect, positive when the units ranked first respond more.
ate	float	Slope of the outcome on the treatment in the whole sample.
curve	DataFrame	One row per step: share of the sample (the top-ranked units), n, cumulative_effect (slope among them) and cumulative_gain (share times that slope, minus share times ate when normalized).
by_quantile	DataFrame	Slope within each quantile group of the ranking, lowest first.
auc_se, auc_ci	float and (float, float), or None	Bootstrap standard error and percentile interval of auc; None unless n_boot was set.

Examples:

>>> import numpy as np, pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> df = pd.DataFrame({"x": rng.uniform(0, 1, 2000),
...                    "dose": rng.uniform(0, 10, 2000)})
>>> df["y"] = (1 + 2 * df["x"]) * df["dose"] + rng.normal(size=2000)
>>> res = sp.cate_gain_curve(df, cate="x", y="y", treat="dose")
>>> type(res).__name__
'CATEGainCurveResult'
>>> list(res.curve.columns)
['share', 'n', 'cumulative_effect', 'cumulative_gain']

CATEEvalResult dataclass

Bases: ResultProtocolMixin

Output of :func:cate_eval.

Attributes:

Name	Type	Description
autoc	float
autoc_se	float
autoc_ci	(float, float)
qini	float
qini_se	float
qini_ci	(float, float)
toc_curve	DataFrame	One row per quantile grid point: q, toc, its standard error se, the pointwise interval ci_lower / ci_upper and the uniform band band_lower / band_upper, which covers the whole curve with probability 1 - alpha. If the band lies above zero anywhere, the score has found heterogeneity.
n_obs	int
target	str	"AUTOC" (default) or "QINI".
method	str	Always "Yadlowsky et al. 2025 (DR-RATE, IF-SE)".

Examples:

>>> import numpy as np
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 400
>>> X = rng.normal(size=(n, 3))
>>> T = rng.integers(0, 2, n)
>>> tau = 1.0 + X[:, 0]
>>> Y = X[:, 1] + tau * T + rng.normal(scale=0.5, size=n)
>>> cate = 1.0 + X[:, 0]  # priority score to evaluate
>>> res = sp.cate_eval(cate, Y, T, X=X, n_folds=3, q_grid=20, random_state=0)
>>> isinstance(res, sp.CATEEvalResult)
True
>>> res.toc_curve.columns.tolist()[:3]
['q', 'toc', 'se']
>>> bool(res.n_obs == n)
True

plot

plot(ax: Any = None, target: Optional[str] = None, figsize: Tuple[float, float] = (6.0, 4.0)) -> tuple[Any, Any]

Plot the TOC curve plus a dashed zero line.

target defaults to self.target; pass "both" to overlay AUTOC's curve with a marker for the QINI weighted area.

ClusterCATEResult dataclass

Bases: ResultProtocolMixin

Per-cluster CATE table.

Returned by :func:cluster_cate; holds the within-cluster CATE table (one row per retained K-means cluster) plus cluster / observation counts.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 400
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> d = rng.integers(0, 2, size=n)
>>> tau = 1.0 + 0.8 * x1  # heterogeneous effect
>>> y = 2.0 + x2 + tau * d + rng.normal(scale=0.5, size=n)
>>> df = pd.DataFrame({"y": y, "d": d, "x1": x1, "x2": x2})
>>> res = sp.cluster_cate(
...     df, y="y", treat="d", covariates=["x1", "x2"],
...     n_clusters=3, seed=0)
>>> isinstance(res, sp.ClusterCATEResult)
True
>>> res.cluster_table.columns.tolist()
['cluster', 'n', 'cate', 'se', 'ci_low', 'ci_high']
>>> bool(res.n_obs == 400)
True

FunctionalCATEResult dataclass

Bases: ResultProtocolMixin

Output of FOCaL functional CATE estimator.

Returned by :func:focal_cate. Holds the estimated CATE surface cate_grid of shape (n_test, n_function_points) together with se_grid, the HC1 standard error of each linear CATE prediction (nuisances treated as known), and the function-evaluation grid.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 300
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> d = rng.integers(0, 2, size=n)
>>> tau = 1.0 + 0.5 * x1
>>> y0 = 2.0 + x2 + rng.normal(scale=0.5, size=n)
>>> y1 = y0 + tau
>>> df = pd.DataFrame({
...     "y_t0": np.where(d == 1, y1, y0),
...     "y_t1": np.where(d == 1, y1 + 0.3, y0),
...     "d": d, "x1": x1, "x2": x2,
... })
>>> res = sp.focal_cate(
...     df, y_columns=["y_t0", "y_t1"], treat="d",
...     covariates=["x1", "x2"], seed=0,
... )
>>> isinstance(res, sp.FunctionalCATEResult)
True
>>> res.cate_grid.shape  # (n_test, n_function_points)
(300, 2)

DRLearner

DR-Learner (Kennedy 2023): doubly robust CATE estimation.

Constructs the doubly robust pseudo-outcome: phi(X) = mu_1(X) - mu_0(X) + D(Y - mu_1(X)) / e(X) - (1-D)(Y - mu_0(X)) / (1-e(X))

Then regresses phi on X to obtain tau(X).

Achieves oracle rates and is robust to mis-specification of either the outcome or propensity model (but not both).

Parameters:

Name	Type	Description	Default
outcome_model	sklearn estimator	Model for mu_d(X) = E[Y|X, D=d].	None
propensity_model	sklearn estimator	Model for e(X) = P(D=1|X).	None
cate_model	sklearn estimator	Final-stage model for tau(X). Default: gradient boosting with at least 20 observations per leaf.	None
n_folds	int	Cross-fitting folds for nuisance estimation.	5
fold_indices	array-like of int	Explicit cross-fitting partition, one label in 0..n_folds-1 per row passed to :meth:fit. Default None keeps the historical KFold(n_folds, shuffle=True, random_state=42).	None

Examples:

>>> import statspai as sp
>>> import numpy as np
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> X = rng.normal(size=(n, 3))
>>> e = 1 / (1 + np.exp(-X[:, 0]))
>>> D = (rng.random(n) < e).astype(float)
>>> Y = X[:, 0] + (1.0 + X[:, 1]) * D + rng.normal(size=n)
>>> learner = sp.DRLearner(n_folds=3).fit(X, Y, D)
>>> cate = learner.effect(X)
>>> bool(cate.shape == (n,))
True

References

kennedy2023towards

RLearner

R-Learner (Nie & Wager 2021).

Based on the Robinson (1988) decomposition: Y - m(X) = tau(X) * (D - e(X)) + epsilon

Estimates nuisance functions m(X) = E[Y|X] and e(X) = E[D|X] via cross-fitting, then minimises the R-loss: L(tau) = sum_i [ (Y_i - m_hat(X_i)) - tau(X_i)*(D_i - e_hat(X_i)) ]^2

Achieves quasi-oracle rates under mild conditions.

The decomposition does not need a binary treatment. With a continuous D the same loss is minimised with e(X) = E[D|X] fitted by a regressor, and tau(X) is the effect of one more unit of D at X under an outcome that is linear in D given X: the conditional version of the partially linear model of double machine learning.

Parameters:

Name	Type	Description	Default
outcome_model	sklearn estimator	Model for m(X) = E[Y|X].	None
propensity_model	sklearn estimator	Model for e(X) = E[D|X]: a classifier for a 0/1 treatment (default gradient-boosted classifier), a regressor for a continuous one (default: the outcome model's default regressor).	None
cate_model	sklearn estimator	Model for tau(X), fitted to the pseudo-outcome with weights (D - e(X))**2. Default: gradient boosting with at least 20 observations per leaf.	None
n_folds	int	Cross-fitting folds for nuisance estimation.	5
fold_indices	array-like of int	Explicit cross-fitting partition, one label in 0..n_folds-1 per row passed to :meth:fit. Default None keeps the historical KFold(n_folds, shuffle=True, random_state=42).	None

Examples:

>>> import statspai as sp
>>> import numpy as np
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> X = rng.normal(size=(n, 3))
>>> e = 1 / (1 + np.exp(-X[:, 0]))
>>> D = (rng.random(n) < e).astype(float)
>>> Y = X[:, 0] + (1.0 + X[:, 1]) * D + rng.normal(size=n)
>>> learner = sp.RLearner(n_folds=3).fit(X, Y, D)
>>> cate = learner.effect(X)
>>> bool(cate.shape == (n,))
True

References

nie2021quasi

SLearner

S-Learner: single model, treatment as a feature.

Fits one model mu(X, D) and estimates CATE as: tau(x) = mu(x, 1) - mu(x, 0)

Simple but may under-regularise the treatment effect when the treatment variable is just one of many features.

Parameters:

Name	Type	Description	Default
model	sklearn estimator	Outcome model mu(X, D). Default: GradientBoostingRegressor.	None

Examples:

>>> import statspai as sp
>>> import numpy as np
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> X = rng.normal(size=(n, 3))
>>> D = (rng.random(n) < 0.5).astype(float)
>>> Y = X[:, 0] + (1.0 + X[:, 1]) * D + rng.normal(size=n)
>>> learner = sp.SLearner().fit(X, Y, D)
>>> cate = learner.effect(X)
>>> bool(cate.shape == (n,))
True

References

kunzel2019metalearners

fit

fit(X: Any, Y: Any, D: Any) -> SLearner

Fit mu(X, D).

effect

effect(X: Any) -> ndarray

Estimate CATE: mu(X,1) - mu(X,0).

TLearner

T-Learner: separate models for each treatment arm.

Fits mu_1(X) on treated units and mu_0(X) on controls: tau(x) = mu_1(x) - mu_0(x)

Simple and flexible but can suffer from regularisation imbalance when treatment/control group sizes differ substantially.

Parameters:

Name	Type	Description	Default
model_0	sklearn estimator	Control outcome model. Default: GradientBoostingRegressor.	None
model_1	sklearn estimator	Treated outcome model. Default: same type as model_0.	None

Examples:

>>> import statspai as sp
>>> import numpy as np
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> X = rng.normal(size=(n, 3))
>>> D = (rng.random(n) < 0.5).astype(float)
>>> Y = X[:, 0] + (1.0 + X[:, 1]) * D + rng.normal(size=n)
>>> learner = sp.TLearner().fit(X, Y, D)
>>> cate = learner.effect(X)
>>> bool(cate.shape == (n,))
True

References

kunzel2019metalearners

XLearner

X-Learner (Künzel et al. 2019).

Two-stage procedure: 1. Fit mu_0, mu_1 (T-Learner first stage). 2. Impute individual treatment effects: - For treated: D1_i = Y_i - mu_0(X_i) - For controls: D0_i = mu_1(X_i) - Y_i 3. Fit CATE models tau_1(X) on D1 and tau_0(X) on D0. 4. Combine: tau(x) = e(x)tau_0(x) + (1-e(x))tau_1(x) where e(x) is the propensity score.

Particularly effective when treatment/control groups are very unbalanced.

Parameters:

Name	Type	Description	Default
model_0	sklearn estimator	Control outcome model.	None
model_1	sklearn estimator	Treated outcome model.	None
cate_model_0	sklearn estimator	CATE model for control-side imputed effects.	None
cate_model_1	sklearn estimator	CATE model for treated-side imputed effects.	None
propensity_model	sklearn estimator	Model for e(x) = P(D=1|X).	None

Examples:

>>> import statspai as sp
>>> import numpy as np
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> X = rng.normal(size=(n, 3))
>>> e = 1 / (1 + np.exp(-X[:, 0]))
>>> D = (rng.random(n) < e).astype(float)
>>> Y = X[:, 0] + (1.0 + X[:, 1]) * D + rng.normal(size=n)
>>> learner = sp.XLearner().fit(X, Y, D)
>>> cate = learner.effect(X)
>>> bool(cate.shape == (n,))
True

References

kunzel2019metalearners

cate_gain_curve

cate_gain_curve(data: DataFrame, cate: Any, y: str, treat: str, n_steps: int = 100, n_quantiles: int = 10, normalize: bool = True, ascending: bool = False, n_boot: int = 0, cluster: Optional[str] = None, alpha: float = 0.05, seed: Optional[int] = None) -> CATEGainCurveResult

Cumulative gain curve of a CATE ranking on randomized data.

The sample is sorted by the prediction, largest first. At each step the effect among the top-ranked units is the OLS slope of the outcome on the treatment in that subset (a difference in means when the treatment is 0/1), and the gain is that slope times the share of the sample used. A ranking that puts responsive units first has a curve above the straight line from zero to the average effect. The normalized curve subtracts that line, and its sum over the steps is the reported area.

Parameters:

Name	Type	Description	Default
data	DataFrame	Evaluation data in which treat was randomized. It should not be the data the predictions were fitted on.	required
cate	str or array - like	Predicted effects (or any score that ranks units by effect): a column name or one value per row of data.	required
y	str	Outcome column.	required
treat	str	Treatment column, binary or continuous.	required
n_steps	int	Number of points on the curve.	100
n_quantiles	int	Number of groups in the effect-by-quantile table.	10
normalize	bool	Subtract the random-ranking line from the gain curve.	True
ascending	bool	Rank the smallest predictions first (for an effect where lower is better).	False
n_boot	int	Bootstrap replications for the standard error and the percentile interval of the area. 0 computes neither.	0
cluster	str	Column whose groups are resampled whole (repeated days of one store, say). Default: rows are resampled.	None
alpha	float	Level of the bootstrap interval.	0.05
seed	int	Seed for the bootstrap.	None

Returns:

Type	Description
CATEGainCurveResult

Notes

The slope identifies the effect in a subset only if the treatment is independent of the potential outcomes there, which randomization delivers and observational data does not. On observational data the curve ranks units by a confounded association. The slope also averages a nonlinear dose response with weights that depend on the distribution of the dose in the subset. The bootstrap of n_boot resamples the evaluation data with the predictions held fixed, so it describes the sampling error of the evaluation, not that of the model behind the predictions.

Examples:

>>> import numpy as np, pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 4000
>>> x = rng.uniform(0, 1, n)
>>> dose = rng.uniform(0, 10, n)
>>> df = pd.DataFrame({"x": x, "dose": dose,
...                    "y": (1 + 2 * x) * dose + rng.normal(size=n)})
>>> good = sp.cate_gain_curve(df, cate="x", y="y", treat="dose")
>>> noise = sp.cate_gain_curve(df, cate=rng.normal(size=n), y="y",
...                            treat="dose")
>>> bool(good.auc > abs(noise.auc))
True

blp_test

blp_test(result: CausalResult, data: DataFrame, y: str, treat: str, covariates: List[str], n_folds: int = 5, alpha: float = 0.05, *, propensity: Any = None, baseline: Any = None, proxy: Any = 'cross_fit', vce: str = 'HC1', seed: int = 42) -> Dict[str, Any]

Best Linear Predictor (BLP) of the CATE, Chernozhukov-Demirer-Duflo-Fernandez-Val.

Weighted least squares with weights w = 1 / (p(Z)(1 - p(Z)))::

Y = a0 + a1 B(Z) + beta_1 (D - p(Z))
    + beta_2 (D - p(Z)) (S(Z) - mean(S)) + e

where S(Z) is the CATE proxy, p(Z) the propensity score and B(Z) an optional baseline proxy (e.g. a prediction of E[Y | Z, D = 0]). beta_1 is the ATE; beta_2 is the slope of the true CATE on the proxy -- beta_2 = 0 means the proxy carries no heterogeneity signal, beta_2 = 1 a perfectly calibrated one. This is GenericML::BLP and, fed the same Y, D, p, B, S, reproduces it to the floating-point floor.

The proxy must not be fit on the outcomes it is evaluated against: in-sample predictions of a flexible learner correlate with the noise in Y and push beta_2 up even when the effect is constant. By default (proxy='cross_fit') the metalearner behind result is refit out of fold (n_folds, seed) to produce S.

.. versionchanged:: 1.30.0 Now the CDDF weighted regression with an out-of-fold proxy. Earlier releases ran unweighted OLS on the in-sample CATE predictions: with no true heterogeneity the default T-learner then reported beta_2 near 1.4 with p < 1e-50. The docstring's claim of equivalence to grf::test_calibration was also wrong; that test is :func:sp.calibration_test.

Parameters:

Name	Type	Description	Default
result	CausalResult, fitted CATE model, or array	A metalearner() result (its learner is refit out of fold), or per-unit CATE predictions used as supplied.	required
data	DataFrame		required
y	str	Outcome and binary treatment columns.	required
treat	str	Outcome and binary treatment columns.	required
covariates	list of str	Covariates Z (for the propensity and the proxy refit).	required
n_folds	int	Folds for the propensity and proxy cross-fits.	5
alpha	float		0.05
propensity	array or column name	Known or estimated p(Z). Default: 5-fold gradient-boosting classifier, clipped to [0.01, 0.99]. In a randomized experiment pass the design probabilities.	None
baseline	array or column name	Baseline proxy B(Z) added as a control.	None
proxy	(cross_fit, in_sample)	Source of S(Z).	'cross_fit'
vce	(HC1, const, HC0, HC2, HC3)	Covariance of the WLS fit (sandwich::vcovHC types; GenericML's default is 'const').	'HC1'
seed	int	Fold seed for the proxy refit.	42

Returns:

Type	Description
dict	beta1 / beta2 with _se, _pvalue (two-sided), _pvalue_right (one-sided, H1: beta > 0) and _ci; heterogeneity_significant (two-sided test of beta2 = 0 at alpha); proxy_source, propensity_source, vcov_type.

References

Chernozhukov, Demirer, Duflo and Fernandez-Val (2025), Econometrica 93(4), 1121-1164. [@chernozhukov2025generic]

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 800
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> d = rng.integers(0, 2, size=n)
>>> y = 0.5 * x1 + 0.3 * x2 + (1.0 + x1) * d + rng.normal(size=n)
>>> df = pd.DataFrame({"y": y, "d": d, "x1": x1, "x2": x2,
...                    "p": np.full(n, 0.5)})
>>> res = sp.metalearner(df, y="y", treat="d",
...                      covariates=["x1", "x2"], learner="t")
>>> blp = sp.blp_test(res, df, y="y", treat="d",
...                   covariates=["x1", "x2"], propensity="p")
>>> blp["heterogeneity_significant"]
True
>>> blp["proxy_source"]
'cross_fit'

cate_by_group

cate_by_group(result: Any, data: DataFrame, by: str, n_groups: int = 4, alpha: float = 0.05) -> DataFrame

Group-level average treatment effects.

Splits the CATE distribution by a covariate (or by CATE quartiles if by='cate') and reports group means with confidence intervals.

Parameters:

Name	Type	Description	Default
result	CausalResult, CausalForest, or array-like	Result from metalearner() / tarnet(), a fitted causal_forest() model, or a raw array of per-unit CATE estimates.	required
data	DataFrame	Original data (same rows as the estimation sample).	required
by	str	Column name to group by, or 'cate' to group by CATE quartiles.	required
n_groups	int	Number of quantile groups when by='cate' or when the grouping variable is continuous.	4
alpha	float	Significance level for CIs.	0.05

Returns:

Type	Description
DataFrame	Columns: group, n, mean_cate, se, ci_lower, ci_upper.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 400
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> d = rng.integers(0, 2, size=n)
>>> y = 0.5 * x1 + 0.3 * x2 + (1.0 + x1) * d + rng.normal(size=n)
>>> df = pd.DataFrame({"y": y, "d": d, "x1": x1, "x2": x2})
>>> res = sp.metalearner(df, y="y", treat="d",
...                      covariates=["x1", "x2"], learner="t")
>>> groups = sp.cate_by_group(res, df, by="cate", n_groups=4)
>>> groups["n"].tolist()
[100, 100, 100, 100]
>>> round(float(groups["mean_cate"].iloc[-1]), 2)
2.7

cate_group_plot

cate_group_plot(group_df: DataFrame, ax: Any = None, figsize: tuple = (8, 5), color: str = '#2C3E50', title: Optional[str] = None) -> Tuple[Any, Any]

Plot group-level CATEs with confidence intervals.

Parameters:

Name	Type	Description	Default
group_df	DataFrame	Output from cate_by_group().	required
ax	matplotlib Axes		None
figsize	tuple		(8, 5)
color	str		'#2C3E50'
title	str		None

Returns:

Type	Description
(fig, ax)

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 400
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> d = rng.integers(0, 2, size=n)
>>> y = 0.5 * x1 + 0.3 * x2 + (1.0 + x1) * d + rng.normal(size=n)
>>> df = pd.DataFrame({"y": y, "d": d, "x1": x1, "x2": x2})
>>> res = sp.metalearner(df, y="y", treat="d",
...                      covariates=["x1", "x2"], learner="t")
>>> groups = sp.cate_by_group(res, df, by="cate", n_groups=4)
>>> fig, ax = sp.cate_group_plot(groups)

cate_plot

cate_plot(result: CausalResult, kind: str = 'hist', ax: Any = None, figsize: tuple = (8, 5), color: str = '#2C3E50', title: Optional[str] = None, **kwargs: Any) -> Tuple[Any, Any]

Plot the CATE distribution.

Parameters:

Name	Type	Description	Default
result	CausalResult	Result from metalearner().	required
kind	str	'hist' for histogram, 'kde' for kernel density, 'both'.	'hist'
ax	matplotlib Axes		None
figsize	tuple		(8, 5)
color	str		'#2C3E50'
title	str		None

Returns:

Type	Description
(fig, ax)

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 400
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> d = rng.integers(0, 2, size=n)
>>> y = 0.5 * x1 + 0.3 * x2 + (1.0 + x1) * d + rng.normal(size=n)
>>> df = pd.DataFrame({"y": y, "d": d, "x1": x1, "x2": x2})
>>> res = sp.metalearner(df, y="y", treat="d",
...                      covariates=["x1", "x2"], learner="t")
>>> fig, ax = sp.cate_plot(res, kind="hist")

cate_summary

cate_summary(result: CausalResult) -> DataFrame

Descriptive statistics of the CATE distribution.

Parameters:

Name	Type	Description	Default
result	CausalResult	Result from metalearner() containing model_info['cate'].	required

Returns:

Type	Description
DataFrame	Summary statistics of the fitted conditional effects: mean, sd, min, q25, median, q75, max, fraction positive, IQR, N. The "Mean" row is the plain average of the fitted CATEs; it is not the estimator's doubly-robust average effect (use average_treatment_effect() on a forest), and for a continuous treatment it averages partial effects. It was labelled "Mean (ATE)" before 1.32.0.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 400
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> d = rng.integers(0, 2, size=n)
>>> y = 0.5 * x1 + 0.3 * x2 + (1.0 + x1) * d + rng.normal(size=n)
>>> df = pd.DataFrame({"y": y, "d": d, "x1": x1, "x2": x2})
>>> res = sp.metalearner(df, y="y", treat="d",
...                      covariates=["x1", "x2"], learner="t")
>>> tab = sp.cate_summary(res)
>>> int(tab.loc["N", "CATE"])
400
>>> round(float(tab.loc["Frac. Positive", "CATE"]), 2)
0.78

compare_metalearners

compare_metalearners(data: DataFrame, y: str, treat: str, covariates: List[str], learners: Optional[List[str]] = None, **kwargs: Any) -> DataFrame

Fit multiple meta-learners and compare their ATE estimates.

Parameters:

Name	Type	Description	Default
data	DataFrame	Input data.	required
y	str	Outcome variable.	required
treat	str	Binary treatment variable (0/1).	required
covariates	list of str	Covariate / effect modifier variables.	required
learners	list of str	Which learners to compare. Default: all five ('s','t','x','r','dr').	None
**kwargs	Any	Additional arguments passed to metalearner().	{}

Returns:

Type	Description
DataFrame	Comparison table with columns: learner, ate, se, ci_lower, ci_upper, pvalue, cate_std, cate_iqr.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 400
>>> age = rng.normal(size=n)
>>> edu = rng.normal(size=n)
>>> training = rng.integers(0, 2, size=n)
>>> wage = 0.5 * age + 0.3 * edu + (1.0 + age) * training + rng.normal(size=n)
>>> df = pd.DataFrame({"wage": wage, "training": training,
...                    "age": age, "edu": edu})
>>> comp = sp.compare_metalearners(df, y='wage', treat='training',
...                                 covariates=['age', 'edu'])
>>> "ate" in comp.columns
True
>>> int(comp.shape[0])
5

gate_test

gate_test(result: CausalResult, data: DataFrame, by: str, n_groups: int = 4, alpha: float = 0.05, *, y: Optional[str] = None, treat: Optional[str] = None, covariates: Optional[List[str]] = None, propensity: Any = None, baseline: Any = None, proxy: Any = 'cross_fit', vce: str = 'HC1', n_folds: int = 5, seed: int = 42, clan: Optional[List[str]] = None) -> Dict[str, Any]

Sorted group average treatment effects (GATES) and heterogeneity tests.

With y, treat and covariates this is the Chernozhukov-Demirer-Duflo-Fernandez-Val GATES regression (GenericML::GATES with monotonize = FALSE): weighted least squares with weights 1 / (p(1 - p)) of Y on an intercept, the optional baseline proxy B(Z) and (D - p) 1{G_k} for groups G_1..G_K. Groups are quantile groups of the CATE proxy when by='cate' (GenericML::quantile_group: type-7 cut points, left-closed) and of the column by otherwise (its values are used directly when it has at most n_groups levels). Each gamma_k is the ATE in group k, estimated from outcomes; top_vs_bottom is gamma_K - gamma_1 with the covariance term; the omnibus test is the Wald test that all gamma_k are equal (chi-square, K - 1 df). The proxy defaults to an out-of-fold refit of the result's learner, as in :func:blp_test.

Without them the function falls back to the pre-1.30 descriptive summary: group means of the predicted CATEs with sd / sqrt(n) SEs. Those p-values treat the model's own predictions as data -- with by='cate' the groups are formed from the same predictions -- so they are not a test of heterogeneity; a warning says so.

.. versionchanged:: 1.30.0 Added the GATES regression (y / treat / covariates).

Parameters:

Name	Type	Description	Default
result	CausalResult, fitted CATE model, or array		required
data	DataFrame		required
by	str	'cate' or a column name.	required
n_groups	int		4
alpha	float		0.05
y	optional	Outcome, binary treatment and covariates; required for GATES.	None
treat	optional	Outcome, binary treatment and covariates; required for GATES.	None
covariates	optional	Outcome, binary treatment and covariates; required for GATES.	None
propensity	Any	As in :func:blp_test.	None
baseline	Any	As in :func:blp_test.	None
proxy	Any	As in :func:blp_test.	None
vce	Any	As in :func:blp_test.	None
n_folds	Any	As in :func:blp_test.	None
seed	Any	As in :func:blp_test.	None
clan	list of str	Classification analysis (CLAN; R GenericML::CLAN): for each of these columns, its mean in every group (standard error sd / sqrt(n_k)) and the difference between the most and the least affected group (Welch standard error), with normal intervals. Returned as clan_table. GATES path only. Equals GenericML::CLAN given the same groups to 1e-12 (tests/reference_parity/test_clan_genericml_parity.py).	None

Returns:

Type	Description
dict	gate_table (group, n, gate, se, ci_lower, ci_upper for GATES; group, n, mean_cate, se, ci_lower, ci_upper for the descriptive path), omnibus_stat / omnibus_pvalue (omnibus_F on the descriptive path), top_vs_bottom_diff / _se / _pvalue, and method.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 800
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> d = rng.integers(0, 2, size=n)
>>> y = 0.5 * x1 + 0.3 * x2 + (1.0 + x1) * d + rng.normal(size=n)
>>> df = pd.DataFrame({"y": y, "d": d, "x1": x1, "x2": x2,
...                    "p": np.full(n, 0.5)})
>>> res = sp.metalearner(df, y="y", treat="d",
...                      covariates=["x1", "x2"], learner="t")
>>> gt = sp.gate_test(res, df, by="cate", y="y", treat="d",
...                   covariates=["x1", "x2"], propensity="p")
>>> gt["method"]
'gates'
>>> bool(gt["top_vs_bottom_pvalue"] < 0.05)
True

predict_cate

predict_cate(result: CausalResult, new_data: DataFrame) -> ndarray

Predict CATE on new (out-of-sample) data.

Parameters:

Name	Type	Description	Default
result	CausalResult	Result from metalearner() containing a fitted estimator.	required
new_data	DataFrame	New data with the same covariate columns used in estimation.	required

Returns:

Type	Description
ndarray	Predicted CATE for each row of new_data.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 400
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> d = rng.integers(0, 2, size=n)
>>> y = 0.5 * x1 + 0.3 * x2 + (1.0 + x1) * d + rng.normal(size=n)
>>> df = pd.DataFrame({"y": y, "d": d, "x1": x1, "x2": x2})
>>> res = sp.metalearner(df, y="y", treat="d",
...                      covariates=["x1", "x2"], learner="t")
>>> new_data = pd.DataFrame({"x1": [-1.0, 0.0, 1.0],
...                          "x2": [0.0, 0.0, 0.0]})
>>> tau_hat = sp.predict_cate(res, new_data)
>>> int(tau_hat.shape[0])
3
>>> bool(tau_hat[2] > tau_hat[0])
True

focal_cate

focal_cate(data: DataFrame, y_columns: List[str], treat: str, covariates: List[str], test_data: Optional[DataFrame] = None, seed: int = 0) -> FunctionalCATEResult

Functional doubly-robust CATE estimator.

Parameters:

Name	Type	Description	Default
data	DataFrame	Training data with outcome columns y_columns (one per function evaluation point t).	required
y_columns	list of str	Outcome columns; len = number of function points.	required
treat	str		required
covariates	list of str		required
test_data	DataFrame	Defaults to data.	None
seed	int		0

Returns:

Type	Description
FunctionalCATEResult

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 400
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> d = rng.integers(0, 2, size=n)
>>> tau = 1.0 + 0.5 * x1  # heterogeneous effect
>>> y0 = 2.0 + x2 + rng.normal(scale=0.5, size=n)
>>> y1 = y0 + tau
>>> df = pd.DataFrame({
...     "y_t0": np.where(d == 1, y1, y0),
...     "y_t1": np.where(d == 1, y1 + 0.3, y0),
...     "d": d, "x1": x1, "x2": x2,
... })
>>> res = sp.focal_cate(
...     df, y_columns=["y_t0", "y_t1"], treat="d",
...     covariates=["x1", "x2"], seed=0,
... )
>>> res.cate_grid.shape  # (n_test, n_function_points)
(400, 2)

metalearner

metalearner(data: DataFrame, y: str, treat: Optional[str] = None, covariates: Optional[List[str]] = None, learner: str = 'dr', outcome_model: Optional[Any] = None, propensity_model: Optional[Any] = None, cate_model: Optional[Any] = None, n_folds: int = 5, n_bootstrap: int = 200, alpha: float = 0.05, d: Optional[str] = None, x: Optional[List[str]] = None, fold_indices: Optional[Any] = None, weights: Optional[str] = None, cluster: Optional[str] = None) -> CausalResult

Estimate heterogeneous treatment effects using meta-learners.

Parameters:

Name	Type	Description	Default
data	DataFrame	Input data.	required
y	str	Outcome variable.	required
treat	str	Binary treatment variable (0/1). learner='r' also takes a continuous treatment; see Notes.	None
covariates	list of str	Covariate / effect modifier variables. A text or category column is expanded to indicator columns.	None
learner	str	Meta-learner type: 's', 't', 'x', 'r', or 'dr'.	'dr'
outcome_model	sklearn estimator or (control, treated) pair	Custom ML model for outcome nuisance. For learner='t' and 'x' a pair (control_model, treated_model) fits a different learner in each arm, which is the point of those two learners when the arms differ in size or in how smooth the response is (Kunzel et al. fit a forest to the large arm and a linear model to the small one). The pair is also used, arm by arm, in the AIPW scores behind estimate and se.	None
propensity_model	sklearn estimator	Custom propensity score model (used by X/R/DR learners).	None
cate_model	sklearn estimator or (control, treated) pair	Custom model for final CATE stage (X/R/DR learners). The default of the R- and DR-learners is gradient boosting with at least 20 observations per leaf, because the pseudo-outcomes they regress on have heavy tails. For learner='x' a pair (model_fit_on_controls, model_fit_on_treated) sets the two second-stage regressions separately.	None
n_folds	int	Cross-fitting folds for nuisance estimation (used by R/DR learners and the unified AIPW SE path; see Notes).	5
n_bootstrap	int	Deprecated and ignored as of v1.11.4. Previously the SE for S/T/X/R-Learner came from a re-sampling bootstrap of the fitted CATE values, which treats τ̂ as fixed and severely under-estimates uncertainty. The function now uses the AIPW influence function for SE regardless of learner=. The argument is kept for backward compatibility and will be removed in a future minor release.	200
alpha	float	Significance level.	0.05
fold_indices	array-like of int	Explicit cross-fitting partition: one label per row of data, the labels being exactly 0, ..., n_folds-1 (the convention sp.dml(fold_indices=...) uses); rows dropped for missing values are dropped from the vector too. It replaces every internal split on this code path -- the R-/DR-Learner nuisance cross-fit and the AIPW cross-fit behind estimate / se for the S/T/X/R learners. Each training complement must contain treated and control units. The default None keeps the historical KFold(n_folds, shuffle=True, random_state=42); passing that split's labels reproduces the default result exactly.	None
weights	str	Sampling-weight column. The AIPW score behind estimate / se is then a weighted cross-fit (weighted outcome and propensity fits, weighted mean, weighted influence function) for every learner=. The CATE function itself is fitted unweighted.	None
cluster	str	Cluster column. Cross-fitting folds are formed over whole clusters and the SE sums the influence function within clusters, centred, with G/(G-1) -- which reduces to the default SE when every row is its own cluster.	None

Returns:

Type	Description
CausalResult	Result with ATE estimate, SE, CI, p-value, and individual CATE predictions accessible via result.model_info['cate']. model_info['estimation_rows'] holds the positions, in data, of the rows behind cate (complete cases); the CATE diagnostics (sp.cate_by_group, sp.blp_test, sp.gate_test) use it to pair each effect with its own row.

Notes

ATE / SE convention (v1.11.4+). Regardless of which CATE estimator the user selects via learner=, the population ATE and its SE are computed via the AIPW (DR) pseudo-outcome:

.. math::

\varphi_i = \hat\mu_1(X_i) - \hat\mu_0(X_i) + \frac{D_i (Y_i - \hat\mu_1(X_i))}{\hat e(X_i)} - \frac{(1-D_i)(Y_i - \hat\mu_0(X_i))}{1 - \hat e(X_i)}

with :math:\hat{\rm ATE} = \bar\varphi, :math:\widehat{\rm SE} = \sigma(\varphi)/\sqrt n. AIPW is the semiparametric-efficient estimating function for :math:E[Y(1) - Y(0)] (van der Laan & Robins 2003; Kennedy 2023), so the SE is valid for any CATE estimator. The chosen learner= determines τ̂(X) (heterogeneity); ATE inference is learner-independent.

Prior to v1.11.4, S/T/X/R-Learner used mean(τ̂) as ATE and a re-sampling bootstrap of τ̂ as SE. The bootstrap silently treated τ̂ as fixed → systematically too small SEs and severe under- coverage. ⚠️ This is a correctness fix; numerical results will change for non-DR learners.

What se does and does not describe. A consequence of the convention above is that result.estimate and result.se are numerically identical across learner='s', 't', 'x' and 'dr' when the same outcome and propensity models are supplied: they describe the AIPW average, not the chosen learner. The learner-specific quantity is the fitted effect function result.model_info['cate'] and its mean result.model_info['cate_mean'], for which no standard error is provided (model_info['cate_mean_se'] is None). Pairing cate_mean with se produces an interval for the wrong functional; the learner-controlled Monte Carlo in the ML4CI companion paper shows the resulting coverage number measures the mismatch rather than the estimator. To compare learners at the average, compare cate_mean across calls and bootstrap the whole procedure for its uncertainty.

Continuous treatment (learner='r'). The R-loss :math:\sum_i [(Y_i - \hat m(X_i)) - \tau(X_i)(D_i - \hat e(X_i))]^2 is defined for any real-valued treatment, with :math:\hat e(X) = \hat E[D \mid X] fitted by a regressor (propensity_model, default the outcome model's default). Then model_info['cate'] is the effect of one more unit of the treatment at each X, under an outcome linear in the treatment given X. estimate is the partially linear coefficient :math:\sum_i \tilde D_i \tilde Y_i / \sum_i \tilde D_i^2 on the cross-fitted residuals (estimand='APE'): the average of :math:\tau(X) weighted by the conditional variance of the treatment, with the standard error of its influence function. There are no propensity scores to clip, so the overlap diagnostics are empty.

References

Künzel, 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. [@kunzel2019metalearners]

Kennedy, E. H. (2023). Towards optimal doubly robust estimation of heterogeneous causal effects. Electronic Journal of Statistics, 17(2), 3008-3049. [@kennedy2023towards]

Examples:

>>> import statspai as sp
>>> import numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> n = 300
>>> age = rng.normal(40, 8, n)
>>> edu = rng.normal(13, 2, n)
>>> training = (rng.random(n) < 0.5).astype(int)
>>> wage = 0.05 * age + 0.3 * edu + (2.0 + 0.1 * edu) * training \
...        + rng.normal(size=n)
>>> df = pd.DataFrame({"wage": wage, "training": training,
...                    "age": age, "edu": edu})
>>> result = sp.metalearner(df, y='wage', treat='training',
...                         covariates=['age', 'edu'],
...                         learner='dr', n_folds=3)
>>> bool(isinstance(result.estimate, float))
True
>>> cate = result.model_info['cate']  # per-unit effects
>>> bool(cate.shape == (n,))
True