statspai.forest

forest

Forest-based causal inference estimators for StatsPAI.

Hosts CausalForest (grf-style honest causal forests) and its companions:

• :func:causal_forest / :class:CausalForest — heterogeneous treatment-effect estimation via honest random forests (Wager-Athey 2018; Athey-Tibshirani-Wager 2019).
• :func:iv_forest / :func:instrumental_forest — instrumental forests for conditional LATEs (Athey-Tibshirani-Wager 2019).
• :func:multi_arm_forest / :func:lm_forest — multi-arm causal forests and conditionally linear models (Nie-Wager 2021 with GRF weights).
• :func:regression_forest, :func:multi_regression_forest, :func:probability_forest, :func:quantile_forest, :func:survival_forest — the prediction forests of the GRF family.
• :func:variable_importance, :func:best_linear_projection, :func:get_scores — grf's post-estimation for any of them.
• :func:forest_group_effects / :func:forest_support / :func:cate_pretrend_test -- group effects with valid SEs (imputation scores for fe= forests), support of counterfactual predictions, and pre-trends by predicted-effect group.
• :func:calibration_test / :func:test_calibration / :func:grf_calibration (the same test on predictions made elsewhere) / :func:rate / :func:honest_variance — post-fit honesty & calibration diagnostics.

This package was previously named statspai.causal — the old name is kept as a deprecation shim for one minor version cycle. Use from statspai.forest import ... going forward.

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).

LMForestResult

Bases: GRFFamilyForest

A fitted linear-model forest.

Attributes:

Name	Type	Description
coefficients	ndarray	Out-of-bag h_k(X_i), shape (n, K, q).
coefficient_variance	ndarray or None	Little-bag variances, (n, K, q) (single outcome only).
regressor_names, outcome_names	list of str

Examples:

>>> import numpy as np, statspai as sp
>>> rng = np.random.default_rng(0)
>>> X = rng.normal(size=(500, 2)); W = rng.normal(size=(500, 1))
>>> Y = (1 + X[:, 0]) * W[:, 0] + rng.normal(size=500)
>>> sp.lm_forest(y=Y, regressors=W, covariates=X,
...              n_estimators=100).regressor_names
['regressors']

predict

predict(newdata: Any = None, estimate_variance: bool = False) -> Dict[str, ndarray]

{"predictions": (m, K, q)[, "variance_estimates": (m, K, q)]}.

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.

PredictionForest

Bases: GRFFamilyForest

A fitted regression / multi-regression / probability / quantile forest.

Attributes:

Name	Type	Description
predictions	ndarray	Out-of-bag predictions for the training rows: (n,) for a regression forest, (n, q) for multi-regression, (n, K) class probabilities, (n, len(quantiles)) quantiles.
variance	ndarray or None	Out-of-bag little-bag variances (regression, multi-regression and probability forests with ci_group_size >= 2).
outcome_names	list of str	Outcome columns, class labels or quantile labels.

Examples:

>>> import numpy as np, statspai as sp
>>> rng = np.random.default_rng(0)
>>> X = rng.normal(size=(300, 2)); y = X[:, 0] + rng.normal(size=300)
>>> sp.regression_forest(y=y, covariates=X, n_estimators=100).forest_type
'regression_forest'

predict

predict(newdata: Any = None, estimate_variance: bool = False, quantiles: Optional[Sequence[float]] = None) -> DataFrame

Predictions at newdata (out-of-bag for the training rows when newdata is None).

SurvivalForestResult

Bases: GRFFamilyForest

A fitted survival forest.

Attributes:

Name	Type	Description
failure_times	ndarray	Grid on which curves are reported.
predictions	ndarray	Out-of-bag survival curves of the training rows, (n, n_times).
prediction_type	str	"Kaplan-Meier" or "Nelson-Aalen".

Examples:

>>> import numpy as np, statspai as sp
>>> rng = np.random.default_rng(0)
>>> X = rng.normal(size=(300, 2)); T = rng.exponential(size=300)
>>> sf = sp.survival_forest(time=T, event=np.ones(300), covariates=X,
...                         n_estimators=100)
>>> sf.prediction_type
'Kaplan-Meier'

predict

predict(newdata: Any = None, failure_times: Any = None, prediction_type: Optional[str] = None) -> DataFrame

Survival curves at newdata (OOB for the training rows when None), evaluated on failure_times (default: the fit grid; other times take the curve's value at the last grid time at or below).

cate_pretrend_test

cate_pretrend_test(forest: Any, *, n_groups: int = 2, leads: Optional[int] = None, groups: Any = None, time_effects: str = 'common', alpha: float = 0.05, covariates: Any = 'none') -> Dict[str, Any]

Test whether predicted-effect groups had different pre-treatment trends.

Heterogeneity found by a forest can be spurious if the units it ranks high were already diverging before treatment [aytug2026euro]. This test sorts units into n_groups quantiles of their mean out-of-bag prediction (or uses groups) and, on the untreated cells only, regresses the outcome on unit and period effects and on lead indicators 1[t - g = k] x 1[group] for k = -1, ..., -leads -- the pre-trend regression of [borusyak2024revisiting], interacted with the groups. Periods earlier than -leads and never-treated units are the reference. Using untreated cells only keeps post-treatment effect heterogeneity out of the test.

Two Wald tests are reported: all leads are zero, and the leads are equal across groups (for each k, every group's lead equals the first group's). The second is the test that bears on heterogeneity. Standard errors are clustered by the forest's clusters with the CR1 factor G/(G-1) (n-1)/(n-K), where K counts the leads and the period effects (unit effects are nested in unit clusters, fixest's fixef.K = "nested").

Parameters:

Name	Type	Description	Default
forest	CausalForest	Fitted with fe="twoway" and a binary absorbing treatment.	required
n_groups	int	Quantile groups of the unit-level mean OOB prediction.	2
leads	int	Number of pre-treatment leads (default: min(4, L - 1) where L is the longest pre-treatment spell; with L leads every pre-period of the earliest-observed units is a lead and the leads are collinear with their unit effects).	None
groups	array - like	One label per training row, constant within units; overrides n_groups.	None
time_effects	('common', 'by_group')	"common" matches the forest and the imputation scores, which use one set of period effects; "by_group" gives each group its own period effects (separate event studies per group).	"common"
alpha	float		0.05
covariates	"none", "auto", list of str or array	Time-varying covariates added linearly, as in the imputation model of :func:forest_group_effects.	"none"

