`statspai.causal`¶

causal ¶

DEPRECATED — use :mod:statspai.forest instead.

This package was renamed in v1.10 to statspai.forest because it only ever housed the four forest-based causal estimators (causal_forest / iv_forest / multi_arm_forest / forest inference helpers). The causal name was misleading — it implied a top-level causal-inference namespace when in fact the content is a single-method family.

This shim keeps both from statspai.causal import X and from statspai.causal.causal_forest import X working by aliasing the deprecated submodule paths through :data:sys.modules to the real :mod:statspai.forest modules. Users see one :class:DeprecationWarning on first import; everything else continues working.

Plan to migrate within one minor cycle.

CausalForest ¶

Bases: BaseModel

Causal Forest for heterogeneous treatment effect estimation

This class implements the Causal Forest algorithm, which uses random forests to estimate conditional average treatment effects (CATE) in a non-parametric way.

The method combines ideas from: 1. Honest estimation to avoid overfitting 2. Double machine learning to handle confounding 3. Random forests for flexible function approximation

Parameters:

Name	Type	Description	Default
`n_estimators`	`int`	Number of trees (grf `num.trees`). Rounded up to a multiple of `ci_group_size`.	`2000`
`min_samples_leaf`	`int`	grf `min.node.size`: a node with at most this many growing samples is not split, and each child of a stabilised split must hold at least this many treated-like and control-like samples.	`5`
`max_depth`	`int`	Optional depth cap (grf has none).	`None`
`max_samples`	`float`	Fraction of clusters drawn per tree (grf `sample.fraction`); at most 0.5 when `ci_group_size > 1`.	`0.5`
`model_y`	`estimator`	scikit-learn estimators for `E[Y \| X, W]` and `E[T \| X, W]`. When omitted, grf regression forests with out-of-bag predictions are used. Supplied estimators are cross-fitted in `nuisance_folds` folds that never split a cluster.	`None`
`model_t`	`estimator`	scikit-learn estimators for `E[Y \| X, W]` and `E[T \| X, W]`. When omitted, grf regression forests with out-of-bag predictions are used. Supplied estimators are cross-fitted in `nuisance_folds` folds that never split a cluster.	`None`
`discrete_treatment`	`bool`	Whether treatment is binary (0/1). Continuous treatments are supported by the causal forest itself; doubly-robust averages require a binary treatment.	`True`
`honest`	`bool`	grf `honesty`: grow the tree on one half of the drawn clusters and estimate leaves on the other.	`True`
`bootstrap`	`None`	Not used by the grf engine (trees are grown on subsamples drawn without replacement). Retained for the legacy engine.	`None`
`random_state`	`int`	Seed. Results are reproducible for a given seed and `n_jobs`- independent.	`None`
`n_jobs`	`int`	Threads used to grow trees (-1 = all cores). Results do not depend on it: every tree draws from its own seed and the trees are combined in a fixed order. .. versionchanged:: 1.40.0 The default was `1`. Estimates are byte-identical for any value (`tests/test_treatment_effects_oct19_fixes.py`).	`-1`
`verbose`	`int`	Verbosity level	`0`
`split_rule`	`(grf, cffe, legacy)`	`"cffe"` (`fe=` forests only) replaces the GRF gradient criterion with the tau-heterogeneity criterion of Kattenberg, Scheer and Thiel (2023), the rule the `causalfe` package implements; see :func:`statspai.causal_forest`. `"legacy"` reproduces the pre-1.29 estimator and is deprecated.	`"grf"`
`mtry`	`int`	Mean number of candidate variables per split (grf default `min(ceil(sqrt(p) + 20), p)`).	`None`
`honesty_fraction`	`float`		`0.5`
`honesty_prune_leaves`	`bool`		`True`
`alpha`	`float`	grf `alpha`: minimum share of the parent's treatment variation each child of a stabilised split must retain.	`0.05`
`imbalance_penalty`	`float`		`0.0`
`stabilize_splits`	`bool`		`True`
`ci_group_size`	`int`	Trees per little bag (`1` disables variance estimates). Defaults to 2 when `max_samples <= 0.5` and to 1, with a warning, otherwise.	`None`
`equalize_cluster_weights`	`bool`	grf `equalize.cluster.weights`: draw the same number of observations from every cluster and weight clusters equally in averages.	`False`
`nuisance_folds`	`int`	Cross-fitting folds for user-supplied nuisance estimators.	`3`
`fe`	`(None, unit, twoway)`	Panel fixed effects removed within every node and leaf. Requires `id=` (and `time=` for `"twoway"`) at fit time; trees are then sampled by unit.	`None`

Attributes:

Name	Type	Description
`fitted_`	`bool`	Whether the model has been fitted
`params`	`Series`	Not applicable for non-parametric methods, returns empty Series
`std_errors`	`Series`	Not applicable for non-parametric methods, returns empty Series
`tvalues`	`Series`	Not applicable for non-parametric methods, returns empty Series
`pvalues`	`ndarray`	Not applicable for non-parametric methods, returns empty array
`diagnostics`	`dict`	Model diagnostics and fit statistics
`data_info`	`dict`	Information about the data used in fitting

Notes

