`statspai.conformal_causal`¶

conformal_causal ¶

Conformal Causal Inference: Distribution-free prediction intervals for ITE.

Provides prediction intervals for individual treatment effects (ITE) without distributional assumptions, using conformal inference.

References

Lei, L. & Candes, E. J. (2021). Conformal Inference of Counterfactuals and Individual Treatment Effects. JRSS-B, 83(5), 911-938. [@lei2021conformal]

Chernozhukov, V., Wuthrich, K., & Zhu, Y. (2021). An Exact and Robust Conformal Inference Method for Counterfactual and Synthetic Controls. JASA, 116(536), 1849-1864. [@chernozhukov2021exact]

DebiasedConformalResult `dataclass` ¶

Bases: ResultProtocolMixin

Debiased ML conformal counterfactual intervals.

Examples:

>>> import statspai as sp
>>> import numpy as np
>>> import pandas as pd
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> d = rng.binomial(1, 1 / (1 + np.exp(-0.5 * x1)))
>>> y = 1.0 + 0.8 * x1 + 0.5 * x2 + 1.5 * d + rng.normal(0, 0.5, n)
>>> df = pd.DataFrame({"y": y, "d": d, "x1": x1, "x2": x2})
>>> res = sp.conformal_debiased_ml(
...     df, y="y", treat="d", covariates=["x1", "x2"], seed=1)
>>> res.n_test
200
>>> bool(res.point_estimate.shape == (200,))
True

ConformalDensityResult `dataclass` ¶

Bases: ResultProtocolMixin

Conditional-density conformal ITE intervals.

Produced by :func:conformal_density_ite. Holds the per-test-point intervals (lower/upper), the ITE point_estimate and a formatted .summary().

Examples:

>>> import statspai as sp
>>> import numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> n = 300
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> d = rng.integers(0, 2, n)
>>> y = 1.5 * d + 0.5 * x1 - 0.3 * x2 + rng.normal(size=n)
>>> df = pd.DataFrame({"y": y, "d": d, "x1": x1, "x2": x2})
>>> res = sp.conformal_density_ite(
...     df, y="y", treat="d", covariates=["x1", "x2"], alpha=0.1, seed=0)
>>> type(res).__name__
'ConformalDensityResult'
>>> res.intervals.shape
(300, 2)
>>> bool((res.intervals[:, 0] <= res.intervals[:, 1]).all())
True
>>> isinstance(res.summary(), str)
True

FairConformalResult `dataclass` ¶

Bases: ResultProtocolMixin

Fairness-aware conformal ITE intervals.

Examples:

>>> import statspai as sp
>>> import numpy as np
>>> import pandas as pd
>>> rng = np.random.default_rng(0)
>>> n = 240
>>> g = rng.integers(0, 2, size=n)
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> d = rng.binomial(1, 0.5, size=n)
>>> y = 1.0 + 0.7 * x1 + 0.4 * x2 + 1.2 * d + 0.5 * g + rng.normal(0, 0.5, n)
>>> df = pd.DataFrame({"y": y, "d": d, "x1": x1, "x2": x2, "grp": g})
>>> res = sp.conformal_fair_ite(
...     df, y="y", treat="d", covariates=["x1", "x2"],
...     protected="grp", seed=1)
>>> bool(res.point_estimate.shape == (240,))
True
>>> sorted(res.group_coverage_targets.keys())
['0', '1']

ConformalCATE ¶

Conformal prediction intervals for individual treatment effects.

Parameters:

Name	Type	Default
`data`	`DataFrame`	required
`y`	`str`	required
`treat`	`str`	required
`covariates`	`list of str`	required
`model`	`sklearn estimator`	`None`
`alpha`	`float`	`0.05`
`calib_fraction`	`float`	`0.25`
`random_state`	`int`	`42`

Examples:

>>> import statspai as sp
>>> import numpy as np, pandas as pd
>>> from sklearn.linear_model import LinearRegression
>>> from statspai.conformal_causal.conformal_ite import ConformalCATE
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> x1 = rng.normal(0, 1, n)
>>> x2 = rng.normal(0, 1, n)
>>> treat = rng.binomial(1, 0.5, n)
>>> y = (1.0 + 0.5 * x1 + treat * (1.0 + 0.5 * x2)
...      + rng.normal(0, 0.5, n))
>>> df = pd.DataFrame({"outcome": y, "treatment": treat,
...                    "x1": x1, "x2": x2})
>>> est = ConformalCATE(df, y="outcome", treat="treatment",
...                     covariates=["x1", "x2"],
...                     model=LinearRegression(), random_state=0)
>>> res = est.fit()
>>> lo = res.model_info["cate_lower"]
>>> hi = res.model_info["cate_upper"]
>>> len(lo) == n
True
>>> bool(np.all(hi >= lo))
True