Returns:

Type	Description
dict	coefficients (DataFrame indexed by group and lead), joint_zero and equal_across_groups (stat, df, p), group_of_unit (Series), n_obs, n_clusters and method.

Examples:

>>> import statspai as sp
>>> df = sp.datasets.currency_union_panel(seed=0)
>>> cf = sp.causal_forest(
...     data=df, y="log_trade", d="euro",
...     x=["pre_trade", "log_gdp_prod", "log_gdppc"],
...     id="pair", time="year", fe="twoway",
...     n_estimators=400, random_state=0,
... )
>>> res = sp.cate_pretrend_test(cf, n_groups=2, leads=3)
>>> sorted(res["equal_across_groups"])
['df', 'p', 'stat']

References

[borusyak2024revisiting], [aytug2026euro]

forest_group_effects

forest_group_effects(forest: Any, 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 treatment effects by group, with valid standard errors.

A causal forest's predictions are regularised: averaging them over a group gives a number that is shrunk toward the overall mean and whose spread says nothing about sampling error. This function averages an unbiased per-row signal instead and uses the forest only to form groups when asked to:

• Forests with two-way fixed effects (fe="twoway"): the imputation score Y - alpha_hat_i - gamma_hat_t of each treated cell, with unit and period effects fitted on the untreated cells [borusyak2024revisiting]. Each group estimate is the ATT over the group's treated cells. Standard errors use the exact linear weights of the imputation estimator; treated residuals are centred on the out-of-bag forest prediction and then on cohort x event-time means (variance="forest") or on those means only (variance="bjs", the conservative convention of :func:statspai.did_imputation).
• Pooled forests: the AIPW score of every row; each estimate is the group ATE.

Rows are weighted by the forest's observation weights, so with equalize_cluster_weights=True every cluster gets equal weight in each group average (a different estimand from :func:statspai.did_imputation, which weights treated cells equally).

Parameters:

Name	Type	Description	Default
forest	CausalForest	Fitted with the GRF engine.	required
by	(None, 'cate_quantile' or array - like)	None pools all eligible rows (or, with members, forms one group per member). "cate_quantile" bins eligible rows into n_groups quantiles of the out-of-bag forest prediction (sorted group average treatment effects, [chernozhukov2025generic]). An array gives one label per training row (a pair id, a period band, df["year"] >= 2009 ...).	None
members	(n, 2) array-like	The two members of each row of dyadic data (exporter and importer, the two countries of a pair). With by=None each member is a group and a row counts for both of its members -- the average effect over all pairs a country belongs to.	None
n_groups	int	Number of quantile groups for by="cate_quantile".	4
cluster	(None, 'dyadic' or array - like)	None clusters by the forest's clusters (default: the panel unit). "dyadic" (requires members) allows any two rows that share a member to be correlated [aronow2015cluster]. An array gives other cluster ids.	None
variance	('forest', 'bjs')	Forests with fixed effects only; see above.	"forest"
alpha	float		0.05
scale	('level', 'percent')	"percent" adds 100 * (exp(x) - 1) columns for log outcomes (interval endpoints are transformed, not re-derived).	"level"
min_rows	int	Drop groups with fewer eligible rows.	1
covariates	"none", "auto", list of str or array	Forests with fixed effects only: covariates entering the untreated outcome model linearly, Y(0) = alpha_i + gamma_t + C' beta. "none" is the pure two-way model of :func:statspai.did_imputation; "auto" adds the forest's effect modifiers and covariates that vary within units (time-invariant ones are absorbed by the unit effect). Controls must not be affected by the treatment.	"none"

Returns:

Type	Description
DataFrame	One row per group: n_rows, n_units, estimate, se, z, p, ci_low, ci_high and forest_mean (the mean out-of-bag prediction over the same rows, for comparison). attrs holds method, estimand, the full vcov and tests: a Wald test that all group effects are equal and the last-minus-first difference with its standard error.

Examples:

>>> import statspai as sp
>>> df = sp.datasets.currency_union_panel(seed=0)
>>> cf = sp.causal_forest(
...     data=df, y="log_trade", d="euro",
...     x=["pre_trade", "log_gdp_prod", "log_gdppc"],
...     id="pair", time="year", fe="twoway",
...     n_estimators=400, random_state=0,
... )
>>> by_country = sp.forest_group_effects(
...     cf, members=df[["country_i", "country_j"]].to_numpy(),
...     scale="percent",
... )
>>> by_country.index.name
'group'

References

[borusyak2024revisiting], [chernozhukov2025generic], [aronow2015cluster], [kattenberg2023causal]

forest_policy_tree

forest_policy_tree(forest: Any, *, depth: int = 2, cost: float = 0.0, x: Any = None, min_leaf_size: Optional[int] = None, n_splits: int = 21, train_frac: float = 0.5, random_state: int = 0, members: Any = None, alpha: float = 0.05, variance: str = 'bjs', cluster: Any = None, covariates: Any = 'none') -> Dict[str, Any]

A treatment rule for a fixed-effects forest, and what it was worth.

Policy learning [athey2021policy] maximises :math:V(\pi) = E[\Gamma_i \pi(X_i)] over a class of rules, with :math:\Gamma a score unbiased for the individual effect. A within-unit design has no propensity, so the doubly-robust score :func:statspai.policy_tree uses does not exist; the imputation scores that give the ATT do, and the rule is learned on those.

Read the population it applies to. Those scores live on treated cells, so this is a retrospective rule: "of the cells that were treated, which should have been". It does not say whether an untreated unit should be treated -- that is an extrapolation of tau(x), which :func:statspai.average_treatment_effect also refuses for these forests. Check :func:statspai.forest_support before carrying the rule to units the design never switched.

With cost=0 and effects that are positive everywhere, treating every cell is optimal and the tree has nothing to find. cost is what usually makes targeting a real question: the rule then treats where tau > cost.

The rule is fitted and priced on disjoint units. A tree chosen to maximise the value of the very scores it is then priced against is priced on its own noise. Measured on a design where every rule is worth exactly the same -- no heterogeneity, and cost equal to the constant effect, so the true gain over treating everyone is exactly 0 (200 replications, N = 200 units, T = 8) -- the same-sample gain averaged +0.048, larger than its own standard error of 0.037, and 13.0% of runs reported a significant benefit from targeting where none existed. Fitted and priced on disjoint halves, with the default 21 splits (400 replications): +0.008 and 0.0%. The split costs little when the heterogeneity is real: with tau = 0.3 + 0.8 z and cost = 0.3 the oracle gain is 0.8 * phi(0) = 0.3192, and the split-sample estimate averaged 0.307 and was significant in every replication; the same-sample estimate overshoots the oracle (0.3316 in the earlier study). The split is therefore not optional here; train_frac moves it, nothing switches it off.

Parameters:

Name	Type	Description	Default
forest	fitted CausalForest with ``fe="twoway"``	Supplies the data and the hyper-parameters; it is not itself used to fit or to price the rule, and is left untouched.	required
depth	int	Tree depth. <= 2 is searched exactly (the policytree guarantee), deeper is greedy and says so in method.	2
cost	float	Cost of treating one cell, in the outcome's units. The rule treats where the effect exceeds it.	0
x	list of str or array	Policy covariates. Defaults to the forest's effect modifiers. Keep this list short and interpretable -- it is the rule a programme would actually be written in.	None
min_leaf_size	int	Smallest leaf; default max(10, 5% of the evaluation cells).	None
train_frac	float	Share of units (or, with members=, dyadic nodes) that fit the rule. The rest price it.	0.5
random_state	int		0
members	array - like	The two members of each dyadic row; splitting is then by member.	None
alpha	float	As :func:statspai.forest_group_effects. variance defaults to 'bjs' here for the reason it does in :func:statspai.rate.	0.05
variance	float	As :func:statspai.forest_group_effects. variance defaults to 'bjs' here for the reason it does in :func:statspai.rate.	0.05
cluster	float	As :func:statspai.forest_group_effects. variance defaults to 'bjs' here for the reason it does in :func:statspai.rate.	0.05
covariates	float	As :func:statspai.forest_group_effects. variance defaults to 'bjs' here for the reason it does in :func:statspai.rate.	0.05

Returns:

Type	Description
dict	rules (the tree, printable), tree (nested dict), policy (0/1 per evaluation cell), value and value_treat_all and gain_over_treat_all (each with estimate, se, ci_low, ci_high, p), share_treated, n_train_units / n_eval_units, split_by, method, diagnostics.

See Also

statspai.policy_tree : the doubly-robust version, for designs with a propensity. statspai.rate_split : is any ranking worth targeting on? statspai.forest_support : is a rule safe to carry to untreated units?

References

[@athey2021policy], [@borusyak2024revisiting], [@kattenberg2023causal]

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> rows = []
>>> for i in range(120):
...     a, z = rng.normal(), rng.normal()
...     g = 4 if i % 3 else 10**6
...     for t in range(1, 8):
...         d = 1.0 * (t >= g)
...         y = a + 0.2 * t + (0.3 + 0.8 * z) * d + rng.normal(0, 0.5)
...         rows.append((i, t, y, d, z))
>>> df = pd.DataFrame(rows, columns=["id", "t", "y", "d", "z"])
>>> cf = sp.causal_forest(
...     "y ~ d | z", data=df, fe="twoway", unit="id", time="t",
...     clusters=df["id"].to_numpy(), n_estimators=100, random_state=0,
... )
>>> res = sp.forest_policy_tree(cf, depth=1, cost=0.3, n_splits=3)
>>> 0.0 <= res["share_treated"] <= 1.0
True
>>> gain = res["gain_over_treat_all"]
>>> sorted(k for k in gain if k in {"estimate", "se", "ci_low", "ci_high"})
['ci_high', 'ci_low', 'estimate', 'se']
>>> gain["n_splits"], res["n_splits"]
(3, 3)
>>> sorted(res["diagnostics"]["split_stability"])[:2]
['root_covariate_counts', 'root_covariate_modal_share']

Many splits, not one. One split is a draw: it fixes both the rule and its price, and with random_state in reach it is easy to keep the pair that agrees with you -- the practice [chernozhukov2025generic] show invalidates inference. n_splits aggregates by their variational estimation and inference (VEIN): the value, the treat-all value and the gain are each the median over splits, their intervals are medians of conditional intervals built at 1 - alpha / 2 (whose median covers at 1 - alpha), and the p-values are twice the median conditional p-value. Each is aggregated on its own, as [chernozhukov2025generic] do, so with n_splits > 1 the three medians do not satisfy gain = value - value_treat_all -- a median of differences is not a difference of medians. That identity holds within each split, and n_splits=1 reports it directly. A rule cannot be averaged either, so the one reported is the rule from the split whose gain is the median, and diagnostics['split_stability'] says how often each covariate was chosen at the root and how far the thresholds and treated shares moved -- if the root covariate changes from split to split, the rule is not identified by this sample however tight the interval on its value looks. n_splits=21 costs about twelve seconds on a 150-unit panel; n_splits=1 reproduces one conditional split and warns.

Parameters:

Name	Type	Description	Default
n_splits	int	Splits to aggregate by VEIN. Odd, so the median is an order statistic.	21
random_state	int	Seeds the sequence of splits; split b uses random_state + b.	0

forest_support

forest_support(forest: Any, X_new: Any, *, k: int = 10, quantile: float = 0.95, alpha: float = 0.05) -> DataFrame

Check whether new rows lie inside the support that informs tau(x).

Predicting the effect a unit would have had (a country that did not adopt a currency, a region that was never treated) applies the fitted tau(x) outside the treated sample. That is a model-based extrapolation, credible only where the forest has seen similar units switch treatment. For every row of X_new this reports

• cate, its little-bag se and a (1 - alpha) interval;
• n_outside_range -- effect modifiers outside the [min, max] of the reference rows;
• knn_ratio -- the distance (on reference-standardised modifiers) to the k-th nearest reference row, divided by the quantile of the reference rows' distance to their k-th nearest row from a different unit (rows of the same unit are excluded, so a unit's other periods do not shrink the benchmark). Above 1 the row is farther from the data than all but 1 - quantile of the reference rows are from other units;
• supported -- n_outside_range == 0 and knn_ratio <= 1.

