statspai.mediation

mediation

Causal Mediation Analysis module for StatsPAI.

Implements modern causal mediation analysis following Imai, Keele, and Tingley (2010), decomposing total treatment effects into: - Average Causal Mediation Effect (ACME): indirect effect through mediator - Average Direct Effect (ADE): direct effect not through mediator

Also supports the classical Baron-Kenny approach for comparison.

References

Imai, K., Keele, L., and Tingley, D. (2010). "A General Approach to Causal Mediation Analysis." Psychological Methods, 15(4), 309-334. [@imai2010general]

Baron, R.M. and Kenny, D.A. (1986). "The Moderator-Mediator Variable Distinction in Social Psychological Research." Journal of Personality and Social Psychology, 51(6), 1173-1182. [@baron1986moderator]

FourWayResult dataclass

Bases: ResultProtocolMixin

Result container for :func:four_way_decomposition.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> treat = np.tile([0.0, 1.0], 50)
>>> mediator = 0.5 + 0.4 * treat + np.repeat(np.linspace(-1, 1, 50), 2)
>>> y = 1 + 2 * treat + 3 * mediator + 4 * treat * mediator
>>> df = pd.DataFrame({"y": y, "a": treat, "m": mediator})
>>> res = sp.four_way_decomposition(df, y="y", treat="a", mediator="m")
>>> type(res).__name__
'FourWayResult'
>>> round(float(res.total_effect), 1)
6.8

to_dict

to_dict() -> dict

JSON-safe dict of every field (agent-native serialization).

MediationAnalysis

Causal Mediation Analysis estimator.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 300
>>> treat = rng.binomial(1, 0.5, n)
>>> x = rng.normal(0, 1, n)
>>> m = 0.4 * treat + 0.3 * x + rng.normal(0, 1, n)
>>> y = 0.5 * treat + 0.6 * m + 0.2 * x + rng.normal(0, 1, n)
>>> df = pd.DataFrame({"treat": treat, "x": x, "m": m, "y": y})
>>> ma = sp.MediationAnalysis(df, y="y", treat="treat", mediator="m",
...                           covariates=["x"], n_boot=100, seed=0)
>>> res = ma.fit()
>>> res.estimand
'ACME'
>>> res.method
'Causal Mediation Analysis'

fit

fit() -> CausalResult

Estimate mediation effects with bootstrap inference.

MediateSensitivityResult dataclass

Bases: ResultProtocolMixin

plot

plot(ax: Any = None, *, fill: bool = True, annotate: bool = True, figsize: Any = (7.0, 4.5), **kwargs: Any) -> Any

Publication-style sensitivity plot.

Shows ACME(ρ) vs the mediator-outcome confounder strength ρ, with a coloured fill for the region of nullability (any ρ in [ρ_at_zero, 1] flips the ACME sign), the ρ at which the ACME crosses zero (annotated), and reference lines at ρ=0 (i.e. sequential-ignorability) and ACME=0.

Parameters:

Name	Type	Description	Default
ax	matplotlib Axes		None
fill	bool	Fill the {ACME(ρ) > 0} region in light blue and the {ACME(ρ) < 0} region in light red, à la sensemakr.	True
annotate	bool	Annotate ρ_at_zero, baseline ACME, and an interpretive note.	True

four_way_decomposition

four_way_decomposition(data: DataFrame, y: str, treat: str, mediator: str, covariates: Optional[Sequence[str]] = None, a0: float = 0.0, a1: float = 1.0, m0: float = 0.0, vce: str = 'ols', yreg: str = 'linear', mreg: str = 'linear', covariate_values: Optional[Mapping[str, float]] = None) -> FourWayResult

Parametric four-way decomposition of TE = CDE + INT_ref + INT_med + PIE.

Parameters:

Name	Type	Description	Default
data	DataFrame		required
y	str		required
treat	str		required
mediator	str		required
covariates	sequence of str	Entered linearly in both the outcome and the mediator model.	None
a0	float	Reference and comparison levels of the treatment (default 0, 1).	0.0
a1	float	Reference and comparison levels of the treatment (default 0, 1).	0.0
m0	float	Mediator reference level at which CDE is evaluated.	0.0
vce	('ols', 'ml')	Variance used for linear models (a linear outcome, a normal mediator). 'ols' uses s^2 = RSS / (n - p) -- R CMAverse::cmest(estimation = "paramfunc"); 'ml' uses RSS / n -- Stata med4way, which fits linear models by maximum likelihood. With a normal mediator and a log-linear outcome the mediator variance enters the components, so the two choices also differ in the point estimates; its sampling variance 2 sigma^4 / dof enters the standard errors.	'ols'
yreg	('linear', 'logistic', 'poisson')	Outcome model, with the exposure-mediator interaction. For 'logistic' and 'poisson' the components are on the excess-relative-risk scale (Stata med4way's ereri_*): they sum to RR - 1. A logistic outcome is read as a risk, which needs a rare outcome (odds ratio close to the risk ratio), as in med4way and CMAverse.	'linear'
mreg	('linear', 'logistic')	Mediator model: normal linear regression, or logistic regression for a 0/1 mediator.	'linear'
covariate_values	mapping	Values of the covariates at which the decomposition is evaluated (med4way's c()); covariates not named are held at their sample means, the default for all of them.	None

Returns:

Type	Description
FourWayResult	se holds delta-method standard errors for total_effect, cde, int_ref, int_med, pie and the proportions (prop_cde ... prop_pie, prop_mediated, prop_interaction, prop_eliminated -- med4way's p_*, op_m, op_ati, op_e). detail['scale'] is 'difference' or 'excess_relative_risk'; on the latter detail['total_rr'] is the ratio itself.

Notes

Every component is a closed-form function of the two regressions (mediation/_four_way_engine.py); the Jacobian for the delta method is computed by complex-step differentiation, so it carries no truncation error. Against Stata med4way 2.3.2 (vce='ml') the estimates agree to 1e-8 and the standard errors to 1e-7 for every outcome / mediator model (tests/reference_parity/ test_four_way_nonlinear_parity.py). With vce='ols' the estimates equal CMAverse's; its standard errors leave out the mediator variance's sampling error.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> treat = np.tile([0.0, 1.0], 50)
>>> mediator = 0.5 + 0.4 * treat + np.repeat(np.linspace(-1, 1, 50), 2)
>>> y = 1 + 2 * treat + 3 * mediator + 4 * treat * mediator
>>> df = pd.DataFrame({"y": y, "a": treat, "m": mediator})
>>> res = sp.four_way_decomposition(df, y="y", treat="a", mediator="m")
>>> round(float(res.cde), 1)
2.0
>>> round(float(res.pie), 1)
1.2

mediate_interventional

mediate_interventional(data: DataFrame, y: str, treat: str, mediator: str, covariates: Optional[List[str]] = None, tv_confounders: Optional[List[str]] = None, n_mc: int = 500, n_boot: int = 500, alpha: float = 0.05, pvalue_method: str = 'bootstrap_sign', seed: int = 42) -> CausalResult

Interventional (in)direct effects (VanderWeele, Vansteelandt, Robins 2014).

Decomposes the total effect into:

• IIE (Interventional Indirect Effect): E[Y(1, G_{M|1})] - E[Y(1, G_{M|0})], i.e. the effect of shifting M's post-treatment distribution from its D=0 draw to its D=1 draw, while holding D fixed at 1.
• IDE (Interventional Direct Effect): E[Y(1, G_{M|0})] - E[Y(0, G_{M|0})].
• Total = IIE + IDE = E[Y(1, G_{M|1})] - E[Y(0, G_{M|0})].

Here :math:G_{M|d} is the random draw from the marginal distribution of :math:M under treatment :math:D = d (integrated over covariates).

Interventional effects are identified under the standard mediation assumptions minus the cross-world independence requirement — which makes them valid even when there are treatment-induced mediator-outcome confounders (tv_confounders). Natural effects are not generally identified in that case.

Parameters:

Name	Type	Description	Default
data	DataFrame		required
y	str	Outcome variable.	required
treat	str	Binary treatment (0/1).	required
mediator	str	Mediator variable.	required
covariates	list of str	Baseline (pre-treatment) covariates.	None
tv_confounders	list of str	Treatment-induced mediator-outcome confounders (variables affected by D that confound the M-Y relationship). These enter the outcome model but not the M-marginalization; each is modelled linearly on [1, D, covariates] and integrated over its law under the intervened treatment, so the D -> L -> Y path is part of the direct effect.	None
n_mc	int	Monte Carlo draws of M for the stochastic intervention.	500
n_boot	int	Nonparametric bootstrap replications.	500
alpha	float	Significance level.	0.05
pvalue_method	(bootstrap_sign, wald)	How the per-effect p-value is computed: 'bootstrap_sign' (default, matches :func:sp.mediate): the bootstrap-CI-inversion p-value, i.e. twice the fraction of bootstrap replications on the opposite side of zero from the point estimate. Respects skewed bootstrap distributions. 'wald': the conventional 2*(1-Φ(|θ̂/ŝe|)) p-value, using the bootstrap SE. Consistent with the Wald pvalue convention used across the rest of StatsPAI's causal-inference surface (e.g. sp.aipw, sp.dml).	'bootstrap_sign'
seed	int	Random seed.	42

Returns:

Type	Description
CausalResult	estimate is IIE; full decomposition lives in detail.

Examples:

>>> import numpy as np
>>> import pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(42)
>>> n = 300
>>> treat = rng.binomial(1, 0.5, n)
>>> x = rng.normal(0, 1, n)
>>> m = 0.4 * treat + 0.3 * x + rng.normal(0, 1, n)
>>> y = 0.5 * treat + 0.6 * m + 0.2 * x + rng.normal(0, 1, n)
>>> df = pd.DataFrame({"treat": treat, "x": x, "m": m, "y": y})
>>> res = sp.mediate_interventional(
...     df, y="y", treat="treat", mediator="m", covariates=["x"],
...     n_mc=100, n_boot=100, seed=0,
... )
>>> res.estimand
'IIE'
>>> res.method
'Interventional Mediation Analysis'
>>> len(res.detail["effect"])     # IIE, IDE, Total
3

Notes

Linear outcome model assumption. The current implementation hard-codes an OLS outcome regression Y ~ D + M + X_base + X_tv. This permits the Monte-Carlo integration to collapse analytically via the linearity of OLS predictions in the treatment-induced confounder block (X_tv), giving an O(n_mc + n) cost instead of O(n × n_mc). Non-linear outcome models (gradient boosting, neural nets, etc.) would break this vectorisation and are not currently supported — passing a custom learner is not exposed via the API.

References

VanderWeele, T.J., Vansteelandt, S. and Robins, J.M. (2014). "Effect decomposition in the presence of an exposure-induced mediator-outcome confounder." Epidemiology, 25(2), 300-306. [@vanderweele2014effect]

mediate_paths

mediate_paths(data: DataFrame, y: str, treat: str, mediators: Sequence[Union[str, Sequence[str]]], covariates: Optional[List[str]] = None, interaction: bool = True, outcome_model: str = 'linear', ps: bool = False, decomposition: str = 'type1', n_boot: int = 500, alpha: float = 0.05, seed: int = 42) -> CausalResult

Path-specific effects of a binary treatment through ordered mediators.

Parameters:

Name	Type	Description	Default
data	DataFrame		required
y	str	Outcome.	required
treat	str	Binary treatment (0/1).	required
mediators	sequence	The mediators in causal order, M_1 first. Each entry is a column name or a list of columns that form one mediator block (R paths' m = list(...)). A mediator may affect every later one; a later mediator may not affect an earlier one.	required
covariates	list of str	Pre-treatment confounders, entered linearly in every model.	None
interaction	bool	Include treatment x mediator interactions in the outcome models.	True
outcome_model	('linear', 'logistic')	Model for E[Y | A, X, M_1..M_j]; 'logistic' for a 0/1 outcome (effects on the risk-difference scale).	'linear'
ps	bool	Reweight with a logistic propensity score of the treatment on the covariates (the imputation-based weighting estimator) instead of the pure imputation estimator.	False
decomposition	('type1', 'type2')	Which decomposition is the headline (both are in detail). Type I switches the mediators from the treated to the control value one at a time starting from M_1; for one mediator it is pure direct + total indirect effect.	'type1'
n_boot	int	Nonparametric bootstrap replications (units resampled, every model refitted) for the standard errors and percentile intervals.	500
alpha	float		0.05
seed	int		42

Returns:

Type	Description
CausalResult	estimate is the total effect. detail has one row per decomposition and component (direct, via_M1 ... via_MK, total) with bootstrap SE, percentile interval and the bootstrap sign p-value. model_info['type1'] / ['type2'] hold the point estimates.

Notes

Identification needs no unmeasured confounding of treatment-outcome, treatment-mediator and mediator-outcome relations given the covariates and earlier mediators, and no other post-treatment confounder of the mediators (an exposure-induced confounder has to be entered as an earlier mediator). The direct and path-specific effects always sum to the total effect.

References

[@zhou2023tracing]

Examples:

>>> import numpy as np, pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> n = 2000
>>> a = rng.integers(0, 2, n)
>>> m1 = 0.6 * a + rng.normal(size=n)
>>> m2 = 0.5 * a + 0.7 * m1 + rng.normal(size=n)
>>> y = 0.4 * a + 0.5 * m1 + 0.8 * m2 + rng.normal(size=n)
>>> df = pd.DataFrame({"y": y, "a": a, "m1": m1, "m2": m2})
>>> r = sp.mediate_paths(df, y="y", treat="a", mediators=["m1", "m2"],
...                      interaction=False, n_boot=50)
>>> sorted(r.model_info["type1"])
['direct', 'via_M1', 'via_M2']

mediate_sensitivity

mediate_sensitivity(data: DataFrame, y: str, treat: str, mediator: str, covariates: Optional[list] = None, rho_range: tuple = (-0.9, 0.9), n_grid: int = 41, *, conf_level: float = 0.95, eps: Optional[float] = None, max_iter: int = 10000) -> MediateSensitivityResult

Sensitivity analysis for causal mediation (linear models).

For each rho on a grid, the mediator model M ~ T + X and the outcome model Y ~ T + M + X are re-estimated jointly with their error correlation fixed at rho, and the ACME a_T * b_M is reported with a delta-method confidence interval (Imai, Keele & Yamamoto 2010, Section 5). rho = 0 is sequential ignorability.

Parameters:

Name	Type	Description	Default
data	DataFrame	Analysis data; rows with a missing value in a used column are dropped.	required
y	str	Outcome column (linear model).	required
treat	str	Treatment column. With a 0/1 treatment the ACME is the effect of moving from control to treatment; with a numeric treatment it is the effect per unit.	required
mediator	str	Mediator column (linear model).	required
covariates	list of str	Pre-treatment covariates entering both models.	None
rho_range	tuple of float	(low, high) with -1 < low <= high < 1.	(-0.9, 0.9)
n_grid	int	Number of equally spaced rho values (numpy.linspace).	41
conf_level	float	Level of the normal-approximation interval at each rho.	0.95
eps	float	Stopping tolerance of the iterated FGLS: iteration stops once the sum of squared changes of the coefficients is below eps (mediation::medsens uses sqrt(.Machine$double.eps)). The default None iterates to the fixed point.	None
max_iter	int	Maximum FGLS steps per rho; a ConvergenceWarning names the rho values that reach it.	10000

Returns:

Type

Description

MediateSensitivityResult

rho_grid, acme_at_rho, se_at_rho, ci_lower / ci_upper; acme_at_zero (least squares, rho = 0); rho_at_zero, the exact rho at which the ACME is zero (Corr(resid(M ~ T + X), resid(Y ~ T + X))); rho_grid_zero, the grid value with the smallest |ACME|; r2_mediator / r2_outcome (R^2 of the two models); r2star_threshold = rho_at_zero^2 and r2tilde_threshold = rho_at_zero^2 (1 - R^2_M)(1 - R^2_Y), the products of the partial and total R^2 an unobserved confounder needs to remove the ACME.

Examples:

>>> import numpy as np, pandas as pd
>>> import statspai as sp
>>> rng = np.random.default_rng(0)
>>> t = rng.binomial(1, 0.5, 400)
>>> m = 0.5 * t + rng.normal(size=400)
>>> df = pd.DataFrame({"t": t, "m": m,
...                    "y": t + 0.8 * m + rng.normal(size=400)})
>>> res = sp.mediate_sensitivity(df, y="y", treat="t", mediator="m",
...                              n_grid=19)
>>> float(round(res.rho_grid[0], 2)), res.rho_at_zero is not None
(-0.9, True)

References

[@imai2010identification], [@imai2010general]

Clean-room implementation; the reference (R mediation 4.5.1, medsens, GPL >= 2) was used only as a black box.