This implementation is inspired by the EconML library's CausalForestDML but adapted to fit the StatsPAI architecture and interface.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> from statspai.forest import CausalForest
>>>
>>> # Generate sample data
>>> np.random.seed(42)
>>> n = 1000
>>> X = np.random.normal(0, 1, (n, 3))
>>> T = np.random.binomial(1, 0.5, n)
>>> Y = X[:, 0] * T + X[:, 1] + np.random.normal(0, 1, n)
>>> data = pd.DataFrame({
...     'Y': Y, 'T': T, 'X1': X[:, 0], 'X2': X[:, 1], 'X3': X[:, 2]
... })
>>>
>>> # Fit Causal Forest
>>> cf = CausalForest(n_estimators=50, random_state=42)
>>> cf = cf.fit('Y ~ T | X1 + X2 + X3', data=data)  # returns self
>>> cf
CausalForest(fitted=True, n_estimators=50)
>>>
>>> # Estimate treatment effects
>>> cate = cf.effect(data[['X1', 'X2', 'X3']])
>>> cate.shape
(1000,)
>>> ate = float(cate.mean())  # average treatment effect

fit ¶

fit(formula: Optional[str] = None, data: Optional[DataFrame] = None, Y: Optional[ndarray] = None, T: Optional[ndarray] = None, X: Optional[ndarray] = None, W: Optional[ndarray] = None, clusters: Optional[Any] = None, Y_hat: Optional[Any] = None, W_hat: Optional[Any] = None, id: Optional[Any] = None, time: Optional[Any] = None) -> CausalForest

Fit the Causal Forest model

Parameters:

Name	Type	Description	Default
`formula`	`str`	Formula specification in the form "Y ~ T \| X1 + X2 + ... [\| W1 + W2 + ...]" where Y is outcome, T is treatment, X are effect modifiers, W are covariates	`None`
`data`	`DataFrame`	Data containing all variables if using formula interface	`None`
`Y`	`array - like`	Outcome variable (n_samples,)	`None`
`T`	`array - like`	Treatment variable (n_samples,)	`None`
`X`	`array - like`	Effect modifier variables (n_samples, n_features)	`None`
`W`	`array - like`	Control variables for confounding adjustment (n_samples, n_controls)	`None`
`clusters`	`array - like`	Cluster ids (grf `clusters`). Trees are subsampled, split into honest halves, and cross-fitted by cluster, and averages report cluster-robust standard errors. With `fe` set it defaults to `id`.	`None`
`Y_hat`	`array - like or scalar`	Precomputed nuisances `E[Y \| X, W]` and `E[T \| X, W]`; must be out-of-sample (cross-fitted or out-of-bag) predictions.	`None`
`W_hat`	`array - like or scalar`	Precomputed nuisances `E[Y \| X, W]` and `E[T \| X, W]`; must be out-of-sample (cross-fitted or out-of-bag) predictions.	`None`
`id`	`array - like`	Panel unit and period identifiers for `fe="unit"` / `fe="twoway"` (`unit=` is accepted as an alias of `id`).	`None`
`time`	`array - like`	Panel unit and period identifiers for `fe="unit"` / `fe="twoway"` (`unit=` is accepted as an alias of `id`).	`None`

Returns:

Name	Type	Description
`self`	`CausalForest`	Fitted estimator

get_nuisances ¶

get_nuisances() -> Dict[str, Any]

Return the forest's internal nuisance predictions (read-only copies).

Returns:

Type Description

dict

"Y_hat": E[Y | X, W] and "W_hat": E[T | X, W] (the propensity for a binary treatment) for every training row, exactly the arrays the forest was grown on and the doubly robust averages use -- out-of-bag regression-forest predictions by default, cross-fitted model_y / model_t predictions, or the caller's Y_hat / W_hat. The arrays are copies with writeable=False. "source" says which, and "overlap" is diagnostics["nuisance_overlap"].

Examples:

>>> import numpy as np
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> X = rng.normal(size=(300, 2))
>>> T = rng.binomial(1, 0.5, size=300)
>>> Y = X[:, 0] * T + rng.normal(size=300)
>>> cf = sp.causal_forest(Y=Y, T=T, X=X, n_estimators=100,
...                       random_state=0)
>>> nu = cf.get_nuisances()
>>> nu["W_hat"].shape
(300,)

effect ¶

effect(X: ndarray) -> ndarray

Estimate conditional average treatment effects

Parameters:

Name	Type	Description	Default
`X`	`(array - like, shape(n_samples, n_features))`	Effect modifier variables	required

Returns:

Name	Type	Description
`effects`	`(array - like, shape(n_samples))`	Estimated conditional average treatment effects

predict ¶

predict(data: Optional[DataFrame] = None) -> ndarray

Generate treatment effect predictions (required by BaseModel)

Parameters:

Name	Type	Description	Default
`data`	`DataFrame`	Data containing effect modifier variables. If None, uses training data.	`None`

Returns:

Type	Description
`ndarray`	Predicted treatment effects

Notes

This method is required by the BaseModel interface. For Causal Forest, "predictions" are treatment effect estimates rather than outcome predictions. With data=None the GRF engine returns out-of-bag predictions for the training rows (each row predicted only by trees that did not see it), which is what every in-sample diagnostic should use.

effect_interval ¶

effect_interval(X: ndarray, alpha: float = 0.05) -> Tuple[ndarray, ndarray]

Pointwise confidence intervals for the CATE.

With the GRF engine the interval is tau(x) +/- z * sqrt(V(x)) where V(x) is the bootstrap-of-little-bags variance (requires ci_group_size >= 2). Intervals are pointwise, not uniform.

Parameters:

Name	Type	Description	Default
`X`	`(array - like, shape(n_samples, n_features))`	Effect modifier variables	required
`alpha`	`float`	Significance level (1-alpha is confidence level)	`0.05`

Returns:

Name	Type	Description
`lower`	`(array - like, shape(n_samples))`	Lower bounds of confidence intervals
`upper`	`(array - like, shape(n_samples))`	Upper bounds of confidence intervals

oob_effect ¶

oob_effect() -> ndarray

Out-of-bag CATE predictions for the training rows (GRF engine).

effect_variance ¶

effect_variance(X: Optional[ndarray] = None) -> ndarray