fit ¶

fit() -> CausalResult

Compute conformal CATE intervals.

predict ¶

predict(X_new: ndarray) -> Dict[str, ndarray]

Predict CATE with conformal intervals for new data.

Returns:

Type	Description
`dict with 'cate', 'lower', 'upper'`

MultiDPConformalResult `dataclass` ¶

Bases: ResultProtocolMixin

Multi-decision-point conformal ITE intervals.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> x1 = rng.normal(0, 1, n)
>>> d1 = rng.binomial(1, 0.5, n)
>>> y1 = 1.0 + 0.5 * x1 + 0.8 * d1 + rng.normal(0, 0.5, n)
>>> x2 = 0.5 * x1 + rng.normal(0, 1, n)
>>> d2 = rng.binomial(1, 0.5, n)
>>> y2 = 0.5 + 0.3 * x2 + 0.6 * d2 + rng.normal(0, 0.5, n)
>>> df = pd.DataFrame({"x1": x1, "d1": d1, "y1": y1,
...                    "x2": x2, "d2": d2, "y2": y2})
>>> res = sp.conformal_ite_multidp(
...     df,
...     y_per_stage=["y1", "y2"],
...     treat_per_stage=["d1", "d2"],
...     history_per_stage=[["x1"], ["x1", "x2"]],
...     alpha=0.1,
... )
>>> isinstance(res, sp.MultiDPConformalResult)
True
>>> res.n_stages
2
>>> res.cumulative_interval.shape
(200, 2)

ConformalCounterfactualResult `dataclass` ¶

Bases: ResultProtocolMixin

Counterfactual prediction intervals under each potential outcome.

Returned by :func:conformal_counterfactual; holds per-arm prediction bands for Y(1) | X and Y(0) | X.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> x1, x2 = rng.normal(size=n), rng.normal(size=n)
>>> t = rng.binomial(1, 0.5, size=n)
>>> y = 1.0 + 2.0 * t + 0.5 * x1 + rng.normal(scale=0.5, size=n)
>>> df = pd.DataFrame({"y": y, "t": t, "x1": x1, "x2": x2})
>>> res = sp.conformal_counterfactual(
...     df, y="y", treat="t", covariates=["x1", "x2"],
...     alpha=0.1, random_state=0)
>>> isinstance(res, sp.ConformalCounterfactualResult)
True
>>> bool(np.all(res.upper_Y1 >= res.lower_Y1))
True

ConformalITEResult `dataclass` ¶

Bases: ResultProtocolMixin

Prediction intervals for the individual treatment effect τ(x).

Returned by :func:conformal_ite_interval; holds the point estimate and lower/upper bound for τ(x) = Y(1) - Y(0).

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> x1, x2 = rng.normal(size=n), rng.normal(size=n)
>>> t = rng.binomial(1, 0.5, size=n)
>>> y = 1.0 + 2.0 * t + 0.5 * x1 + rng.normal(scale=0.5, size=n)
>>> df = pd.DataFrame({"y": y, "t": t, "x1": x1, "x2": x2})
>>> res = sp.conformal_ite_interval(
...     df, y="y", treat="t", covariates=["x1", "x2"],
...     alpha=0.1, random_state=0)
>>> isinstance(res, sp.ConformalITEResult)
True
>>> list(res.to_frame().columns)
['tau', 'tau_lower', 'tau_upper']

ContinuousConformalResult `dataclass` ¶

Bases: ResultProtocolMixin

Output of :func:conformal_continuous.

Examples:

>>> import numpy as np, pandas as pd, statspai as sp
>>> rng = np.random.default_rng(0)
>>> t = rng.uniform(0, 5, 300); x = rng.normal(size=300)
>>> train = pd.DataFrame({"y": 2 + 0.7 * t + 0.5 * x + rng.normal(0, 0.5, 300),
...                       "t": t, "x": x})
>>> res = sp.conformal_continuous(train, y="y", treatment="t",
...     covariates=["x"], test_data=pd.DataFrame({"t": [1.0], "x": [0.0]}))
>>> isinstance(res, sp.ContinuousConformalResult)
True

InterferenceConformalResult `dataclass` ¶

Bases: ResultProtocolMixin

Output of :func:conformal_interference.

Examples:

>>> import numpy as np, pandas as pd, statspai as sp
>>> rng = np.random.default_rng(0)
>>> df = pd.DataFrame({"cluster": np.repeat(np.arange(12), 8),
...                    "treat": rng.integers(0, 2, 96), "x": rng.normal(size=96)})
>>> df["y"] = 1 + 0.6 * df.treat + 0.4 * df.x + rng.normal(0, 0.3, 96)
>>> res = sp.conformal_interference(df, y="y", treatment="treat",
...     cluster="cluster", covariates=["x"], test_clusters=[0, 1])
>>> isinstance(res, sp.InterferenceConformalResult)
True

conformal_debiased_ml ¶

conformal_debiased_ml(data: DataFrame, y: str, treat: str, covariates: List[str], test_data: Optional[DataFrame] = None, alpha: float = 0.1, n_folds: int = 5, seed: int = 0) -> DebiasedConformalResult

Debiased ML conformal counterfactual intervals.

Parameters:

Name	Type	Description	Default
`data`	`DataFrame`		required
`y`	`str`		required
`treat`	`str`		required
`covariates`	`list of str`		required
`test_data`	`DataFrame`		`None`
`alpha`	`float`		`0.1`
`n_folds`	`int`	Cross-fitting folds (debiased step).	`5`
`seed`	`int`		`0`

Returns:

Type	Description
`DebiasedConformalResult`

Examples:

>>> import statspai as sp
>>> import numpy as np
>>> import pandas as pd
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> ps = 1 / (1 + np.exp(-0.5 * x1))
>>> d = rng.binomial(1, ps)
>>> y = 1.0 + 0.8 * x1 + 0.5 * x2 + 1.5 * d + rng.normal(0, 0.5, n)
>>> df = pd.DataFrame({"y": y, "d": d, "x1": x1, "x2": x2})
>>> res = sp.conformal_debiased_ml(
...     df, y="y", treat="d", covariates=["x1", "x2"],
...     alpha=0.1, n_folds=5, seed=1)
>>> isinstance(res, sp.DebiasedConformalResult)
True
>>> res.intervals.shape
(200, 2)
>>> widths = res.intervals[:, 1] - res.intervals[:, 0]
>>> bool((widths > 0).all())
True
>>> print(res.summary())

conformal_density_ite ¶

conformal_density_ite(data: DataFrame, y: str, treat: str, covariates: List[str], test_data: Optional[DataFrame] = None, alpha: float = 0.1, bandwidth: Optional[float] = None, seed: int = 0) -> ConformalDensityResult

Conditional-density conformal ITE intervals.

Parameters:

Name	Type	Description	Default
`data`	`DataFrame`	Training data with both treated and control units.	required
`y`	`str`		required
`treat`	`str`		required
`covariates`	`list of str`		required
`test_data`	`DataFrame`	Test set; defaults to `data` (in-sample intervals).	`None`
`alpha`	`float`	Miscoverage; intervals target 1 − α coverage.	`0.1`
`bandwidth`	`float`	KDE bandwidth; defaults to Silverman's rule on calibration Y.	`None`
`seed`	`int`		`0`

Returns:

Type	Description
`ConformalDensityResult`

Examples:

>>> import statspai as sp
>>> import numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> n = 300
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> d = rng.integers(0, 2, n)
>>> y = 1.5 * d + 0.5 * x1 - 0.3 * x2 + rng.normal(size=n)
>>> df = pd.DataFrame({"y": y, "d": d, "x1": x1, "x2": x2})
>>> res = sp.conformal_density_ite(
...     df, y="y", treat="d", covariates=["x1", "x2"], alpha=0.1, seed=0)
>>> res.intervals.shape
(300, 2)
>>> res.point_estimate.shape
(300,)
>>> bool((res.intervals[:, 0] <= res.intervals[:, 1]).all())
True

conformal_fair_ite ¶

conformal_fair_ite(data: DataFrame, y: str, treat: str, covariates: List[str], protected: str, test_data: Optional[DataFrame] = None, alpha: float = 0.1, seed: int = 0) -> FairConformalResult

Counterfactual-fair conformal ITE intervals.

Parameters:

Name	Type	Description	Default
`data`	`DataFrame`		required
`y`	`str`		required
`treat`	`str`		required
`covariates`	`list of str`	Predictive features. `protected` is excluded from the outcome regression to preserve counterfactual fairness, but used downstream for stratified calibration.	required
`protected`	`str`	Protected attribute column (categorical).	required
`test_data`	`DataFrame`		`None`
`alpha`	`float`	Miscoverage of the effect interval. Within each protected group, each arm's residuals get a split-conformal quantile at `1 - alpha / 2` and the half-width is their sum, so the interval targets coverage of `Y(1) - Y(0)` of at least `1 - alpha` in each group (union bound; conservative) when treatment is randomised given the covariates.	`0.1`
`seed`	`int`		`0`

Returns:

Type	Description
`FairConformalResult`

Examples:

>>> import statspai as sp
>>> import numpy as np
>>> import pandas as pd
>>> rng = np.random.default_rng(0)
>>> n = 240
>>> g = rng.integers(0, 2, size=n)
>>> x1 = rng.normal(size=n)
>>> x2 = rng.normal(size=n)
>>> d = rng.binomial(1, 0.5, size=n)
>>> y = 1.0 + 0.7 * x1 + 0.4 * x2 + 1.2 * d + 0.5 * g + rng.normal(0, 0.5, n)
>>> df = pd.DataFrame({"y": y, "d": d, "x1": x1, "x2": x2, "grp": g})
>>> res = sp.conformal_fair_ite(
...     df, y="y", treat="d", covariates=["x1", "x2"],
...     protected="grp", alpha=0.1, seed=1)
>>> isinstance(res, sp.FairConformalResult)
True
>>> res.intervals.shape
(240, 2)
>>> sorted(res.group_widths.keys())
['0', '1']
>>> widths = res.intervals[:, 1] - res.intervals[:, 0]
>>> bool((widths > 0).all())
True
>>> print(res.summary())

conformal_cate ¶

conformal_cate(data: DataFrame, y: str, treat: str, covariates: List[str], model: Optional[BaseEstimator] = None, alpha: float = 0.05, calib_fraction: float = 0.25, random_state: int = 42) -> CausalResult

Compute conformal prediction intervals for CATE.

Parameters:

Name	Type	Description	Default
`data`	`DataFrame`	Input data.	required
`y`	`str`	Outcome variable.	required
`treat`	`str`	Binary treatment variable (0/1).	required
`covariates`	`list of str`	Covariate names.	required
`model`	`sklearn estimator`	Outcome model for mu_d(X). If None, uses GBM.	`None`
`alpha`	`float`	Miscoverage level. Intervals have (1-alpha) coverage.	`0.05`
`calib_fraction`	`float`	Fraction of data used for calibration.	`0.25`
`random_state`	`int`		`42`

Returns:

Type Description

CausalResult

estimate / se are the average treatment effect and its standard error from the cross-fitted AIPW score (model for the outcome in each arm, a gradient-boosting propensity, 5 folds), as in :func:sp.metalearner. The mean of the split-sample CATE predictions is model_info['cate_mean'] and has no standard error. Includes CATE point estimates and prediction intervals. model_info contains: - 'cate': point estimates - 'cate_lower': lower bounds of prediction intervals - 'cate_upper': upper bounds of prediction intervals - 'interval_width': average width of prediction intervals

Examples:

>>> import statspai as sp
>>> import numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> x1 = rng.normal(0, 1, n)
>>> x2 = rng.normal(0, 1, n)
>>> treat = rng.binomial(1, 0.5, n)
>>> y = (1.0 + 0.5 * x1 + treat * (1.0 + 0.5 * x2)
...      + rng.normal(0, 0.5, n))
>>> df = pd.DataFrame({"outcome": y, "treatment": treat,
...                    "x1": x1, "x2": x2})
>>> result = sp.conformal_cate(df, y="outcome", treat="treatment",
...                            covariates=["x1", "x2"], random_state=0)
>>> cate_lower = result.model_info["cate_lower"]
>>> cate_upper = result.model_info["cate_upper"]
>>> len(cate_lower) == n
True
>>> bool(np.all(cate_upper >= cate_lower))
True

conformal_ite_multidp ¶

conformal_ite_multidp(data: DataFrame, y_per_stage: List[str], treat_per_stage: List[str], history_per_stage: List[List[str]], test_data: Optional[DataFrame] = None, alpha: float = 0.1, seed: int = 0) -> MultiDPConformalResult

Multi-decision-point conformal ITE.

Parameters:

Name	Type	Description	Default
`data`	`DataFrame`	Long-format wide data: each row = subject; columns include stage-specific outcomes / treatments / histories.	required
`y_per_stage`	`list of str`	Outcome column at each stage k = 1, ..., K.	required
`treat_per_stage`	`list of str`	Binary treatment at each stage.	required
`history_per_stage`	`list of list of str`	History (covariates available at decision time k).	required
`test_data`	`DataFrame`		`None`
`alpha`	`float`		`0.1`
`seed`	`int`		`0`

Returns:

Type	Description
`MultiDPConformalResult`

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> x1 = rng.normal(0, 1, n)
>>> d1 = rng.binomial(1, 0.5, n)
>>> y1 = 1.0 + 0.5 * x1 + 0.8 * d1 + rng.normal(0, 0.5, n)
>>> x2 = 0.5 * x1 + rng.normal(0, 1, n)
>>> d2 = rng.binomial(1, 0.5, n)
>>> y2 = 0.5 + 0.3 * x2 + 0.6 * d2 + rng.normal(0, 0.5, n)
>>> df = pd.DataFrame({"x1": x1, "d1": d1, "y1": y1,
...                    "x2": x2, "d2": d2, "y2": y2})
>>> res = sp.conformal_ite_multidp(
...     df,
...     y_per_stage=["y1", "y2"],
...     treat_per_stage=["d1", "d2"],
...     history_per_stage=[["x1"], ["x1", "x2"]],
...     alpha=0.1,
... )
>>> res.n_stages
2
>>> len(res.intervals_per_stage)
2

conformal_counterfactual ¶

conformal_counterfactual(data: DataFrame, y: str, treat: str, covariates: list, X_test: Optional[ndarray] = None, *, alpha: float = 0.1, calib_frac: float = 0.3, model: Optional[BaseEstimator] = None, propensity_model: Optional[BaseEstimator] = None, random_state: Optional[int] = None) -> ConformalCounterfactualResult

Prediction intervals for the counterfactual potential outcomes Y(1) | X and Y(0) | X (Lei & Candès 2021 Theorem 1).

Uses weighted split-conformal separately for each treatment arm, with the propensity score providing the covariate-shift weight between the treated sub-population and the overall population.

Parameters:

Name	Type	Description	Default
`data`	`DataFrame`		required
`y`	`str`	Outcome and 0/1 treatment column names.	required
`treat`	`str`	Outcome and 0/1 treatment column names.	required
`covariates`	`list of str`	Columns used as features for both the outcome and propensity models.	required
`X_test`	`array`	Points at which to return intervals. Defaults to `data[covariates]`.	`None`
`alpha`	`float`	Miscoverage level.	`0.1`
`calib_frac`	`float`	Fraction of each arm used for calibration.	`0.3`
`model`	`sklearn estimators`	Defaults: :class:`RandomForestRegressor` / :class:`LogisticRegression`.	`None`
`propensity_model`	`sklearn estimators`	Defaults: :class:`RandomForestRegressor` / :class:`LogisticRegression`.	`None`
`random_state`	`int`		`None`

Returns:

Type	Description
`ConformalCounterfactualResult`	A band is infinite at a test point whose own weight (`1/e(x)` for Y(1), `1/(1-e(x))` for Y(0)) carries more than `alpha` of the total weight: the calibration arm has too little overlap there for a finite interval with the coverage guarantee. `point_Y1` / `point_Y0` hold the interval centres.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> x1, x2 = rng.normal(size=n), rng.normal(size=n)
>>> t = rng.binomial(1, 0.5, size=n)
>>> y = 1.0 + 2.0 * t + 0.5 * x1 + rng.normal(scale=0.5, size=n)
>>> df = pd.DataFrame({"y": y, "t": t, "x1": x1, "x2": x2})
>>> res = sp.conformal_counterfactual(
...     df, y="y", treat="t", covariates=["x1", "x2"],
...     alpha=0.1, random_state=0)
>>> list(res.to_frame().columns)
['Y1_lower', 'Y1_upper', 'Y0_lower', 'Y0_upper']
>>> len(res.X)
200

References

lei2021conformal

conformal_ite_interval ¶

conformal_ite_interval(data: DataFrame, y: str, treat: str, covariates: list, X_test: Optional[ndarray] = None, *, alpha: float = 0.1, calib_frac: float = 0.3, model: Optional[BaseEstimator] = None, propensity_model: Optional[BaseEstimator] = None, random_state: Optional[int] = None) -> ConformalITEResult

Conformal prediction intervals for the individual treatment effect τ(x) = Y(1) - Y(0).

Implements the Lei-Candès (2021) nested counterfactual bound (Eq. 3.4):

.. math::

[\hat τ(x) - \Delta_1(x) - \Delta_0(x),
 \hat τ(x) + \Delta_1(x) + \Delta_0(x)]

where Δ_a(x) is the half-width of the split-conformal counterfactual interval for arm a at x. This is conservative but finite-sample valid under the usual overlap and SUTVA conditions.

Accepts the same arguments as :func:conformal_counterfactual.

Returns:

Type	Description
`ConformalITEResult`

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> x1, x2 = rng.normal(size=n), rng.normal(size=n)
>>> t = rng.binomial(1, 0.5, size=n)
>>> y = 1.0 + 2.0 * t + 0.5 * x1 + rng.normal(scale=0.5, size=n)
>>> df = pd.DataFrame({"y": y, "t": t, "x1": x1, "x2": x2})
>>> res = sp.conformal_ite_interval(
...     df, y="y", treat="t", covariates=["x1", "x2"],
...     alpha=0.1, random_state=0)
>>> list(res.to_frame().columns)
['tau', 'tau_lower', 'tau_upper']
>>> bool(np.all(res.upper >= res.lower))  # valid intervals
True