The reference rows are the rows that inform the effect: for a forest with fixed effects, the rows of units whose treatment varies (only they identify the within effect); for a pooled forest, all training rows.

Parameters:

Name	Type	Description	Default
forest	CausalForest	Fitted with the GRF engine.	required
X_new	array - like or DataFrame	Effect modifiers of the rows to predict (DataFrame columns are matched to the fitted feature names).	required
k	int		10
quantile	float		0.95
alpha	float		0.05

Returns:

Type	Description
DataFrame	One row per row of X_new; attrs["summary"] holds the share of supported rows and the reference-sample size.

Examples:

>>> import statspai as sp
>>> df = sp.datasets.currency_union_panel(seed=0)
>>> x = ["pre_trade", "log_gdp_prod", "log_gdppc"]
>>> cf = sp.causal_forest(
...     data=df, y="log_trade", d="euro", x=x,
...     id="pair", time="year", fe="twoway",
...     n_estimators=400, random_state=0,
... )
>>> outs = df[df["ever_euro"] == 0]
>>> sup = sp.forest_support(cf, outs[x])
>>> bool(sup["supported"].mean() > 0)
True

rate_split

rate_split(forest: Any, target: str = 'AUTOC', *, n_splits: int = 21, train_frac: float = 0.5, random_state: int = 0, members: Any = None, alpha: float = 0.05, variance: str = 'bjs', cluster: Any = None, covariates: Any = 'none', se_method: str = 'auto', q_grid: int = 100) -> Dict[str, Any]