Little-bag variance of the CATE at X (OOB rows when None).

split_frequencies ¶

split_frequencies(max_depth: int = 4) -> DataFrame

Number of splits on each covariate at each depth (GRF engine).

average_treatment_effect ¶

average_treatment_effect(X: Optional[ndarray] = None, T: Optional[ndarray] = None, target_sample: str = 'all', alpha: float = 0.05, clip: float = 0.01, variance: str = 'forest', covariates: Any = 'none', subset: Any = None) -> Dict[str, float]

GRF-style ATE/ATT/ATC/ATO aggregation of CATE predictions.

Forests with fixed effects support target_sample='treated' only, estimated by imputation; variance applies to them (see :func:statspai.average_treatment_effect).

The target may be given positionally, as in grf: cf.average_treatment_effect("treated").

group_effects ¶

group_effects(by: Any = None, *, members: Any = None, n_groups: int = 4, cluster: Any = None, variance: str = 'forest', alpha: float = 0.05, scale: str = 'level', min_rows: int = 1, covariates: Any = 'none') -> DataFrame

Average effects by group; see :func:statspai.forest_group_effects.

forest_diagnostics ¶

forest_diagnostics(X: Optional[ndarray] = None, T: Optional[ndarray] = None, propensity_bounds: Tuple[float, float] = (0.05, 0.95)) -> Dict[str, object]

Overlap and CATE-distribution diagnostics for this fitted forest.

summary ¶

summary() -> str

Return a summary of the fitted model

Returns:

Name	Type	Description
`summary`	`str`	Model summary string

variable_importance ¶

variable_importance(method: Optional[str] = None, decay_exponent: float = 2.0, max_depth: int = 4) -> Series

Variable importance of the causal forest.

Parameters:

Name	Type	Description	Default
`method`	`(split, permutation)`	`"split"` is `grf::variable_importance`: the depth-weighted share of splits on each covariate (see :func:`statspai.variable_importance`). `"permutation"` is the earlier StatsPAI measure: shuffle one covariate and record the mean squared change of `effect(X)` on the training rows (which uses in-bag trees, so it reflects the fitted function, not out-of-sample relevance). A forest grown with `split_rule="legacy"` records no split depths; for it the default is `"permutation"`. .. versionchanged:: 1.39.0 The default is `"split"`. It was `"permutation"`, with a `FutureWarning` announcing the change since 1.31.	`"split"`
`decay_exponent`	`float`	Options of the `"split"` measure.	`2.0`
`max_depth`	`float`	Options of the `"split"` measure.	`2.0`

Returns:

Type	Description
`Series`	Normalised importance, indexed by covariate. The permutation measure sums to 1; the split measure sums to 1 when every counted depth has at least one split, and to less otherwise (a depth without splits contributes zero while its weight stays in the normaliser, grf's convention).

best_linear_projection ¶

best_linear_projection(X_test: Optional[ndarray] = None, alpha: float = 0.05, clip: float = 0.01, covariates: Any = 'none') -> DataFrame

Best Linear Projection (BLP) of CATE on features (Semenova-Chernozhukov 2021).

For GRF-engine forests (the default) this is grf::best_linear_projection: the doubly-robust scores built from the out-of-bag CATE are regressed on (1, X) on the covariates' original scale, with observation weights and an HC3 (cluster-robust when the forest has clusters) covariance. X_test must then be the training rows' effect modifiers. Legacy forests keep the standardised-covariate HC1 regression described below.

Constructs the augmented inverse-propensity-weighted (AIPW) doubly-robust score :math:\Gamma_i and regresses it on :math:X_i with HC1 standard errors:

.. math:: \Gamma_i = \hat{\tau}(X_i) + \frac{T_i - \hat{e}(X_i)}{\hat{e}(X_i)(1-\hat{e}(X_i))} \bigl(Y_i - \hat{m}(X_i) - (T_i - \hat{e}(X_i))\hat{\tau}(X_i)\bigr).

:math:\Gamma_i is unbiased for :math:\tau(X_i) under standard cross-fitting / overlap conditions; OLS of :math:\Gamma_i on :math:(1, X_i) recovers the population BLP coefficients with valid heteroscedasticity-robust inference (HC1).

This replaces the earlier plug-in OLS of :math:\hat{\tau}(X_i) on :math:X_i, which produces anti-conservative SEs (the SE on a fitted model is not the SE on the population BLP).

Parameters:

Name	Type	Description	Default
`X_test`	`array - like`	Features to evaluate the BLP at; defaults to in-sample X.	`None`
`alpha`	`float`	Significance level for CIs (reported alongside coef/SE).	`0.05`
`clip`	`float`	Propensity clip (binary discrete treatment only) to prevent the inverse-propensity term from blowing up under near-violations of overlap. Counts of clipped units are exposed via `self.diagnostics`.	`0.01`

Returns:

Type	Description
`DataFrame`	Index `["Intercept", *features]` with columns `[coef, se, t, p, ci_lower, ci_upper]`. HC1 SEs.

Notes

Forests with fixed effects (fe=) have no propensity, so the response is the imputation score Y - alpha_hat_i - gamma_hat_t of every treated cell (unit and period effects fitted on untreated cells) and the regression runs over treated cells; standard errors come from the exact linear weights of that estimator, clustered by the forest's clusters. covariates adds covariates to that untreated model (see :func:statspai.average_treatment_effect). See :func:statspai.forest_group_effects.

References

[@semenova2021debiased], [@athey2019generalized], [@borusyak2024revisiting]

ate ¶

ate(X: Optional[ndarray] = None) -> ScalarEffect

Average Treatment Effect.

Returns a :class:~statspai.core.results.ScalarEffect — a float whose value is the GRF-style doubly-robust (AIPW) average, the quantity .se, .ci and .pvalue describe, with the full aggregation payload in .detail and the plug-in mean CATE in .detail['plug_in_estimate']. Printing the result shows the inference; arithmetic and formatting behave exactly like a bare float.

.. versionchanged:: 1.29.0 The float value was the plug-in mean CATE before 1.29; it is now the doubly-robust estimate, so the point estimate and its interval refer to the same quantity. See MIGRATION.md.

att ¶

att(X: Optional[ndarray] = None, T: Optional[ndarray] = None) -> ScalarEffect

Average Treatment Effect on the Treated.

Returns a :class:~statspai.core.results.ScalarEffect whose float value is the GRF-style doubly-robust ATT score, with its own inference attached as .se / .ci / .pvalue / .detail and the plug-in mean CATE over treated rows in .detail['plug_in_estimate'].

.. versionchanged:: 1.29.0 The float value was the plug-in average before 1.29. See MIGRATION.md.

.. versionchanged:: 1.30.0 Raises :class:MethodIncompatibility for a non-binary treatment instead of averaging over the rows with T == 1.

IVForestResult ¶

Bases: GRFFamilyForest

A fitted instrumental forest.

Attributes:

Name	Type	Description
`late`	`float`	Average conditional LATE, `E[tau(X)]`, from the doubly-robust scores (grf `average_treatment_effect` on an instrumental forest).
`se, ci, pvalue`	`(float, tuple, float)`	Its (cluster-robust) standard error, confidence interval and p-value.
`cate`	`ndarray`	Out-of-bag `tau(X_i)` for the training rows.
`cate_variance`	`ndarray or None`	Little-bag variance of each out-of-bag `tau(X_i)`.
`n_obs`	`int`
`detail`	`dict`	Nuisance sources, first-stage strength, compliance-score range.

Examples:

>>> import numpy as np, statspai as sp
>>> rng = np.random.default_rng(0)
>>> X = rng.normal(size=(600, 2)); Z = rng.binomial(1, 0.5, 600)
>>> W = (Z + rng.normal(size=600) > 0.5).astype(float)
>>> Y = (1 + X[:, 0]) * W + rng.normal(size=600)
>>> fit = sp.iv_forest(y=Y, treat=W, instrument=Z, covariates=X,
...                    n_estimators=200)
>>> fit.cate.shape
(600,)

predict ¶

predict(newdata: Any = None, estimate_variance: bool = False) -> DataFrame

tau(x) (and optionally its variance) at newdata.

newdata=None returns the out-of-bag predictions for the training rows.

get_scores ¶

get_scores(compliance_score: Any = None, debiasing_weights: Any = None) -> ndarray

Doubly-robust scores Gamma_i for E[tau(X)].

compliance_score (Delta(X_i)) or the full debiasing_weights may be supplied; otherwise the stored estimates from the fit are used.

average_treatment_effect ¶

average_treatment_effect(alpha: float = 0.05, compliance_score: Any = None, debiasing_weights: Any = None) -> Dict[str, Any]

Average conditional LATE with its (cluster-robust) SE.

best_linear_projection ¶

best_linear_projection(A: Any = None, vce: str = 'HC3', alpha: float = 0.05, compliance_score: Any = None, debiasing_weights: Any = None) -> DataFrame

Regress the doubly-robust scores on (1, A) (grf::best_linear_projection).

MultiArmForestResult ¶

Bases: GRFFamilyForest

A fitted multi-arm causal forest.

Attributes:

Name	Type	Description
`arms`	`list`	Arm labels, reference first.
`reference`	`Any`	The reference arm.
`ate, ate_se, ci, pvalue`	`dict`	Doubly-robust average effect of each non-reference arm against the reference (keyed by arm label).
`cate`	`dict`	Out-of-bag `tau_k(X_i)` per non-reference arm.
`cate_variance`	`dict or None`	Little-bag variances of `cate`.
`propensities`	`DataFrame`	Out-of-bag `e_k(X_i)`, one column per arm.

Examples:

>>> import numpy as np, statspai as sp
>>> rng = np.random.default_rng(0)
>>> X = rng.normal(size=(600, 2)); A = rng.integers(0, 3, 600)
>>> Y = X[:, 1] + (A == 1) - 0.5 * (A == 2) + rng.normal(size=600)
>>> fit = sp.multi_arm_forest(y=Y, treat=A, covariates=X, n_estimators=200)
>>> fit.contrasts
[1, 2]

predict ¶

predict(newdata: Any = None, estimate_variance: bool = False) -> DataFrame

tau_k(x) for every contrast (columns "<arm> - <reference>").

get_scores ¶

get_scores() -> ndarray

AIPW scores, (n, K-1), one column per contrast.

average_treatment_effect ¶

average_treatment_effect(alpha: float = 0.05) -> DataFrame

Doubly-robust average effect of each arm against the reference.

best_linear_projection ¶

best_linear_projection(A: Any = None, vce: str = 'HC3', alpha: float = 0.05) -> DataFrame

BLP of each contrast's CATE on (1, A), stacked with a contrast index level.

average_treatment_effect ¶

average_treatment_effect(forest: 'CausalForest', X: Optional[ndarray] = None, T: Optional[ndarray] = None, target_sample: str = 'all', alpha: float = 0.05, clip: float = 0.01, variance: str = 'forest', covariates: Any = 'none', subset: Any = None) -> Dict[str, Any]

Aggregate CATE predictions into ATE/ATT/ATC/ATO targets.

This mirrors the most-used grf::average_treatment_effect targets: "all" (ATE), "treated" (ATT), "control" (ATC), and "overlap" (ATO, weighted by e(X)(1-e(X))).

The estimate is the doubly-robust AIPW influence-function mean (the estimator grf reports), not a plug-in average of the CATE predictions. Using the forest's own cross-fitted nuisances :math:\hat m(X)=\hat E[Y\mid X] and :math:\hat e(X)=\hat E[T\mid X], the ATE score is

.. math:: \Gamma_i = \hat\tau(X_i) + \frac{T_i-\hat e(X_i)}{\hat e(X_i)(1-\hat e(X_i))} \bigl(Y_i-\hat m(X_i)-(T_i-\hat e(X_i))\hat\tau(X_i)\bigr),

(:func:aipw_scores). ATT / ATC follow grf: the plug-in CATE mean over the target arm plus a Hajek-normalised doubly-robust correction, with the two variance components added (:func:grf_att_atc). se for the ATE is the influence-function standard error :math:\mathrm{sd}(\Gamma)/\sqrt n. The overlap target (ATO) is grf's Robinson / R-learner regression of Y - Y_hat on W - W_hat with an intercept and an HC3 standard error (:func:grf_overlap_ate); it needs no propensity inversion and is reported for continuous treatments as well. Given the same forest outputs all four targets reproduce grf exactly (tests/reference_parity/test_ml_causal_R_parity.py).

For forests fitted with the GRF engine (the default) the CATE entering the score is the out-of-bag prediction, observation weights follow equalize_cluster_weights, standard errors are cluster-robust when the forest has clusters, and continuous treatments use grf's debiasing weights -- the definitions used by grf::average_treatment_effect. X / T may only restate the training data.

Forests with fixed effects (fe=) have no propensity. For them only target_sample='treated' is identified, and it is estimated by imputation [borusyak2024revisiting]: unit and period effects are fitted on the untreated cells, each treated cell's effect is imputed as Y - alpha_hat_i - gamma_hat_t, and the ATT is their mean. The standard error uses the exact linear weights of that estimator, clustered by the forest's clusters; variance chooses how treated residuals are centred ('forest': out-of-bag forest prediction, then cohort x event-time means; 'bjs': cohort x event-time means only, the conservative Borusyak-Jaravel-Spiess convention, identical to :func:statspai.did_imputation). The payload also reports forest_plug_in, the mean OOB prediction over the same cells, which is shrunk toward zero when effects are heterogeneous.

.. versionchanged:: 1.30.0 Forests with fixed effects return the imputation ATT for target_sample='treated' instead of raising.

For other forests the doubly-robust scores exist only for the training sample; for other rows the function falls back to the plug-in CATE average, sets method='plug_in' and warns.

.. versionchanged:: 1.30.0 target_sample='overlap' is grf's R-learner estimator for every forest. Earlier releases averaged the AIPW scores with e(1-e) weights for legacy-engine forests, a different (also consistent) estimator of the same estimand.

Parameters:

Name	Type	Description	Default
`forest`	`fitted CausalForest`		required
`X`	`arrays`	Rows to aggregate over (default: the training sample).	`None`
`T`	`arrays`	Rows to aggregate over (default: the training sample).	`None`
`target_sample`	`('all', 'treated', 'control', 'overlap')`		`'all'`
`alpha`	`float`		`0.05`
`clip`	`float`	Propensity scores are clipped to `[clip, 1-clip]` before the inverse-propensity terms of the ATE / ATT / ATC scores. `grf` does not clip; `clip=0` reproduces it, and the clip is inert whenever all propensities already lie inside the band.	`0.01`
`variance`	`('forest', 'bjs')`	Forests with fixed effects only (see above); ignored otherwise.	`'forest'`
`covariates`	`'none', 'auto', list of str or array`	Forests with fixed effects only: covariates in the untreated outcome model `Y(0) = alpha_i + gamma_t + C' beta`. `'none'` is :func:`statspai.did_imputation` without covariates (same estimate; same standard error with `variance='bjs'`). `'auto'` adds the effect modifiers and covariates that vary within units; use it when a time-varying covariate drives the untreated outcome and is not itself affected by the treatment.	`'none'`
`subset`	`boolean mask or row positions`	Average over some of the training rows only, as `grf::average_treatment_effect(subset=)` does. The forest, its out-of-bag predictions and its nuisances stay those of the full fit; the target is the effect among the selected rows. Choose the rows from covariates. A subset built from the forest's own predictions on these rows (`oob > median`) describes the fit but its standard error does not account for the selection; see :func:`statspai.rate_split` for a test of that kind. GRF-engine forests without fixed effects only; for fixed-effects forests use :func:`statspai.forest_group_effects`.	`None`

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 400
>>> X = rng.normal(size=(n, 3))
>>> T = rng.binomial(1, 0.5, size=n)
>>> tau = 1.0 + X[:, 0]  # heterogeneous effect
>>> Y = X[:, 1] + tau * T + rng.normal(scale=0.5, size=n)
>>> df = pd.DataFrame({
...     "y": Y, "d": T,
...     "x0": X[:, 0], "x1": X[:, 1], "x2": X[:, 2],
... })
>>> cf = sp.causal_forest(
...     data=df, formula="y ~ d | x0 + x1 + x2",
...     n_estimators=50, random_state=0,
... )
>>> ate = sp.average_treatment_effect(cf, target_sample="all")
>>> ate["estimand"]
'ATE'
>>> att = sp.average_treatment_effect(cf, target_sample="treated")
>>> att["estimand"]
'ATT'

calibration_test ¶

calibration_test(forest: 'CausalForest', X: Optional[ndarray] = None, Y: Optional[ndarray] = None, T: Optional[ndarray] = None, alpha: float = 0.05, vce: str = 'HC3', method: str = 'auto', covariates: Any = 'none') -> DataFrame

Best-linear-predictor calibration test of CATEs [@chernozhukov2025generic].

The form of grf::test_calibration: the outcome residual Y - Y_hat is regressed (no intercept) on the treatment residual scaled by the mean forest prediction and on the treatment residual times the demeaned prediction,

Y - Y_hat = b1 (W - W_hat) mean(tau_hat)
            + b2 (W - W_hat) (tau_hat - mean(tau_hat)) + e.

b1 = 1 indicates a calibrated mean prediction; b2 = 1 a calibrated heterogeneity signal, and a significantly positive b2 is evidence of heterogeneity -- the headline finding, showing the forest captures real heterogeneity rather than noise. Following grf, t tests each coefficient against 0 and p is one-sided (alternative > 0) from Student's t with n - 2 degrees of freedom; the default vce='HC3' is grf's. beta_differential is also the slope by which the predictions should be rescaled around their mean (see :func:calibrate_cate).

For forests fitted with the GRF engine (the default) tau_hat is the out-of-bag prediction, the regression uses the forest's observation weights and a cluster-robust covariance when the forest has clusters, and the test is defined on the training sample, so X / Y / T may be omitted (the stored arrays are used) and, if given, must equal them. The regression itself is the pure operator :func:grf_calibration.

Forests with two-way fixed effects (fe='twoway') and a binary treatment use method='imputation' by default: the imputation score Y - alpha_hat_i - gamma_hat_t of each treated cell (unit and period effects fitted on untreated cells, [borusyak2024revisiting]) is regressed on mean(tau_oob) and tau_oob - mean(tau_oob) over the treated cells. The score is unbiased for each cell's effect, so beta_differential is the best-linear-predictor slope of the true effect on the forest prediction (it corrects the forest's shrinkage) and the one-sided test of it is a heterogeneity test. Standard errors use the exact linear weights of the imputation estimator, clustered by the forest's clusters, with a normal reference; vce does not apply. method='within' keeps the earlier regression on globally within-transformed variables [aytug2026attenuated], which tests for heterogeneity but whose slope is not a de-attenuation factor.