References

lei2021conformal

weighted_conformal_prediction ¶

weighted_conformal_prediction(X_train: ndarray, y_train: ndarray, X_calib: ndarray, y_calib: ndarray, X_test: ndarray, weights_calib: Optional[ndarray] = None, model: Optional[BaseEstimator] = None, alpha: float = 0.1, weights_test: Optional[ndarray] = None) -> tuple

Split conformal prediction with per-calibration-point weights.

Implements the Tibshirani-Barber-Candès-Ramdas (2019) weighted split-conformal procedure. When weights_calib is None, this reduces to standard split conformal.

Parameters:

Name	Type	Description	Default
`X_train`	`arrays`	Training fold used to fit the base regression model.	required
`y_train`	`arrays`	Training fold used to fit the base regression model.	required
`X_calib`	`arrays`	Calibration fold used to compute non-conformity scores.	required
`y_calib`	`arrays`	Calibration fold used to compute non-conformity scores.	required
`X_test`	`array`	Points at which to produce prediction intervals.	required
`weights_calib`	`array`	Per-calibration-point likelihood-ratio weights `w_i = f_test(X_i) / f_train(X_i)` for covariate-shift correction. If None, uniform weights.	`None`
`model`	`sklearn-style estimator`	Defaults to `RandomForestRegressor(n_estimators=200, min_samples_leaf=5, random_state=0)`.	`None`
`alpha`	`float`	Miscoverage level (interval targets `1-alpha` coverage).	`0.1`
`weights_test`	`array`	Likelihood-ratio weight of each test point, `w(X_test[j])`, on the same scale as `weights_calib` (same function). The quantile for test point `j` puts mass `w(X_test[j])` on `+inf` next to the calibration weights, so each test point gets its own interval width. Required for the finite-sample coverage guarantee under covariate shift. When `weights_calib` is given and this is not, every test point gets weight 1 (the behaviour before 1.41), which undercovers wherever the test density is high relative to the training density, and an `AssumptionWarning` says so.	`None`

Returns:

Type	Description
(lower, upper, point) : tuple of arrays, each length ``len(X_test)``

Notes

The guarantee is marginal coverage of at least 1 - alpha for a test point drawn from the shifted distribution, provided the weights are the true density ratio up to a common constant [@tibshirani2019conformal]. Under covariate shift, giving the test point weight 1 instead of its own ratio broke it: with training covariates N(0, 1), test covariates N(1.5, 1), 30 calibration points and alpha=0.1, the old intervals covered 81% of test outcomes, the corrected ones 97% (tests/test_te_ml_scan_round5.py).

Examples:

>>> import numpy as np
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> X_train = rng.normal(size=(100, 2))
>>> y_train = X_train[:, 0] + rng.normal(scale=0.3, size=100)
>>> X_calib = rng.normal(size=(50, 2))
>>> y_calib = X_calib[:, 0] + rng.normal(scale=0.3, size=50)
>>> X_test = rng.normal(size=(10, 2))
>>> lower, upper, point = sp.weighted_conformal_prediction(
...     X_train, y_train, X_calib, y_calib, X_test, alpha=0.1)
>>> bool(np.all(upper - lower > 0))  # every band has positive width
True
>>> len(lower)
10

conformal_available_kinds ¶

conformal_available_kinds() -> list[str]

Return the full list of registered conformal kind names.

Examples:

>>> import statspai as sp
>>> kinds = sp.conformal_available_kinds()
>>> bool("cate" in kinds)
True
>>> bool("ite" in kinds and "weighted" in kinds)
True

conformal ¶

conformal(kind: str = 'cate', /, **kwargs: Any) -> Any

Unified entry point for the conformal causal inference family.

Parameters:

Name	Type	Description	Default
`kind`	`str`	The conformal estimator to run. Supported values are listed by `sp.conformal_available_kinds()`.	``"cate"``
`**kwargs`	`Any`	Passed through unchanged to the target function.	`{}`

Returns:

Type	Description
The underlying estimator's return object — e.g. ``CausalResult`` for
func:`conformal_cate`, ``ConformalITEResult`` for ``kind="ite"``,
``ContinuousConformalResult`` for ``kind="continuous"``.

Examples:

>>> import statspai as sp
>>> import numpy as np, pandas as pd
>>> rng = np.random.default_rng(0)
>>> n = 200
>>> x1, x2 = rng.normal(size=n), rng.normal(size=n)
>>> d = rng.integers(0, 2, size=n)
>>> y = 1.0 + 0.5 * x1 + d * (1.0 + 0.5 * x2) + rng.normal(0, 0.5, n)
>>> df = pd.DataFrame({"y": y, "d": d, "x1": x1, "x2": x2})
>>> r = sp.conformal("cate", data=df, y="y", treat="d",
...                   covariates=["x1", "x2"])
>>> type(r).__name__
'CausalResult'

conformal_continuous ¶

conformal_continuous(data: DataFrame, *, y: str, treatment: str, covariates: Sequence[str], test_data: DataFrame, dose_grid: Optional[Sequence[float]] = None, alpha: float = 0.1, estimator: Optional[Any] = None, calibration_frac: float = 0.5, random_state: int = 0) -> ContinuousConformalResult

Split-conformal bands for a continuous-treatment dose response.

Parameters:

Name	Type	Description	Default
`data`	`DataFrame`	Training sample with continuous `treatment` and outcome `y`.	required
`y`	`str`		required
`treatment`	`str`		required
`covariates`	`sequence of str`		required
`test_data`	`DataFrame`	Test rows to predict on.	required
`dose_grid`	`sequence of float`	If provided, return conformal bands for the entire curve `E[Y \| T=t, X]` at each `t` on the grid (one curve per test row).	`None`
`alpha`	`float`	Target miscoverage rate (`1 - alpha` is the nominal coverage).	`0.1`
`estimator`	`sklearn-style regressor`	Model for `E[Y \| T, X]`. Must accept `fit(X, y)` and `predict(X)`. Defaults to a gradient-boosting regressor.	`None`
`calibration_frac`	`float`		`0.5`
`random_state`	`int`		`0`

Returns:

Type	Description
`ContinuousConformalResult`

Examples:

>>> import statspai as sp
>>> import numpy as np
>>> import pandas as pd
>>> rng = np.random.default_rng(0)
>>> n = 300
>>> t = rng.uniform(0, 5, n)            # continuous treatment / dose
>>> x = rng.normal(size=n)
>>> y = 2.0 + 0.7 * t + 0.5 * x + rng.normal(0, 0.5, n)
>>> train = pd.DataFrame({'y': y, 't': t, 'x': x})
>>> test = pd.DataFrame({'t': [1.0, 3.0], 'x': [0.0, 0.5]})
>>> res = sp.conformal_continuous(
...     train, y='y', treatment='t', covariates=['x'],
...     test_data=test, alpha=0.1, random_state=0,
... )
>>> list(res.predictions.columns)
['prediction', 'lo', 'hi']
>>> bool((res.predictions['hi'] > res.predictions['lo']).all())
True

References

Schröder et al. (arXiv:2407.03094, 2024).

conformal_interference ¶

conformal_interference(data: DataFrame, *, y: str, treatment: str, cluster: str, covariates: Sequence[str], test_clusters: Sequence, alpha: float = 0.1, estimator: Optional[Any] = None, calibration_frac: float = 0.5, random_state: int = 0) -> InterferenceConformalResult

Cluster-exchangeable split-conformal prediction under interference.

When units within a cluster interfere (spillover, networks) but clusters are exchangeable, the exchangeable-data guarantee of split conformal survives at the cluster level. We compute a cluster-level absolute-residual score by averaging per-unit residuals inside each cluster, then build the usual split-conformal quantile over clusters.

Parameters:

Name	Type	Description	Default
`data`	`DataFrame`	Full sample with cluster identifier `cluster`.	required
`y`	`str`		required
`treatment`	`str`		required
`cluster`	`str`		required
`covariates`	`sequence of str`		required
`test_clusters`	`sequence`	Cluster IDs to predict on. Must appear in `data`.	required
`alpha`	`float`		`0.1`
`estimator`	`sklearn-style regressor`		`None`
`calibration_frac`	`float`		`0.5`
`random_state`	`int`		`0`

Returns:

Type	Description
`InterferenceConformalResult`

Examples:

>>> import statspai as sp
>>> import numpy as np
>>> import pandas as pd
>>> rng = np.random.default_rng(0)
>>> rows = []   # 12 clusters; units within a cluster interfere
>>> for c in range(12):
...     cluster_shock = rng.normal(0, 0.5)
...     for _ in range(rng.integers(5, 12)):
...         treat = rng.integers(0, 2)
...         x = rng.normal()
...         y = (1.0 + 0.6 * treat + 0.4 * x + cluster_shock
...              + rng.normal(0, 0.3))
...         rows.append({'cluster': c, 'treat': treat, 'x': x, 'y': y})
>>> df = pd.DataFrame(rows)
>>> res = sp.conformal_interference(
...     df, y='y', treatment='treat', cluster='cluster',
...     covariates=['x'], test_clusters=[0, 1], alpha=0.1, random_state=0,
... )
>>> list(res.predictions['cluster'])
[0, 1]
>>> bool((res.predictions['hi'] > res.predictions['lo']).all())
True

