import numpy as np
import pandas as pd
from scipy.special import expit
from scipy.stats import norm
from sklearn.linear_model import LinearRegression, LogisticRegression
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import SplineTransformer, StandardScaler
def as_array(x):
return np.asarray(x, dtype=float).reshape(-1)
def lm_features(*cols):
return np.column_stack([as_array(col) for col in cols])
def linear_fit(x, y):
model = LinearRegression()
model.fit(np.asarray(x, dtype=float), as_array(y))
return model
def logit_fit(x, y):
model = make_pipeline(
StandardScaler(),
LogisticRegression(C=1e6, max_iter=4000, solver="lbfgs"),
)
model.fit(np.asarray(x, dtype=float), as_array(y).astype(int))
return model
def predict_prob(model, x):
return model.predict_proba(np.asarray(x, dtype=float))[:, 1]
def clip_prob(p, eps=0.025):
return np.clip(as_array(p), eps, 1.0 - eps)Mediation estimands, estimators, and Python code
This note is a compact Python companion to the SER 2026 mediation workshop. It separates three layers that are easy to blur:
- the causal estimand, written in counterfactual notation;
- the statistical estimand, written as a functional of the observed data law;
- the estimator, written as runnable Python code.
The notebook is organized as literate programming. Each estimand section puts the executable cells next to the algebra they implement: first the section’s toy data, then the nuisance functions named in the formula, then the plug-in or one-step summands, and finally the sample average.
Download the self-contained Quarto source.
1 Minimal Python setup
The setup chunk only defines imports and small modeling utilities. It does not hide any mediation estimator. The actual estimators live in the sections where their formulas are introduced.
2 Observed data structure
The basic cross-sectional mediation setup is
\[ O = (W, A, M, Y), \]
where \(W\) are baseline covariates, \(A\) is treatment or exposure, \(M\) is the mediator, and \(Y\) is the outcome. When there is an intermediate mediator-outcome confounder affected by treatment, write
\[ O = (W, A, Z, M, Y). \]
The structural story used throughout the simple examples is:
\[ W \rightarrow A,\quad W \rightarrow M,\quad W \rightarrow Y,\quad A \rightarrow M,\quad A \rightarrow Y,\quad M \rightarrow Y. \]
For the natural direct and indirect effects, identification typically requires:
- no unmeasured treatment-outcome confounding conditional on \(W\);
- no unmeasured mediator-outcome confounding conditional on \((A,W)\);
- no unmeasured treatment-mediator confounding conditional on \(W\);
- the cross-world independence condition needed for natural effects;
- positivity for the treatment and mediator mechanisms.
3 Estimand map
| Target | Counterfactual estimand | Identification idea |
|---|---|---|
| Controlled direct effect | \(\mathbb{E}[Y_{1,m} - Y_{0,m}]\) | Compare outcome regressions at fixed mediator value \(m\). |
| Natural direct effect | \(\mathbb{E}[Y_{1,M_0} - Y_{0,M_0}]\) | Average controlled direct effects over the mediator distribution under control. |
| Natural indirect effect | \(\mathbb{E}[Y_{1,M_1} - Y_{1,M_0}]\) | Compare mediator distributions under treatment and control while fixing treatment. |
| Randomized interventional mean | \(\mathbb{E}[Y_{a,G_{a'}}]\) | Replace the mediator draw by a stochastic draw from its distribution under \(A=a'\). |
| Stochastic direct effect | \(\mathbb{E}[Y_{A_\delta,M} - Y_{A,M}]\) | Change the treatment distribution while holding the observed mediator fixed. |
| Stochastic indirect effect | \(\mathbb{E}[Y_{A_\delta,M_{A_\delta}} - Y_{A_\delta,M}]\) | Compare the mediator under the stochastic intervention to the observed mediator. |
4 Product-of-coefficients is not the estimand
The classical linear mediation heuristic multiplies the coefficient of \(A\) in a mediator model by the coefficient of \(M\) in an outcome model. This can be a useful descriptive regression summary, but it is not generally a causal mediation estimand.
The workshop’s counterexample has an intermediate variable \(Z\) on the pathway:
\[ A \rightarrow Z \rightarrow M \rightarrow Y. \]
The mediated pathway is present, but the product-of-coefficients regression summary is approximately zero after conditioning on \(Z\).
First simulate data from the counterexample.
poc_n = 100_000
poc_rng = np.random.default_rng(1)
poc_w = poc_rng.normal(size=poc_n)
poc_a = poc_rng.binomial(1, 0.5, size=poc_n)
poc_z = poc_rng.binomial(1, 0.2 * poc_a + 0.3)
poc_m = poc_rng.normal(poc_w + poc_z, 1.0)
poc_y = poc_rng.normal(poc_m + poc_w - poc_a + poc_z, 1.0)The product heuristic is literally the product of two fitted slopes:
\[ \widehat{\mathrm{product}} = \hat\beta_M \hat\alpha_A, \]
where \(\hat\beta_M\) is the coefficient of \(M\) in the outcome regression and \(\hat\alpha_A\) is the coefficient of \(A\) in the mediator regression.
poc_y_fit = linear_fit(lm_features(poc_m, poc_a, poc_w, poc_z), poc_y)
poc_m_fit = linear_fit(lm_features(poc_a, poc_w, poc_z), poc_m)
poc_beta_M = poc_y_fit.coef_[0]
poc_alpha_A = poc_m_fit.coef_[0]
poc_product = poc_beta_M * poc_alpha_A
pd.Series(
{
"beta_M_from_outcome_model": poc_beta_M,
"alpha_A_from_mediator_model": poc_alpha_A,
"product_of_coefficients": poc_product,
}
)beta_M_from_outcome_model 0.999192
alpha_A_from_mediator_model 0.000616
product_of_coefficients 0.000615
dtype: float645 Controlled direct effect
5.1 Estimand
The controlled direct effect at mediator value \(m\) is
\[ \psi_{\mathrm{CDE}}(m) = \mathbb{E}[Y_{1,m} - Y_{0,m}]. \]
Under the usual adjustment conditions, it is identified by
\[ \psi_{\mathrm{CDE}}(m) = \mathbb{E}\{b(1,m,W) - b(0,m,W)\}, \]
where
\[ b(a,m,w) = \mathbb{E}[Y \mid A=a, M=m, W=w]. \]
5.2 Data for this section
The simple data-generating process makes the CDE transparent: \(Y\) is linear in \(W\), \(A\), and \(M\).
cde_n = 100_000
cde_m_value = 0.0
cde_rng = np.random.default_rng(2)
cde_w = cde_rng.normal(size=cde_n)
cde_a = cde_rng.binomial(1, 0.5, size=cde_n)
cde_m = cde_rng.normal(cde_w + cde_a, 1.0)
cde_y = cde_rng.normal(cde_w + cde_a + cde_m, 1.0)5.3 Estimate \(b(a,m,w)\)
Fit the outcome regression named in the identifying formula:
\[ \hat b(a,m,w) \approx \mathbb{E}[Y \mid A=a,M=m,W=w]. \]
cde_b_hat = linear_fit(lm_features(cde_m, cde_a, cde_w), cde_y)5.4 Evaluate the two terms in the summand
For each observed \(W_i\), the formula needs both \(\hat b(1,m,W_i)\) and \(\hat b(0,m,W_i)\) at the same fixed mediator value \(m\).
cde_b_hat_1mW = cde_b_hat.predict(
lm_features(np.full(cde_n, cde_m_value), np.ones(cde_n), cde_w)
)
cde_b_hat_0mW = cde_b_hat.predict(
lm_features(np.full(cde_n, cde_m_value), np.zeros(cde_n), cde_w)
)
cde_summand = cde_b_hat_1mW - cde_b_hat_0mW5.5 Average the summand
The estimator is the empirical analogue of the outer expectation:
\[ \hat\psi_{\mathrm{CDE}}(m) = \frac{1}{n}\sum_{i=1}^n \{\hat b(1,m,W_i)-\hat b(0,m,W_i)\}. \]
cde_estimate = np.mean(cde_summand)
float(cde_estimate)0.98858240720617446 Natural direct effect by g-computation
6.1 Estimand
The natural direct effect is
\[ \psi_{\mathrm{NDE}} = \mathbb{E}[Y_{1,M_0} - Y_{0,M_0}]. \]
With the natural-effect assumptions, the identifying functional is
\[ \psi_{\mathrm{NDE}} = \mathbb{E}\left[ \mathbb{E}\{b(1,M,W)-b(0,M,W) \mid A=0,W\} \right]. \]
Define
\[ Q(M,W) = b(1,M,W) - b(0,M,W), \]
and
\[ r(a,w) = \mathbb{E}[Q(M,W) \mid A=a,W=w]. \]
Then
\[ \psi_{\mathrm{NDE}} = \mathbb{E}[r(0,W)]. \]
6.2 Workshop simulation for this section
The estimation chapter’s simulation deliberately creates misspecification for simple linear g-computation. The true outcome regression has an \(A \times M\) interaction and nonlinear dependence on \(W\).
def nde_mean_y(m, a, w):
return np.abs(as_array(w)) + as_array(a) * as_array(m)
def nde_mean_m(a, w):
return expit(as_array(w) ** 2 - as_array(a))
def nde_pscore(w):
return expit(1.0 - np.abs(as_array(w)))
def simulate_nde_data(n=1000, seed=None):
rng = np.random.default_rng(seed)
w = rng.uniform(-1.0, 1.0, n)
a = rng.binomial(1, nde_pscore(w))
m = rng.binomial(1, nde_mean_m(a, w))
y = rng.normal(nde_mean_y(m, a, w), 1.0)
return pd.DataFrame({"y": y, "m": m, "a": a, "w": w})
def true_nde(grid_size=200_000):
w_grid = np.linspace(-1.0, 1.0, grid_size)
q_m1 = nde_mean_y(1, 1, w_grid) - nde_mean_y(1, 0, w_grid)
q_m0 = nde_mean_y(0, 1, w_grid) - nde_mean_y(0, 0, w_grid)
return float(np.mean(q_m1 * nde_mean_m(0, w_grid) + q_m0 * (1 - nde_mean_m(0, w_grid))))
nde_data = simulate_nde_data(n=1000, seed=4235243)
nde_y = nde_data["y"]
nde_m = nde_data["m"]
nde_a = nde_data["a"]
nde_w = nde_data["w"]
nde_n = len(nde_data)
nde_truth = true_nde()6.3 Estimate \(b(a,m,w)\)
The first plug-in nuisance is the outcome regression:
\[ \hat b(a,m,w) \approx \mathbb{E}[Y \mid A=a,M=m,W=w]. \]
nde_b_hat = linear_fit(lm_features(nde_m, nde_a, nde_w), nde_y)6.4 Construct \(\hat Q(M_i,W_i)\)
The inner contrast
\[ Q(M_i,W_i)=b(1,M_i,W_i)-b(0,M_i,W_i) \]
becomes the observed-sample vector below.
nde_b_hat_1MW = nde_b_hat.predict(lm_features(nde_m, np.ones(nde_n), nde_w))
nde_b_hat_0MW = nde_b_hat.predict(lm_features(nde_m, np.zeros(nde_n), nde_w))
nde_Q_hat = nde_b_hat_1MW - nde_b_hat_0MW6.5 Estimate \(r(a,w)\) and average \(r(0,W_i)\)
The second-stage regression estimates
\[ r(a,w) = \mathbb{E}[Q(M,W) \mid A=a,W=w]. \]
The final estimator averages \(\hat r(0,W_i)\).
nde_r_hat = linear_fit(lm_features(nde_a, nde_w), nde_Q_hat)
nde_r_hat_0W = nde_r_hat.predict(lm_features(np.zeros(nde_n), nde_w))
nde_gcomp_estimate = np.mean(nde_r_hat_0W)
pd.Series(
{
"truth": nde_truth,
"gcomp_estimate": nde_gcomp_estimate,
"gcomp_bias": nde_gcomp_estimate - nde_truth,
}
)truth 0.580534
gcomp_estimate 0.383509
gcomp_bias -0.197025
dtype: float646.6 Reusable version for simulation
The function below is just the same four algebraic steps packaged so that the Monte Carlo comparison can call it repeatedly.
def gcomp_nde(y, m, a, w):
y = as_array(y)
m = as_array(m)
a = as_array(a)
w = as_array(w)
n = len(y)
b_hat = linear_fit(lm_features(m, a, w), y)
b_hat_1MW = b_hat.predict(lm_features(m, np.ones(n), w))
b_hat_0MW = b_hat.predict(lm_features(m, np.zeros(n), w))
Q_hat = b_hat_1MW - b_hat_0MW
r_hat = linear_fit(lm_features(a, w), Q_hat)
r_hat_0W = r_hat.predict(lm_features(np.zeros(n), w))
return float(np.mean(r_hat_0W))7 One-step estimator for the natural direct effect
7.1 Estimand
The target is still
\[ \psi_{\mathrm{NDE}} = \mathbb{E}[r(0,W)]. \]
The one-step estimator augments the plug-in regression with inverse-probability and residual correction terms from the efficient influence function.
Use the nuisance functions
\[ g(a \mid w) = \mathbb{P}(A=a \mid W=w), \]
\[ e(a \mid m,w) = \mathbb{P}(A=a \mid M=m,W=w), \]
\[ b(a,m,w) = \mathbb{E}[Y \mid A=a,M=m,W=w], \]
and
\[ r(w) = \mathbb{E}[Q(M,W) \mid A=0,W=w]. \]
7.2 Estimator
Using the Bayes-rule reparameterization of the mediator density ratio, the uncentered EIF-style estimating function is
\[ \begin{aligned} \phi(O_i;\hat\eta) = &\left[ \frac{A_i}{\hat g(0\mid W_i)} \frac{\hat e(0\mid M_i,W_i)}{\hat e(1\mid M_i,W_i)} - \frac{1-A_i}{\hat g(0\mid W_i)} \right] \{Y_i-\hat b(A_i,M_i,W_i)\} \\ &+ \frac{1-A_i}{\hat g(0\mid W_i)} \{\hat Q(M_i,W_i)-\hat r(W_i)\} + \hat r(W_i). \end{aligned} \]
The one-step estimator is
\[ \hat\psi_{\mathrm{NDE,1step}} = \frac{1}{n}\sum_{i=1}^n \phi(O_i;\hat\eta). \]
7.3 Basis expansions for the nuisance regressions
The code uses spline-expanded linear and logistic regressions to mimic the workshop’s illustrative mgcv implementation. The feature builders define only regression design matrices; they are not estimators by themselves.
def fit_spline(w, n_knots=6):
transformer = SplineTransformer(
n_knots=n_knots,
degree=3,
include_bias=False,
extrapolation="continue",
)
transformer.fit(as_array(w).reshape(-1, 1))
return transformer
def spline_features(transformer, w):
return transformer.transform(as_array(w).reshape(-1, 1))
def b_features(a, m, w, spline):
a = as_array(a)
m = as_array(m)
w = as_array(w)
s = spline_features(spline, w)
return np.column_stack(
[
a,
m,
w,
np.abs(w),
a * m,
s,
s * a[:, None],
s * m[:, None],
]
)
def e_features(m, w, spline):
m = as_array(m)
w = as_array(w)
s = spline_features(spline, w)
return np.column_stack([m, w, np.abs(w), m * w, s, s * m[:, None]])
def g_features(w, spline):
w = as_array(w)
s = spline_features(spline, w)
return np.column_stack([w, np.abs(w), s])
def r_features(a, w, spline):
a = as_array(a)
w = as_array(w)
s = spline_features(spline, w)
return np.column_stack([a, w, np.abs(w), s, s * a[:, None]])7.4 Estimate the nuisance functions
These three fitted objects correspond directly to \(\hat b\), \(\hat e\), and \(\hat g\) in the displayed EIF.
nde_spline = fit_spline(nde_w)
os_b_hat = linear_fit(b_features(nde_a, nde_m, nde_w, nde_spline), nde_y)
os_e_hat = logit_fit(e_features(nde_m, nde_w, nde_spline), nde_a)
os_g_hat = logit_fit(g_features(nde_w, nde_spline), nde_a)Convert the probability models into the terms that appear in the denominators.
os_g1_hat = clip_prob(predict_prob(os_g_hat, g_features(nde_w, nde_spline)))
os_g0_hat = clip_prob(1.0 - os_g1_hat)
os_e1_hat = clip_prob(predict_prob(os_e_hat, e_features(nde_m, nde_w, nde_spline)))
os_e0_hat = clip_prob(1.0 - os_e1_hat)7.5 Construct \(\hat Q\) and \(\hat r(W_i)\)
The plug-in part still needs
\[ \hat Q(M_i,W_i)=\hat b(1,M_i,W_i)-\hat b(0,M_i,W_i) \]
and
\[ \hat r(W_i)=\hat{\mathbb{E}}\{\hat Q(M,W)\mid A=0,W=W_i\}. \]
os_b_hat_AMW = os_b_hat.predict(b_features(nde_a, nde_m, nde_w, nde_spline))
os_b_hat_1MW = os_b_hat.predict(b_features(np.ones(nde_n), nde_m, nde_w, nde_spline))
os_b_hat_0MW = os_b_hat.predict(b_features(np.zeros(nde_n), nde_m, nde_w, nde_spline))
os_Q_hat = os_b_hat_1MW - os_b_hat_0MW
os_r_hat = linear_fit(r_features(nde_a, nde_w, nde_spline), os_Q_hat)
os_r_hat_W = os_r_hat.predict(r_features(np.zeros(nde_n), nde_w, nde_spline))7.6 Translate the EIF line by line
The first line of \(\phi\) is the inverse-probability weighted outcome residual.
os_outcome_residual_weight = (
nde_a / os_g0_hat * os_e0_hat / os_e1_hat
- (1.0 - nde_a) / os_g0_hat
)
os_outcome_residual_term = os_outcome_residual_weight * (nde_y - os_b_hat_AMW)The second line residualizes \(\hat Q\) around \(\hat r(W)\) among control units.
os_mediator_residual_term = (1.0 - nde_a) / os_g0_hat * (os_Q_hat - os_r_hat_W)The final line adds back the plug-in component \(\hat r(W_i)\).
os_phi_hat = os_outcome_residual_term + os_mediator_residual_term + os_r_hat_W
one_step_estimate = np.mean(os_phi_hat)
pd.Series(
{
"truth": nde_truth,
"gcomp_estimate": nde_gcomp_estimate,
"one_step_estimate": one_step_estimate,
"gcomp_abs_error": abs(nde_gcomp_estimate - nde_truth),
"one_step_abs_error": abs(one_step_estimate - nde_truth),
}
)truth 0.580534
gcomp_estimate 0.383509
one_step_estimate 0.575216
gcomp_abs_error 0.197025
one_step_abs_error 0.005318
dtype: float647.7 Reusable version for simulation
Again, the function below is the same displayed calculation packaged for repeated sampling.
def one_step_nde(y, m, a, w, clip=0.025):
y = as_array(y)
m = as_array(m)
a = as_array(a)
w = as_array(w)
n = len(y)
spline = fit_spline(w)
b_hat = linear_fit(b_features(a, m, w, spline), y)
e_hat = logit_fit(e_features(m, w, spline), a)
g_hat = logit_fit(g_features(w, spline), a)
g1_hat = clip_prob(predict_prob(g_hat, g_features(w, spline)), clip)
g0_hat = clip_prob(1.0 - g1_hat, clip)
e1_hat = clip_prob(predict_prob(e_hat, e_features(m, w, spline)), clip)
e0_hat = clip_prob(1.0 - e1_hat, clip)
b_hat_AMW = b_hat.predict(b_features(a, m, w, spline))
b_hat_1MW = b_hat.predict(b_features(np.ones(n), m, w, spline))
b_hat_0MW = b_hat.predict(b_features(np.zeros(n), m, w, spline))
Q_hat = b_hat_1MW - b_hat_0MW
r_hat = linear_fit(r_features(a, w, spline), Q_hat)
r_hat_W = r_hat.predict(r_features(np.zeros(n), w, spline))
outcome_residual_weight = a / g0_hat * e0_hat / e1_hat - (1.0 - a) / g0_hat
outcome_residual_term = outcome_residual_weight * (y - b_hat_AMW)
mediator_residual_term = (1.0 - a) / g0_hat * (Q_hat - r_hat_W)
phi_hat = outcome_residual_term + mediator_residual_term + r_hat_W
return float(np.mean(phi_hat))7.8 Small repeated-sampling comparison
def run_estimator_simulation(estimator, reps=50, n=1000, seed=123):
rng = np.random.default_rng(seed)
estimates = []
for _ in range(reps):
data = simulate_nde_data(n=n, seed=int(rng.integers(0, 2**31 - 1)))
estimates.append(estimator(data["y"], data["m"], data["a"], data["w"]))
return pd.DataFrame({"rep": np.arange(reps), "estimate": estimates})
def summarize_simulation(estimates, truth):
est = estimates["estimate"].to_numpy()
se = np.std(est, ddof=1)
lower = est - norm.ppf(0.975) * se
upper = est + norm.ppf(0.975) * se
return pd.Series(
{
"truth": truth,
"mean": float(np.mean(est)),
"bias": float(np.mean(est) - truth),
"sd": float(se),
"coverage": float(np.mean((lower < truth) & (truth < upper))),
}
)
gcomp_sims = run_estimator_simulation(gcomp_nde, reps=50, n=1000, seed=1)
one_step_sims = run_estimator_simulation(one_step_nde, reps=50, n=1000, seed=2)
pd.DataFrame(
{
"gcomp": summarize_simulation(gcomp_sims, nde_truth),
"one_step": summarize_simulation(one_step_sims, nde_truth),
}
)| gcomp | one_step | |
|---|---|---|
| truth | 0.580534 | 0.580534 |
| mean | 0.443133 | 0.570955 |
| bias | -0.137402 | -0.009580 |
| sd | 0.082041 | 0.073914 |
| coverage | 0.620000 | 0.940000 |
8 Randomized interventional effects
8.1 Estimand
When there is a post-treatment mediator-outcome confounder \(Z\), natural effects can fail to be identified. A randomized interventional estimand uses a draw \(G_{a'}\) from the mediator distribution under treatment value \(a'\).
The mean counterfactual is
\[ \psi(a,a') = \mathbb{E}[Y_{a,G_{a'}}]. \]
The decomposition is
\[ \mathbb{E}[Y_{1,G_1} - Y_{0,G_0}] = \{\mathbb{E}[Y_{1,G_1}] - \mathbb{E}[Y_{1,G_0}]\} + \{\mathbb{E}[Y_{1,G_0}] - \mathbb{E}[Y_{0,G_0}]\}. \]
The first bracket is the randomized interventional indirect effect; the second is the randomized interventional direct effect.
8.2 Estimator
For binary \(Z\), one identifying expression is
\[ \psi(a,a') = \mathbb{E}\left[ \mathbb{E}\left\{ \sum_z \mathbb{E}[Y \mid A=a,Z=z,M,W] \mathbb{P}(Z=z \mid A=a,W) \mid A=a',W \right\} \right]. \]
8.3 Data for this section
rie_n = 100_000
rie_rng = np.random.default_rng(4)
rie_w = rie_rng.normal(size=rie_n)
rie_a = rie_rng.binomial(1, 0.5, size=rie_n)
rie_z = rie_rng.binomial(1, 0.5 + 0.2 * rie_a)
rie_m = rie_rng.normal(rie_w + rie_a - rie_z, 1.0)
rie_y = rie_rng.normal(rie_w + rie_a + rie_z + rie_m, 1.0)8.4 Work through \(\psi(1,0)\)
Set \(a=1\) and \(a'=0\). The outcome regression in the inner summation is
\[ \hat\mu(a,z,m,w) \approx \mathbb{E}[Y \mid A=a,Z=z,M=m,W=w]. \]
rie_a_value = 1
rie_a_prime = 0
rie_mu_hat = linear_fit(lm_features(rie_m, rie_a, rie_z, rie_w), rie_y)Evaluate \(\hat\mu(a,z,M_i,W_i)\) at \(z=0\) and \(z=1\).
rie_mu_hat_a_z0_MW = rie_mu_hat.predict(
lm_features(rie_m, np.full(rie_n, rie_a_value), np.zeros(rie_n), rie_w)
)
rie_mu_hat_a_z1_MW = rie_mu_hat.predict(
lm_features(rie_m, np.full(rie_n, rie_a_value), np.ones(rie_n), rie_w)
)The weight for the sum over \(z\) is
\[ \hat p_Z(1 \mid a,W_i) \approx \mathbb{P}(Z=1 \mid A=a,W_i). \]
rie_z_hat = linear_fit(lm_features(rie_a, rie_w), rie_z)
rie_p_z1_given_aW = clip_prob(
rie_z_hat.predict(lm_features(np.full(rie_n, rie_a_value), rie_w)),
eps=1e-6,
)
rie_p_z0_given_aW = 1.0 - rie_p_z1_given_aWNow form the inner summation:
\[ \sum_z \hat\mu(a,z,M_i,W_i)\hat{\mathbb{P}}(Z=z\mid A=a,W_i). \]
rie_inner_sum = rie_mu_hat_a_z0_MW * rie_p_z0_given_aW
rie_inner_sum += rie_mu_hat_a_z1_MW * rie_p_z1_given_aWThe outer conditional expectation is a regression of that pseudo-outcome on \((A,W)\), evaluated at \(A=a'\) and averaged over the sample.
rie_eta_hat = linear_fit(lm_features(rie_a, rie_w), rie_inner_sum)
rie_eta_hat_aprimeW = rie_eta_hat.predict(
lm_features(np.full(rie_n, rie_a_prime), rie_w)
)
psi_10 = float(np.mean(rie_eta_hat_aprimeW))
psi_101.19538810775477178.5 Package the same calculation for the decomposition
The direct and indirect decomposition needs the same plug-in map at \((a,a')=(1,1)\) and \((0,0)\).
def randomized_interventional_mean(a_value, a_prime):
mu_hat = linear_fit(lm_features(rie_m, rie_a, rie_z, rie_w), rie_y)
mu_hat_a_z0_MW = mu_hat.predict(
lm_features(rie_m, np.full(rie_n, a_value), np.zeros(rie_n), rie_w)
)
mu_hat_a_z1_MW = mu_hat.predict(
lm_features(rie_m, np.full(rie_n, a_value), np.ones(rie_n), rie_w)
)
z_hat = linear_fit(lm_features(rie_a, rie_w), rie_z)
p_z1_given_aW = clip_prob(
z_hat.predict(lm_features(np.full(rie_n, a_value), rie_w)),
eps=1e-6,
)
p_z0_given_aW = 1.0 - p_z1_given_aW
inner_sum = mu_hat_a_z0_MW * p_z0_given_aW
inner_sum += mu_hat_a_z1_MW * p_z1_given_aW
eta_hat = linear_fit(lm_features(rie_a, rie_w), inner_sum)
eta_hat_aprimeW = eta_hat.predict(lm_features(np.full(rie_n, a_prime), rie_w))
return float(np.mean(eta_hat_aprimeW))
psi_11 = randomized_interventional_mean(a_value=1, a_prime=1)
psi_00 = randomized_interventional_mean(a_value=0, a_prime=0)
pd.Series(
{
"psi_10": psi_10,
"psi_11": psi_11,
"psi_00": psi_00,
"randomized_direct": psi_10 - psi_00,
"randomized_indirect": psi_11 - psi_10,
"total": psi_11 - psi_00,
}
)psi_10 1.195388
psi_11 1.995820
psi_00 -0.006821
randomized_direct 1.202209
randomized_indirect 0.800432
total 2.002641
dtype: float649 Stochastic effects under an incremental propensity score intervention
9.1 Estimand
For binary \(A\), let
\[ g(1\mid w) = \mathbb{P}(A=1 \mid W=w). \]
An incremental propensity score intervention with multiplier \(\delta\) changes the odds of treatment:
\[ g_\delta(1\mid w) = \frac{\delta g(1\mid w)} {\delta g(1\mid w) + 1 - g(1\mid w)}. \]
The stochastic total effect can be decomposed as
\[ \mathbb{E}[Y_{A_\delta,M_{A_\delta}} - Y_{A,M_A}] = \mathbb{E}[Y_{A_\delta,M_{A_\delta}} - Y_{A_\delta,M}] + \mathbb{E}[Y_{A_\delta,M} - Y_{A,M}]. \]
9.2 Estimator
Let
\[ b_1(a,m,w) = \mathbb{E}[Y \mid A=a,M=m,W=w] \]
and
\[ b_2(a,w) = \mathbb{E}[Y \mid A=a,W=w]. \]
For binary \(A\), the stochastic indirect plug-in estimator averages
\[ \sum_a \{ \hat b_2(a,W_i) - \hat b_1(a,M_i,W_i) \} \hat g_\delta(a\mid W_i), \]
and the direct plug-in estimator averages
\[ \sum_a \{ \hat b_1(a,M_i,W_i) - Y_i \} \hat g_\delta(a\mid W_i). \]
9.3 Data for this section
ipsi_n = 100_000
ipsi_delta = 2.0
ipsi_rng = np.random.default_rng(5)
ipsi_w = ipsi_rng.normal(size=ipsi_n)
ipsi_a = ipsi_rng.binomial(1, expit(1.0 + ipsi_w))
ipsi_m = ipsi_rng.normal(ipsi_w + ipsi_a, 1.0)
ipsi_y = ipsi_rng.normal(ipsi_w + ipsi_a + ipsi_m, 1.0)9.4 Estimate \(b_1(a,m,w)\)
ipsi_b1_hat = linear_fit(lm_features(ipsi_m, ipsi_a, ipsi_w), ipsi_y)
ipsi_b1_hat_1MW = ipsi_b1_hat.predict(lm_features(ipsi_m, np.ones(ipsi_n), ipsi_w))
ipsi_b1_hat_0MW = ipsi_b1_hat.predict(lm_features(ipsi_m, np.zeros(ipsi_n), ipsi_w))9.5 Estimate \(b_2(a,w)\)
ipsi_b2_hat = linear_fit(lm_features(ipsi_a, ipsi_w), ipsi_y)
ipsi_b2_hat_1W = ipsi_b2_hat.predict(lm_features(np.ones(ipsi_n), ipsi_w))
ipsi_b2_hat_0W = ipsi_b2_hat.predict(lm_features(np.zeros(ipsi_n), ipsi_w))9.6 Estimate \(g_\delta(a\mid W_i)\)
First estimate the observed propensity score, then apply the incremental odds-shift formula.
ipsi_g_hat = logit_fit(lm_features(ipsi_w), ipsi_a)
ipsi_g1_hat = clip_prob(predict_prob(ipsi_g_hat, lm_features(ipsi_w)), eps=1e-6)
ipsi_g_delta_1W = (
ipsi_delta * ipsi_g1_hat
/ (ipsi_delta * ipsi_g1_hat + 1.0 - ipsi_g1_hat)
)
ipsi_g_delta_0W = 1.0 - ipsi_g_delta_1W9.7 Construct the indirect summand
This cell is the binary-\(A\) expansion of
\[ \sum_a \{ \hat b_2(a,W_i) - \hat b_1(a,M_i,W_i) \} \hat g_\delta(a\mid W_i). \]
ipsi_indirect_i = (ipsi_b2_hat_1W - ipsi_b1_hat_1MW) * ipsi_g_delta_1W
ipsi_indirect_i += (ipsi_b2_hat_0W - ipsi_b1_hat_0MW) * ipsi_g_delta_0W9.8 Construct the direct summand
This cell is the binary-\(A\) expansion of
\[ \sum_a \{ \hat b_1(a,M_i,W_i) - Y_i \} \hat g_\delta(a\mid W_i). \]
ipsi_direct_i = (ipsi_b1_hat_1MW - ipsi_y) * ipsi_g_delta_1W
ipsi_direct_i += (ipsi_b1_hat_0MW - ipsi_y) * ipsi_g_delta_0W9.9 Average the summands
ipsi_odds_ratio = (
ipsi_g_delta_1W / ipsi_g_delta_0W
/ (ipsi_g1_hat / (1.0 - ipsi_g1_hat))
)
pd.Series(
{
"direct": float(np.mean(ipsi_direct_i)),
"indirect": float(np.mean(ipsi_indirect_i)),
"total": float(np.mean(ipsi_direct_i + ipsi_indirect_i)),
"observed_ps_mean": float(np.mean(ipsi_g1_hat)),
"intervention_ps_mean": float(np.mean(ipsi_g_delta_1W)),
"odds_ratio_mean": float(np.mean(ipsi_odds_ratio)),
}
)direct 0.109234
indirect 0.109234
total 0.218468
observed_ps_mean 0.699220
intervention_ps_mean 0.807988
odds_ratio_mean 2.000000
dtype: float6410 Real workshop data
The workshop’s running data example is mma::weight_behavior. The code below downloads the current CRAN source package for mma, extracts weight_behavior.rda, drops missing observations, and one-hot encodes the factor columns to match the workshop variables. It is marked eval: false because loading the R data requires the optional pyreadr package.
from pathlib import Path
import gzip
import tarfile
import urllib.request
import pyreadr
def cran_package_version(package):
packages_url = "https://cran.r-project.org/src/contrib/PACKAGES.gz"
with urllib.request.urlopen(packages_url, timeout=30) as response:
text = gzip.decompress(response.read()).decode("utf-8", errors="replace")
for block in text.split("\n\n"):
lines = dict(line.split(": ", 1) for line in block.splitlines() if ": " in line)
if lines.get("Package") == package:
return lines["Version"]
raise RuntimeError(f"Could not find CRAN package metadata for {package!r}.")
def preprocess_weight_behavior(data):
df = data.dropna().copy()
if "sex" in df.columns:
df["sex_F"] = (df["sex"].astype(str) == "F").astype(int)
df = df.drop(columns=["sex"])
for col in ["sports", "snack"]:
if col in df.columns:
dummies = pd.get_dummies(df[col], prefix=col, drop_first=True, dtype=int)
df = pd.concat([df.drop(columns=[col]), dummies], axis=1)
return df
def load_weight_behavior(cache_dir=None):
cache = Path(cache_dir or Path.home() / ".cache" / "ser2026_mediation_workshop")
cache.mkdir(parents=True, exist_ok=True)
rda_path = cache / "weight_behavior.rda"
if not rda_path.exists():
version = cran_package_version("mma")
url = f"https://cran.r-project.org/src/contrib/mma_{version}.tar.gz"
archive_path = cache / f"mma_{version}.tar.gz"
urllib.request.urlretrieve(url, archive_path)
with tarfile.open(archive_path, "r:gz") as tar:
member = next(
item
for item in tar.getmembers()
if item.name.endswith("data/weight_behavior.rda")
)
extracted = tar.extractfile(member)
if extracted is None:
raise RuntimeError("Could not extract weight_behavior.rda.")
rda_path.write_bytes(extracted.read())
result = pyreadr.read_r(str(rda_path))
data = next(iter(result.values()))
return preprocess_weight_behavior(data)
weight_behavior = load_weight_behavior()
weight_behavior[
["age", "sex_F", "tvhours", "sports_2", "snack_2", "exercises", "bmi"]
].head()The corresponding crumble() call in R targets the natural direct and indirect effects of sports participation on BMI through snacking and exercise, adjusting for age, sex, and TV hours. A full Python analogue of that call would require a Python implementation of crumble’s Riesz-representer learning and cross-fitted one-step estimator for high-dimensional mediators. The code in this directory instead exposes the estimator pieces directly.