.. versionchanged:: 1.30.0 method= / covariates= added; forests with two-way fixed effects and a binary treatment default to the imputation regression. t / p are grf's: each coefficient against 0, one-sided, Student t (the null column is gone; t_vs_zero / p_one_sided are kept as synonyms). Legacy-engine forests use grf's regression too: earlier releases regressed an inverse-propensity pseudo-outcome on [tau_hat, tau_hat - mean(tau_hat)], so differential_forest_prediction estimated b2 - b1 rather than b2, and paired rows other than the training sample with mean-of-Y / mean-of-T stand-in nuisances.

Parameters:

Name	Type	Description	Default
`forest`	`fitted CausalForest`		required
`X`	`optional arrays`	The training sample (defaults to the stored arrays).	`None`
`Y`	`optional arrays`	The training sample (defaults to the stored arrays).	`None`
`T`	`optional arrays`	The training sample (defaults to the stored arrays).	`None`
`alpha`	`float`	Level for `ci_low` / `ci_high`.	`0.05`
`vce`	`('HC3', 'HC2', 'HC1', 'HC0')`	Heteroskedasticity-robust covariance (`grf`'s `vcov.type`).	`'HC3'`
`method`	`('auto', 'imputation', 'within')`	Forests with fixed effects only: `'auto'` is `'imputation'` for `fe='twoway'` with a binary treatment and `'within'` otherwise. Other forests accept only `'auto'`.	`'auto'`
`covariates`	`'none', 'auto', list of str or array`	Covariates in the untreated outcome model of the imputation regression (see :func:`average_treatment_effect`).	`'none'`

Returns:

Type	Description
`DataFrame`	Rows `mean_forest_prediction` and `differential_forest_prediction` with `coef`, `se`, `t`, `p`, `ci_low`, `ci_high`, `t_vs_zero`, `p_one_sided`.

Examples:

sp.test_calibration is an alias of this function.

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 400
>>> X = rng.normal(size=(n, 3))
>>> T = rng.binomial(1, 0.5, size=n)
>>> tau = 1.0 + X[:, 0]  # heterogeneous effect
>>> Y = X[:, 1] + tau * T + rng.normal(scale=0.5, size=n)
>>> df = pd.DataFrame({
...     "y": Y, "d": T,
...     "x0": X[:, 0], "x1": X[:, 1], "x2": X[:, 2],
... })
>>> cf = sp.causal_forest(
...     data=df, formula="y ~ d | x0 + x1 + x2",
...     n_estimators=50, random_state=0,
... )
>>> ct = sp.calibration_test(cf)
>>> ct.index.tolist()
['mean_forest_prediction', 'differential_forest_prediction']
>>> sp.test_calibration is sp.calibration_test
True

References

[@chernozhukov2025generic], [@athey2019generalized]

forest_diagnostics ¶

forest_diagnostics(forest: 'CausalForest', X: Optional[ndarray] = None, T: Optional[ndarray] = None, propensity_bounds: Tuple[float, float] = (0.05, 0.95)) -> Dict[str, object]

Return overlap and CATE-distribution diagnostics for a fitted forest.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 400
>>> X = rng.normal(size=(n, 3))
>>> T = rng.binomial(1, 0.5, size=n)
>>> tau = 1.0 + X[:, 0]  # heterogeneous effect
>>> Y = X[:, 1] + tau * T + rng.normal(scale=0.5, size=n)
>>> df = pd.DataFrame({
...     "y": Y, "d": T,
...     "x0": X[:, 0], "x1": X[:, 1], "x2": X[:, 2],
... })
>>> cf = sp.causal_forest(
...     data=df, formula="y ~ d | x0 + x1 + x2",
...     n_estimators=50, random_state=0,
... )
>>> diag = sp.forest_diagnostics(cf)
>>> bool(diag["n_treated"] + diag["n_control"] == diag["n"])
True

honest_variance ¶

honest_variance(forest: 'CausalForest', X: Optional[ndarray] = None, n_splits: int = 25, seed: Optional[int] = None) -> Dict[str, float]

Half-sample spread of the mean CATE prediction (deprecated).

Draws n_splits random halves (without replacement) of the CATE predictions and reports the standard deviation of the half-sample means as se. For a half-sample drawn without replacement the finite-population correction is exactly one half, so this spread estimates the sampling standard error of the full-sample mean, sd(tau_hat) / sqrt(n).

It is descriptive: the predictions are held fixed, so the forest's own estimation error is not propagated. grf has no counterpart to this function. For GRF-engine forests with X omitted it returns the doubly-robust ATE of :func:average_treatment_effect and its influence-function SE instead; use that function directly.

.. versionchanged:: 1.30.0 se was the standard deviation of the half-sample means divided by sqrt(n_splits) -- the Monte Carlo error of their average, which shrinks to zero as n_splits grows. It is now the spread itself.

Parameters:

Name	Type	Description	Default
`forest`	`fitted CausalForest`		required
`X`	`ndarray`		`None`
`n_splits`	`int`	Number of random half-sample draws.	`25`
`seed`	`int`		`None`

Returns:

Type	Description
dict with ``ate``, ``se``, ``ci_low``, ``ci_high`` (95 %).

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 400
>>> X = rng.normal(size=(n, 3))
>>> T = rng.binomial(1, 0.5, size=n)
>>> tau = 1.0 + X[:, 0]  # heterogeneous effect
>>> Y = X[:, 1] + tau * T + rng.normal(scale=0.5, size=n)
>>> df = pd.DataFrame({
...     "y": Y, "d": T,
...     "x0": X[:, 0], "x1": X[:, 1], "x2": X[:, 2],
... })
>>> cf = sp.causal_forest(
...     data=df, formula="y ~ d | x0 + x1 + x2",
...     n_estimators=50, random_state=0,
... )
>>> hv = sp.honest_variance(cf, n_splits=25, seed=0)
>>> bool(hv["se"] >= 0 and hv["ci_low"] <= hv["ate"] <= hv["ci_high"])
True

rate ¶

rate(forest: 'CausalForest', X: Optional[ndarray] = None, Y: Optional[ndarray] = None, T: Optional[ndarray] = None, target: str = 'AUTOC', q_grid: int = 100, alpha: float = 0.05, seed: Optional[int] = None, priorities: Optional[ndarray] = None, se_method: str = 'auto', n_bootstrap: int = 200, variance: str = 'bjs', cluster: Any = None, members: Any = None, covariates: Any = 'none') -> Dict[str, Any]

Rank-weighted average treatment effect (RATE) [@yadlowsky2025evaluating].

``grf::rank_average_treatment_effect``: evaluates a prioritisation
rule -- by default the forest's own CATE predictions (out-of-bag for
GRF-engine forests), or any ``priorities`` vector (e.g. from a model
trained on another sample) -- on the forest's doubly-robust scores
(:func:`aipw_scores`, ``grf::get_scores``).  Units are sorted by
decreasing priority, scores are averaged within tied priorities, and