Notes

This is the cluster-exchangeable variant used in the Memmesheimer, Heuveline & Hesser (arXiv:2509.21660, 2025) systematic review as the recommended default when SUTVA is violated within clusters.

conformal_regression ¶

conformal_regression(formula: str, data: DataFrame, newdata: Optional[DataFrame] = None, method: str = 'jackknife+', alpha: float = 0.1, train_frac: float = 0.5, seed: Optional[int] = None) -> DataFrame

Distribution-free prediction intervals for a least-squares regression.

Returns, for each row of newdata, the least-squares prediction and an interval that contains the outcome of a new, exchangeable observation with probability at least 1 - alpha, without assuming normal or homoskedastic errors or a correctly specified mean.

Parameters:

Name	Type	Description	Default
`formula`	`str`	Regression formula, e.g. `"medv ~ rm + lstat"`.	required
`data`	`DataFrame`	The sample the model is fitted on.	required
`newdata`	`DataFrame`	Rows to predict. When omitted, each row of `data` is predicted from all the others (its own outcome plays no part in its interval), which is the way to check coverage on a sample.	`None`
`method`	`('jackknife+', 'split', 'full')`	`"split"` fits on a random `train_frac` of the sample and calibrates on the rest: the cheapest, exact `1 - alpha` coverage, wider intervals because each step sees part of the data. `"jackknife+"` uses every observation for both: coverage at least `1 - 2 alpha` in the worst case and close to `1 - alpha` in practice. `"full"` is the exact conformal set for least squares: coverage at least `1 - alpha`, the most work per new point (it grows with the square of the sample size).	`"jackknife+"`
`alpha`	`float`	One minus the target coverage.	`0.1`
`train_frac`	`float`	Share of the sample used for fitting under `method="split"`.	`0.5`
`seed`	`int`	Seed of the random split (`method="split"` only).	`None`

Returns:

Type	Description
`DataFrame`	Columns `yhat`, `lower`, `upper`, indexed like `newdata` (or `data`). `attrs` records `method`, `alpha` and `n`. An interval end is infinite when the sample is too small for the requested level.

Notes

The guarantee is marginal: it holds on average over new points, not at each value of the regressors. With heteroskedastic errors the intervals have the right coverage overall and are too wide where the noise is small and too narrow where it is large. They cover an outcome, not a mean or a coefficient, and say nothing about causal effects; for those see sp.conformal("cate", ...).

Exchangeability fails for time series and for clustered samples whose new point comes from a new cluster.

Examples:

>>> import numpy as np, pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> x = rng.uniform(-2, 2, size=400)
>>> df = pd.DataFrame({"x": x, "y": 1 + x + rng.standard_t(3, size=400)})
>>> new = pd.DataFrame({"x": [0.0, 1.0]})
>>> out = sp.conformal_regression("y ~ x", df, new, alpha=0.1)
>>> list(out.columns)
['yhat', 'lower', 'upper']
>>> bool((out["lower"] < out["yhat"]).all() and (out["yhat"] < out["upper"]).all())
True

References

[@lei2018distribution], [@barber2021predictive]

statspai.conformal_causal¶

conformal_causal ¶

DebiasedConformalResult dataclass ¶

ConformalDensityResult dataclass ¶

FairConformalResult dataclass ¶

ConformalCATE ¶

fit ¶

predict ¶

MultiDPConformalResult dataclass ¶

ConformalCounterfactualResult dataclass ¶

ConformalITEResult dataclass ¶

ContinuousConformalResult dataclass ¶

InterferenceConformalResult dataclass ¶

conformal_debiased_ml ¶

conformal_density_ite ¶

conformal_fair_ite ¶

conformal_cate ¶

conformal_ite_multidp ¶

conformal_counterfactual ¶

conformal_ite_interval ¶

weighted_conformal_prediction ¶

conformal_available_kinds ¶

conformal ¶

conformal_continuous ¶

conformal_interference ¶

conformal_regression ¶

`statspai.conformal_causal`¶

DebiasedConformalResult `dataclass` ¶

ConformalDensityResult `dataclass` ¶

FairConformalResult `dataclass` ¶

MultiDPConformalResult `dataclass` ¶

ConformalCounterfactualResult `dataclass` ¶

ConformalITEResult `dataclass` ¶

ContinuousConformalResult `dataclass` ¶

InterferenceConformalResult `dataclass` ¶