RATE of the forest's targeting rule, fitted and evaluated on disjoint units.

:func:statspai.rate ranks the rows it also scores. For a forest with fixed effects that is not a valid test even though the predictions are out-of-bag: every imputation score carries :math:-\hat\gamma_t, estimated from the same periods the forest was trained on, so the ranking and the scores stay correlated. Measured on a design with no heterogeneity whatsoever (400 replications, N = 150 units, T = 8, staggered adoption selected on the unit effect, tau = 0.3): AUTOC averaged -0.028 instead of 0 and a nominal 5% test rejected 16.0% of the time; QINI averaged -0.006 and rejected 12.2%.

This function removes the overlap. The units (or, with members=, the nodes of a dyadic panel) are split in two; a forest with the same hyper-parameters is refitted on each half; the training half's forest ranks the evaluation half's treated cells, whose imputation scores come from an untreated two-way model fitted on the evaluation half alone. Nothing the rule saw enters the score it is graded on. On the same null design, with the default 21 splits, AUTOC then averaged -0.0002 and rejected 0.5% of the time (QINI +0.0004 and 0.0%): the aggregation over splits is conservative by construction, so the test errs on the side of not rejecting. Against tau = 0.3 + 0.5 z it rejected in every replication. There the estimate averaged 0.383 for AUTOC where the ideal ranking has 0.452, because a forest grown on half the units ranks imperfectly: the quantity estimated is the RATE of the rule that was fitted, not of the best rule. variance='forest' gives the same rejection rates under the null but narrower intervals under the alternative (they covered the ideal ranking's AUTOC 89.5% of the time against 94.5%), which is why 'bjs' is the default here. The study is tests/reliability/forest_split_evaluation.py; single-split figures quoted through 1.39.3 (7.5% and 4.0%) came from halves that were not refitted the way the forest was.

The price is sample: each forest sees half the units, so the rule is noisier than the one fitted on everything and the RATE it earns is a lower bound on what the full-sample rule is worth. Report :func:statspai.rate for the point estimate of the full rule if you like, but report this one when the claim is that targeting pays.

Parameters:

Name	Type	Description	Default
forest	fitted CausalForest	Supplies the data and the hyper-parameters; it is not itself used to rank or to score, and is left untouched.	required
target	('AUTOC', 'QINI')		'AUTOC'
train_frac	float	Share of units (members) that fit the rule. The rest evaluate it.	0.5
random_state	int	Seed for the split. The refitted forests keep the original's seed.	0
members	array - like	The two members of each dyadic row. Splitting is then by member and rows that straddle the two halves are dropped (counted in the result), because a flow between a training and an evaluation country belongs to neither.	None
alpha	float	Passed to :func:statspai.rate on the evaluation half.	0.05
variance	float	Passed to :func:statspai.rate on the evaluation half.	0.05
cluster	float	Passed to :func:statspai.rate on the evaluation half.	0.05
covariates	float	Passed to :func:statspai.rate on the evaluation half.	0.05
se_method	float	Passed to :func:statspai.rate on the evaluation half.	0.05
q_grid	float	Passed to :func:statspai.rate on the evaluation half.	0.05

Returns:

Type	Description
dict	As :func:statspai.rate, plus n_train_units / n_eval_units, n_rows_dropped, split_by and split_random_state.

See Also

statspai.rate : the same curve without the split (a diagnostic). statspai.forest_group_effects : effects for named groups rather than a curve.

References

[@yadlowsky2025evaluating]

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> rows = []
>>> for i in range(80):
...     a, z = rng.normal(), rng.normal()
...     g = 4 if i % 3 else 10**6
...     for t in range(1, 7):
...         d = 1.0 * (t >= g)
...         y = a + 0.2 * t + (0.3 + 0.5 * z) * d + rng.normal(0, 0.5)
...         rows.append((i, t, y, d, z))
>>> df = pd.DataFrame(rows, columns=["id", "t", "y", "d", "z"])
>>> cf = sp.causal_forest(
...     "y ~ d | z", data=df, fe="twoway", unit="id", time="t",
...     clusters=df["id"].values, n_estimators=100, random_state=0,
... )
>>> res = sp.rate_split(cf, target="AUTOC")
>>> res["split_by"], res["n_eval_units"] > 0
('units', True)

Many splits, not one. A single split is a draw, not an estimator: on the guide's 15-country trade panel six of them ranged over +/-0.08, and with random_state in reach it is easy to keep the one that agrees with you -- the practice [chernozhukov2025generic] show invalidates inference. Their variational estimation and inference (VEIN) is what n_splits does: the estimate is the median over splits, the interval is the median of the conditional intervals built at 1 - alpha / 2 (whose median covers at 1 - alpha), and the p-value is twice the median conditional p-value. n_splits=21 costs about twelve seconds on a 150-unit panel; 100, as they use, costs about a minute. n_splits=1 reproduces a single conditional split and warns. estimate_min / estimate_max / estimate_iqr in the result show how much the split was moving the answer.

Parameters:

Name	Type	Description	Default
forest	Any		required
target	Any		required
train_frac	Any		required
members	Any		required
alpha	Any		required
variance	Any		required
cluster	Any		required
covariates	Any	As described above.	'none'
se_method	Any	As described above.	'none'
q_grid	Any	As described above.	'none'
n_splits	int	Splits to aggregate by VEIN. Odd, so the median is an order statistic.	21
random_state	int	Seeds the sequence of splits; split b uses random_state + b.	0

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'

calibrate_cate

calibrate_cate(forest: 'CausalForest', newdata: Optional[ndarray] = None, alpha: float = 0.05, method: str = 'auto', covariates: Any = 'none') -> Dict[str, Any]

Rescale CATE predictions by their best-linear-predictor calibration.

Forest CATE predictions are shrunk toward their mean: leaf averaging and honest subsampling regularise, so the spread of tau_hat understates the spread of tau. The calibration regression of :func:calibration_test estimates the slope beta_differential of the effect signal on the (out-of-bag) predictions; this function returns

tau_cal(x) = beta_mean * mean(tau_oob)
             + beta_differential * (tau_hat(x) - mean(tau_oob))

for the training rows (newdata=None, out-of-bag) or for new rows. Because the slope is estimated on out-of-bag predictions it tends to 0 when the effect is homogeneous, so the correction does not manufacture heterogeneity.

Forests with two-way fixed effects and a binary treatment (method='auto' or 'imputation'): the slope comes from regressing imputation scores on the out-of-bag predictions over the treated cells (see :func:calibration_test). The scores are unbiased for each cell's effect, so the slope is a genuine de-attenuation factor, and tau_bar is the mean prediction over treated cells, so beta_mean * tau_bar is the imputation ATT. In StatsPAI's simulations (N = 300 units, T = 8, staggered adoption selected on the unit effect, 200 replications) this reduced the RMSE of the treated cells' predictions from 0.634 to 0.570 with effects (1 + x1)(1 + 0.2 e) and from 0.213 to 0.118 with a constant effect. The map is estimated on treated cells; applying it to other rows assumes it carries over.

method='within' uses globally within-transformed variables, as proposed by [@aytug2026attenuated]. Global two-way demeaning is exact only when the effect is homogeneous, so the heterogeneity test keeps its size, but the slope is not a de-attenuation factor: in StatsPAI's simulations (N = 300, T = 6) the OOB predictions had a slope of 0.65-0.91 on the true effect while beta_differential stayed at 0.88-1.07, and calibrated predictions did not reduce RMSE. A warning is emitted in that case.

.. versionchanged:: 1.30.0 Forests with fixed effects and a binary treatment calibrate against imputation scores by default (method='within' restores the earlier regression).

Parameters:

Name	Type	Description	Default
forest	CausalForest	A forest fitted with the GRF engine (split_rule="grf").	required
newdata	array - like	Effect-modifier rows to calibrate. None calibrates the out-of-bag predictions of the training rows. X= is accepted as an alias.	None
alpha	float	Level for the reported slope confidence intervals.	0.05
method	('auto', 'imputation', 'within')	Forests with fixed effects only; see :func:calibration_test.	'auto'
covariates	'none', 'auto', list of str or array	Imputation regression only; see :func:average_treatment_effect.	'none'

Returns:

Type	Description
dict	cate (calibrated predictions), raw_cate, beta_mean, beta_differential, their *_se, calibration (the full :func:calibration_test frame), heterogeneity_detected (one-sided 5% test of beta_differential > 0) and warnings.

Examples:

>>> import numpy as np
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 1000
>>> X = rng.normal(size=(n, 3))
>>> T = rng.integers(0, 2, n)
>>> Y = X[:, 1] + (1 + X[:, 0]) * T + rng.normal(size=n)
>>> cf = sp.causal_forest(Y=Y, T=T, X=X, n_estimators=400, random_state=0)
>>> cal = sp.calibrate_cate(cf)
>>> cal["cate"].shape
(1000,)

References

[@chernozhukov2025generic], [@aytug2026attenuated]

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

grf_calibration

grf_calibration(*, Y: ndarray, W: ndarray, Y_hat: ndarray, W_hat: ndarray, tau_hat: ndarray, vce: str = 'HC3', alpha: float = 0.05, weights: Optional[ndarray] = None, clusters: Optional[ndarray] = None) -> DataFrame

Calibration test of any CATE predictions [@chernozhukov2025generic].

The grf::test_calibration regression, given its inputs rather than a fitted forest. Fits, by (weighted) least squares without an intercept,

.. math:: Y_i - \hat m(X_i) = \beta_1\,(W_i - \hat e(X_i))\,\bar\tau + \beta_2\,(W_i - \hat e(X_i))\,(\hat\tau(X_i) - \bar\tau) + \varepsilon_i,

with :math:\bar\tau the (observation-weighted) mean CATE prediction. beta_1 = 1 says the mean prediction is calibrated; beta_2 = 1 says the heterogeneity is calibrated, and beta_2 > 0 significantly is evidence that the predictions order units by their effect at all.

:func:calibration_test runs this regression on a fitted forest's own out-of-bag predictions. Call this function when the predictions come from somewhere else: a forest fitted on other folds, a meta-learner, or a forest whose clustering should be imposed after the fact. The test is only valid if tau_hat[i] (and Y_hat[i], W_hat[i]) was not fitted on row i -- or, with clustered data, on any row of i's cluster. In-sample predictions have learned the noise in Y and make beta_2 look significant.

t tests each coefficient against 0 and p is the one-sided p-value for the alternative beta > 0 from Student's t with n - 2 degrees of freedom -- grf's reporting rule (lmtest::coeftest then halving). vce is grf's vcov.type (default "HC3"); grf passes it to sandwich::vcovCL with clusters (one cluster per observation when there are none), which is what the covariance here reproduces (:func:statspai.forest._grf_inference.cluster_robust_vcov). ci_low / ci_high use the same t quantile. t_vs_zero and p_one_sided repeat t and p under their earlier names.

Parameters:

Name	Type	Description	Default
Y	array-like of shape (n,)	Outcome and treatment.	required
W	array-like of shape (n,)	Outcome and treatment.	required
Y_hat	array-like of shape (n,)	Out-of-sample predictions of E[Y | X] and E[W | X].	required
W_hat	array-like of shape (n,)	Out-of-sample predictions of E[Y | X] and E[W | X].	required
tau_hat	array-like of shape (n,)	Out-of-sample CATE predictions to grade.	required
vce	('HC3', 'HC2', 'HC1', 'HC0')	Sandwich covariance type, clustered when clusters is given.	'HC3'
alpha	float	Level for ci_low / ci_high.	0.05
weights	array-like of shape (n,)	Non-negative observation weights; rows with weight 0 are left out. 1 / (size of the row's cluster) gives every cluster the same total weight (grf's equalize.cluster.weights).	None
clusters	array-like of shape (n,)	Cluster labels for the covariance.	None

Returns:

Type	Description
DataFrame	Index mean_forest_prediction / differential_forest_prediction; columns coef, se, t, p, ci_low, ci_high, t_vs_zero, p_one_sided.

Examples:

Grade forests that did not see a school when predicting for it: fit on four folds of schools, predict the fifth, then test with a covariance clustered by school.

>>> import numpy as np
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n, n_schools = 600, 30
>>> school = rng.integers(0, n_schools, size=n)
>>> X = rng.normal(size=(n, 3))
>>> W = rng.binomial(1, 0.5, size=n).astype(float)
>>> Y = X[:, 1] + (1.0 + X[:, 0]) * W + rng.normal(size=n)
>>> kw = dict(n_estimators=100, random_state=0)
>>> fold = school % 5
>>> Y_hat, W_hat, tau_hat = np.empty(n), np.empty(n), np.empty(n)
>>> for k in range(5):
...     tr, te = fold != k, fold == k
...     my = sp.regression_forest(y=Y[tr], covariates=X[tr], **kw)
...     mw = sp.regression_forest(y=W[tr], covariates=X[tr], **kw)
...     Y_hat[te] = np.ravel(my.predict(X[te]))
...     W_hat[te] = np.ravel(mw.predict(X[te]))
...     cf = sp.causal_forest(Y=Y[tr], T=W[tr], X=X[tr], **kw)
...     tau_hat[te] = cf.effect(X[te])
>>> out = sp.grf_calibration(Y=Y, W=W, Y_hat=Y_hat, W_hat=W_hat,
...                          tau_hat=tau_hat, clusters=school)
>>> out.index.tolist()
['mean_forest_prediction', 'differential_forest_prediction']

References

[@chernozhukov2025generic], [@athey2019generalized]

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)