- **AUTOC** = mean over ``k`` of ``TOC_k``,
- **QINI** = mean over ``k`` of ``(k/n) TOC_k``,

with ``TOC_k`` the mean score of the top ``k`` units minus the overall
mean -- ``grf``'s estimator, reproduced exactly given the same scores
and priorities (the pure operator :func:`rate_from_scores`).

Standard error: ``se_method='influence'`` (default) is the analytic
influence-function SE including the term from estimating the priority
ranks, summed within clusters when the forest has clusters;
``se_method='half_sample'`` is ``grf``'s half-sample bootstrap with
``n_bootstrap`` draws (``R`` in ``grf``, default 200; whole clusters
are drawn), seeded by ``seed``.  The two agree to Monte Carlo error.

Because priorities and evaluation share one sample when
``priorities`` is omitted, RATE is then a *diagnostic*; for a
pre-registered test of targeting value, fit the forest on a training
split and evaluate priorities on a held-out split.

**Forests with fixed effects** have no propensity, so the curve is read
on the population they identify -- the treated cells -- through the
imputation scores :math:`\Gamma_{it} = Y_{it} - \hat\alpha_i -
\hat\gamma_t` that :func:`statspai.average_treatment_effect` already
uses there:  ``TOC(q) = ATT(top q by priority) - ATT``, a
*retrospective* targeting curve ("the effect was this much larger among
the cells the rule would have prioritised"), not the population RATE a
randomised design gives.  ``variance=``, ``cluster=``, ``members=`` and
``covariates=`` are then passed through exactly as in
:func:`statspai.forest_group_effects`, and ``se_method`` becomes:

- ``'imputation'`` (what ``'auto'`` selects): the rank weights and the
  imputation weights are both linear, so their composition makes RATE
  one more linear functional of ``y`` with the same exact, cluster- or
  dyad-robust variance as the ATT.  The weights annihilate the fixed
  effects exactly (``max |Z'V| < 1e-15`` in the committed fixture) and
  sum to zero over treated cells, RATE being a contrast.
- ``'influence'``: the rank-corrected influence function applied to the
  imputation scores -- it pays for estimating the ranking but treats
  the fixed effects as known.

``variance`` defaults to ``'bjs'`` here, not to the ``'forest'`` that
:func:`statspai.average_treatment_effect` uses, and the reason is
measured.  ``'forest'`` centres treated residuals on the out-of-bag
prediction, which nets out the effect heterogeneity and so estimates
the variance of the RATE *of this sample*; the ATT converges to its
population value fast enough that this does not show, but RATE loads on
the tail of the effect distribution and it does.  Over 120 replications
(N = 150 units, T = 8, staggered adoption selected on the unit effect,
tau = 0.3 + 0.5 z, priorities held out) nominal 95% intervals covered
the *population* AUTOC 97.5% of the time with ``'bjs'`` and 90.8% with
``'forest'``; against the RATE of the realised sample it was the other
way round, 99.2% and 95.8%.  So ``'bjs'`` is conservative for the
population statement most readers want, ``'forest'`` is exact for the
sample one.  Bias was -0.006 on 0.452 (AUTOC) and -0.0002 on 0.141
(QINI).

``se_method='half_sample'`` raises for these forests: resampling units
would have to refit the untreated two-way model in every draw, and the
exact variance supersedes it.

**Do not read the default ranking as a test.** With the forest's own
out-of-bag predictions the ranking and the scores share
``gamma_hat_t``, so they stay correlated even though no unit predicts
itself.  On the null design above (400 replications, no heterogeneity
at all) AUTOC averaged -0.028 rather than 0 and a nominal 5% test
rejected 16.0% of the time, QINI 12.2%.  :func:`statspai.rate_split`,
which refits the rule on one half of the units and scores it on the
other, averaged -0.0002 and rejected 0.5% and 0.0%, with 100% power
against ``tau = 0.3 + 0.5 z``
(``tests/reliability/forest_split_evaluation.py``).  A warning fires on
this path;
passing ``priorities=`` from a rule fitted elsewhere also silences it.

.. versionchanged:: 1.31.0
   Forests with fixed effects are supported (they previously raised);
   the default ``se_method`` is ``'auto'``, which is ``'influence'`` for
   every forest that has a propensity, as before.

.. versionchanged:: 1.30.0
   Scores are ``grf``'s (unclipped) AIPW scores; tied priorities are
   averaged as in ``grf``; QINI uses ``grf``'s ``k/n`` weights, so a
   flat TOC now gives exactly 0 (previously ``-mean(score) / (2n)``);
   the TOC curve interpolates between cut points as ``grf`` does; the
   influence-function SE includes the rank-estimation term
   (previously ~30% too large).

Parameters

forest : fitted CausalForest
X, Y, T : arrays, optional
    The training sample (defaults to the stored arrays).
target : {'AUTOC', 'QINI'}
q_grid : int
    Number of equally spaced points in ``(0, 1]`` at which the TOC
    curve is reported. Does not affect the estimate.
alpha : float
seed : int, optional
    Seed for ``se_method='half_sample'``.
priorities : array, optional
    Priority score per training unit (higher = treat first). Defaults
    to the forest's CATE predictions.
se_method : {'auto', 'influence', 'half_sample', 'imputation'}
    ``'auto'`` is ``'influence'``, or ``'imputation'`` for a forest with
    fixed effects. ``'imputation'`` is defined only for those.
n_bootstrap : int
    Half-sample draws for ``se_method='half_sample'``.
variance : {'bjs', 'forest'}
    Fixed-effects forests only: how treated residuals are centred
    before the cohort x event-time blocks. ``'bjs'`` (default here) is
    conservative for the population RATE, ``'forest'`` exact for the
    sample's -- see above.
cluster : array-like or 'dyadic', optional
    Fixed-effects forests only; ``'dyadic'`` needs ``members=``.
members : array-like, optional
    Fixed-effects forests only: the two members of each dyadic row.
covariates : {'none', 'auto'}, list of str or array, default 'none'
    Fixed-effects forests only: time-varying covariates for the
    untreated model behind the imputation scores.

Returns

dict with keys ``estimate``, ``se``, ``ci_low``, ``ci_high``,
``target``, ``toc_curve`` (``(q_grid, 2)``), ``n``, ``method``,
``priority_source``, ``n_clusters``.

References

[@yadlowsky2025evaluating]

Examples

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 400
>>> X = rng.normal(size=(n, 3))
>>> T = rng.binomial(1, 0.5, size=n)
>>> tau = 1.0 + X[:, 0]  # heterogeneous effect
>>> Y = X[:, 1] + tau * T + rng.normal(scale=0.5, size=n)
>>> df = pd.DataFrame({
...     "y": Y, "d": T,
...     "x0": X[:, 0], "x1": X[:, 1], "x2": X[:, 2],
... })
>>> cf = sp.causal_forest(
...     data=df, formula="y ~ d | x0 + x1 + x2",
...     n_estimators=50, random_state=0,
... )
>>> res = sp.rate(cf, target="AUTOC")
>>> sorted(res.keys())
['ci_high', 'ci_low', 'estimate', 'method', 'n', 'n_clusters', \

'priority_source', 'se', 'target', 'toc_curve'] >>> res["toc_curve"].shape (100, 2)

statspai.causal¶

causal ¶

CausalForest ¶

fit ¶

get_nuisances ¶

effect ¶

predict ¶

effect_interval ¶

oob_effect ¶

effect_variance ¶

split_frequencies ¶

average_treatment_effect ¶

group_effects ¶

forest_diagnostics ¶

summary ¶

variable_importance ¶

best_linear_projection ¶

ate ¶

att ¶

IVForestResult ¶

predict ¶

get_scores ¶

average_treatment_effect ¶

best_linear_projection ¶

MultiArmForestResult ¶

predict ¶

get_scores ¶

average_treatment_effect ¶

best_linear_projection ¶

average_treatment_effect ¶

calibration_test ¶

forest_diagnostics ¶

honest_variance ¶

rate ¶

`statspai.causal`¶