best_linear_projection

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

Best linear projection of the conditional effect on covariates.

Regresses the forest's doubly-robust scores on (1, A) with a heteroskedasticity- (cluster-, when the forest has clusters) robust covariance -- grf::best_linear_projection. The coefficients are the population best linear predictor of tau(X) given A (Semenova and Chernozhukov 2021); A=None gives the doubly-robust average effect.

Parameters:

Name	Type	Description	Default
forest	fitted forest, or a DML result	sp.causal_forest (GRF engine), sp.iv_forest, sp.multi_arm_forest (one projection per contrast, stacked), sp.causal_survival_forest. Also the result of sp.dml(model='irm') with the ATE score: its cross-fitted doubly-robust scores are projected the same way (DoubleML's cate() / gate() with A a basis or group indicators), and the joint covariance is in .attrs['vcov'].	required
A	array - like or DataFrame	Projection covariates, one row per training observation (after missing-value removal).	None
vce	('HC3', 'HC0', 'HC1', 'HC2')	Covariance type (sandwich::vcovCL conventions).	"HC3"
alpha	float	Significance level for the intervals.	0.05
**kwargs	Any	Forwarded to the forest's score construction.	{}

Returns:

Type	Description
DataFrame	Rows Intercept and one per column of A with coef, se, t, p, ci_lower, ci_upper. For causal forests and DML fits .attrs['vcov'] holds the covariance of the coefficients, from which the standard error of the projected effect at any covariate value, sqrt(a' V a), follows.

Examples:

>>> import numpy as np, statspai as sp
>>> rng = np.random.default_rng(0)
>>> X = rng.normal(size=(1000, 3)); T = rng.binomial(1, 0.5, 1000)
>>> Y = (1 + X[:, 0]) * T + rng.normal(size=1000)
>>> cf = sp.causal_forest(Y=Y, T=T, X=X, n_estimators=400, random_state=1)
>>> sp.best_linear_projection(cf, A=X[:, :1])

References

[@semenova2021debiased], [@athey2019generalized]

get_scores

get_scores(forest: Any, **kwargs: Any) -> ndarray

Doubly-robust scores of a forest's average effect (grf::get_scores).

The mean of the scores is the forest's doubly-robust average effect; regressing them on covariates is the best linear projection. Every score uses the forest's out-of-bag conditional effect.

Parameters:

Name	Type	Description	Default
forest	fitted forest	sp.causal_forest (GRF engine, binary or continuous treatment), sp.iv_forest (average conditional LATE), sp.multi_arm_forest (one column per contrast), sp.causal_survival_forest.	required
**kwargs	Any	Forwarded to the forest's own get_scores (e.g. compliance_score= for an instrumental forest).	{}

Returns:

Type	Description
ndarray	(n,) scores, or (n, K-1) for a multi-arm forest.

Examples:

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

References

[@athey2019generalized], [@robins1994estimation]

variable_importance

variable_importance(forest: Any, decay_exponent: float = 2.0, max_depth: int = 4) -> Series

Split-frequency variable importance of a forest (grf::variable_importance).

For depths d = 1..max_depth the share of splits made on each covariate at that depth is weighted by d^(-decay_exponent); the weighted shares are summed and divided by the total weight, so the importances sum to one. This measures how often the forest uses a covariate near the root -- for a causal forest, how often it splits on it to separate effect heterogeneity -- and is not a test of anything.

Parameters:

Name	Type	Description	Default
forest	fitted GRF-engine forest	Any of sp.causal_forest (GRF engine), sp.iv_forest, sp.multi_arm_forest, sp.lm_forest, sp.regression_forest, sp.multi_regression_forest, sp.probability_forest, sp.quantile_forest, sp.survival_forest, sp.causal_survival_forest.	required
decay_exponent	float	How quickly deeper splits lose weight.	2
max_depth	int	Deepest level counted.	4

Returns:

Type	Description
Series	Importance per covariate, in covariate order. The importances sum to one 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).

Examples:

>>> import numpy as np, statspai as sp
>>> rng = np.random.default_rng(0)
>>> X = rng.normal(size=(1000, 3)); T = rng.binomial(1, 0.5, 1000)
>>> Y = X[:, 0] * T + rng.normal(size=1000)
>>> cf = sp.causal_forest(Y=Y, T=T, X=X, n_estimators=200, random_state=1)
>>> sp.variable_importance(cf).idxmax()
'X0'

References

[@athey2019generalized]

instrumental_forest

instrumental_forest(*args: Any, **kwargs: Any) -> IVForestResult

Alias of :func:iv_forest under grf's name.

multi_regression_forest

multi_regression_forest(data: Optional[DataFrame] = None, y: Any = None, covariates: Any = None, *, clusters: Any = None, weights: Any = None, equalize_cluster_weights: bool = False, n_estimators: int = 2000, min_samples_leaf: int = 5, max_samples: float = 0.5, mtry: Optional[int] = None, honest: bool = True, honesty_fraction: float = 0.5, honesty_prune_leaves: bool = True, split_alpha: float = 0.05, imbalance_penalty: float = 0.0, ci_group_size: int = 1, max_depth: Optional[int] = None, random_state: Optional[int] = 42, n_jobs: int = 1) -> PredictionForest

Multi-task regression forest for E[Y_1, ..., Y_q | X] (grf::multi_regression_forest).

Splits maximise the summed between-child variation of all outcomes, so one tree serves every outcome; outcomes should be on comparable scales.

Parameters:

Name	Type	Description	Default
y	list of str or 2-D array	Outcomes. (see below for the shared forest options; ci_group_size defaults to 1 as grf does not report variances here)	None

Returns:

Type	Description
PredictionForest	.predictions of shape (n, q).

Examples:

>>> import numpy as np, statspai as sp
>>> rng = np.random.default_rng(0)
>>> X = rng.normal(size=(500, 2))
>>> Y = np.c_[X[:, 0], -X[:, 0]] + rng.normal(size=(500, 2))
>>> sp.multi_regression_forest(y=Y, covariates=X,
...                            n_estimators=100).predictions.shape
(500, 2)

References

[@athey2019generalized]

probability_forest

probability_forest(data: Optional[DataFrame] = None, y: Any = None, covariates: Any = None, *, clusters: Any = None, weights: Any = None, equalize_cluster_weights: bool = False, n_estimators: int = 2000, min_samples_leaf: int = 5, max_samples: float = 0.5, mtry: Optional[int] = None, honest: bool = True, honesty_fraction: float = 0.5, honesty_prune_leaves: bool = True, split_alpha: float = 0.05, imbalance_penalty: float = 0.0, ci_group_size: int = 2, max_depth: Optional[int] = None, random_state: Optional[int] = 42, n_jobs: int = 1) -> PredictionForest

Probability forest for P[Y = k | X] (grf::probability_forest).

Trees split on class counts (the Gini-type criterion of a regression forest on the class indicators); leaves hold class frequencies of the honest estimation sample, so every prediction is a proper probability vector.

Parameters:

Name	Type	Description	Default
y	str or array - like	Class labels (any hashable values; classes are sorted). (see below for the shared forest options)	None

Returns:

Type	Description
PredictionForest	.predictions (n, K) out-of-bag probabilities with columns named by class; .classes.

Examples:

>>> import numpy as np, statspai as sp
>>> rng = np.random.default_rng(0)
>>> X = rng.normal(size=(600, 2)); cls = (X[:, 0] > 0).astype(int)
>>> pf = sp.probability_forest(y=cls, covariates=X, n_estimators=200)
>>> pf.predictions.shape
(600, 2)

References

[@athey2019generalized]

quantile_forest

quantile_forest(data: Optional[DataFrame] = None, y: Any = None, covariates: Any = None, *, quantiles: Sequence[float] = (0.1, 0.5, 0.9), regression_splitting: bool = False, clusters: Any = None, weights: Any = None, equalize_cluster_weights: bool = False, n_estimators: int = 2000, min_samples_leaf: int = 5, max_samples: float = 0.5, mtry: Optional[int] = None, honest: bool = True, honesty_fraction: float = 0.5, honesty_prune_leaves: bool = True, split_alpha: float = 0.05, imbalance_penalty: float = 0.0, max_depth: Optional[int] = None, random_state: Optional[int] = 42, n_jobs: int = 1) -> PredictionForest

Quantile forest for conditional quantiles of Y (grf::quantile_forest).

At each node the outcome is relabelled by the interval it falls in between the node's empirical quantiles, and the split maximises the between-child separation of those class counts -- so trees chase changes anywhere in the conditional distribution, not only the mean (ATW 2019, Sec. 3.3). Predictions are quantiles of the training outcomes under the forest weights (Meinshausen 2006).

Parameters:

Name	Type	Description	Default
y	str or array - like	Outcome.	None
quantiles	sequence of float	Quantiles that calibrate the splits and are predicted by default.	(0.1, 0.5, 0.9)
regression_splitting	bool	Use ordinary regression splits instead (grf regression.splitting; the quantile regression forest of Meinshausen 2006). (see below for the shared forest options)	False

Returns:

Type	Description
PredictionForest	.predictions (n, len(quantiles)) out-of-bag quantiles; .predict(newdata, quantiles=...) for other levels.

Examples:

>>> import numpy as np, statspai as sp
>>> rng = np.random.default_rng(0)
>>> X = rng.normal(size=(800, 2))
>>> y = X[:, 0] + (1 + (X[:, 1] > 0)) * rng.normal(size=800)
>>> qf = sp.quantile_forest(y=y, covariates=X, n_estimators=200)
>>> qf.predict(X[:3]).columns.tolist()
['q0.1', 'q0.5', 'q0.9']

References

[@athey2019generalized], [@meinshausen2006quantile]

regression_forest

regression_forest(data: Optional[DataFrame] = None, y: Any = None, covariates: Any = None, *, clusters: Any = None, weights: Any = None, equalize_cluster_weights: bool = False, n_estimators: int = 2000, min_samples_leaf: int = 5, max_samples: float = 0.5, mtry: Optional[int] = None, honest: bool = True, honesty_fraction: float = 0.5, honesty_prune_leaves: bool = True, split_alpha: float = 0.05, imbalance_penalty: float = 0.0, ci_group_size: int = 2, max_depth: Optional[int] = None, random_state: Optional[int] = 42, n_jobs: int = 1) -> PredictionForest

Honest regression forest for E[Y | X] (grf::regression_forest).

Parameters:

Name	Type	Description	Default
y	str or array - like	Outcome. (see below for the shared forest options)	None

Returns:

Type	Description
PredictionForest	.predictions (out-of-bag), .variance (little bags), .predict(newdata, estimate_variance=True), .variable_importance().

Examples:

>>> import numpy as np, statspai as sp
>>> rng = np.random.default_rng(0)
>>> X = rng.normal(size=(500, 3)); y = X[:, 0] ** 2 + rng.normal(size=500)
>>> rf = sp.regression_forest(y=y, covariates=X, n_estimators=200)
>>> rf.predict(X[:5]).shape
(5, 1)

References

[@athey2019generalized], [@wager2018estimation]

tune_causal_forest

tune_causal_forest(formula: Optional[str] = None, data: Optional[DataFrame] = None, Y: Optional[ndarray] = None, T: Optional[ndarray] = None, X: Optional[ndarray] = None, *, parameters: Sequence[str] = ('min_samples_leaf', 'max_samples', 'mtry'), n_draws: int = 40, tune_trees: int = 200, tune_reps: int = 2, n_estimators: int = 2000, random_state: Optional[int] = None, **kwargs: Any) -> Dict[str, Any]

Choose causal-forest settings by out-of-bag R-loss, then fit.

Parameters:

Name	Type	Description	Default
formula	Optional[str]	As in :func:statspai.causal_forest. covariates= is accepted for X.	None
data	Optional[str]	As in :func:statspai.causal_forest. covariates= is accepted for X.	None
Y	Optional[str]	As in :func:statspai.causal_forest. covariates= is accepted for X.	None
T	Optional[str]	As in :func:statspai.causal_forest. covariates= is accepted for X.	None
X	Optional[str]	As in :func:statspai.causal_forest. covariates= is accepted for X.	None
parameters	sequence of str	Which settings to tune, from 'min_samples_leaf', 'max_samples', 'mtry', 'honesty_fraction', 'honesty_prune_leaves' and 'imbalance_penalty'. The default tunes the three that matter most.	('min_samples_leaf', 'max_samples', 'mtry')
n_draws	int	Random settings tried, besides the defaults.	40
tune_trees	int	Trees in each trial forest. Small forests are noisy, which is why a tuned setting is adopted only if it beats the defaults clearly (see Notes).	200
tune_reps	int	Trial forests per setting, with different seeds; their losses are averaged and their spread gives the noise of the comparison.	2
n_estimators	int	Trees in the final forest.	2000
random_state	int		None
**kwargs	Any	Passed to every forest (for example clusters=, model_y=).	{}

Returns:

Type	Description
dict	forest (the final :class:~statspai.forest.causal_forest.CausalForest), best_params, tuned (whether a tuned setting replaced the defaults), default_error and tuned_error (out-of-bag R-loss), noise (standard deviation of the loss across repetitions of the same setting) and trials (a DataFrame, one row per setting, sorted by loss).

Notes

The nuisance functions m and e are estimated once, out of fold, by the forest with the default settings, and shared by every trial, so settings are compared on the same residuals. When no trial is adopted that default forest is what is returned. A tuned setting is adopted only when its loss is below the defaults' by more than twice the repetition noise; otherwise the defaults are kept and tuned is False. Picking the minimum of forty noisy losses without such a margin would select noise.

What to expect, from 12 replications each of four designs with 1,500 rows and 30 draws (root mean squared error of the out-of-bag effects against the truth, tuned against default): a constant effect, 0.082 against 0.119, with a tuned setting adopted 10 times out of 12, always with much larger leaves; a step and a fast-varying effect, no difference, the defaults being kept almost every time; a linear effect, 0.184 against 0.172, slightly worse. The R-loss is a noisy criterion and a lower loss is not always a lower error of the effect itself. Tuning pays when the effect varies little; it is not a free improvement, and trials is there to be looked at.

grf fits a smooth surface to the trial losses and minimises it. Here the best trial itself is taken, with the margin above; the two choose from the same kind of evidence but will not return the same setting.

Examples:

>>> import numpy as np
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 800
>>> X = rng.normal(size=(n, 4))
>>> W = rng.binomial(1, 0.5, n)
>>> Y = X[:, 0] + W * (1 + 2 * (X[:, 1] > 0)) + rng.normal(size=n)
>>> out = sp.tune_causal_forest(Y=Y, T=W, X=X, n_draws=10,
...                             n_estimators=200, random_state=1)
>>> sorted(out)[:3]
['best_params', 'default_error', 'forest']
>>> bool(out["tuned_error"] <= out["default_error"])
True

References

[@nie2021quasi], [@athey2019generalized]