--- title: "Mediation estimands, estimators, and Python code" format: html: toc: true code-fold: false number-sections: true jupyter: python3 --- 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](source.txt). ## 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. ```{python} 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) ``` ## 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. ## 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. | ## 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. ```{python} 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. ```{python} 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, } ) ``` ## Controlled direct effect ### 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]. $$ ### Data for this section The simple data-generating process makes the CDE transparent: $Y$ is linear in $W$, $A$, and $M$. ```{python} 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) ``` ### 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]. $$ ```{python} cde_b_hat = linear_fit(lm_features(cde_m, cde_a, cde_w), cde_y) ``` ### 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$. ```{python} 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_0mW ``` ### 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)\}. $$ ```{python} cde_estimate = np.mean(cde_summand) float(cde_estimate) ``` ## Natural direct effect by g-computation ### 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)]. $$ ### 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$. ```{python} 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() ``` ### 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]. $$ ```{python} nde_b_hat = linear_fit(lm_features(nde_m, nde_a, nde_w), nde_y) ``` ### 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. ```{python} 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_0MW ``` ### 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)$. ```{python} 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, } ) ``` ### 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. ```{python} 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)) ``` ## One-step estimator for the natural direct effect ### 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]. $$ ### 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). $$ ### 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. ```{python} 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]]) ``` ### Estimate the nuisance functions These three fitted objects correspond directly to $\hat b$, $\hat e$, and $\hat g$ in the displayed EIF. ```{python} 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. ```{python} 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) ``` ### 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\}. $$ ```{python} 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)) ``` ### Translate the EIF line by line The first line of $\phi$ is the inverse-probability weighted outcome residual. ```{python} 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. ```{python} 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)$. ```{python} 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), } ) ``` ### Reusable version for simulation Again, the function below is the same displayed calculation packaged for repeated sampling. ```{python} 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)) ``` ### Small repeated-sampling comparison ```{python} 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), } ) ``` ## Randomized interventional effects ### 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. ### 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]. $$ ### Data for this section ```{python} 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) ``` ### 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]. $$ ```{python} 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$. ```{python} 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). $$ ```{python} 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_aW ``` Now form the inner summation: $$ \sum_z \hat\mu(a,z,M_i,W_i)\hat{\mathbb{P}}(Z=z\mid A=a,W_i). $$ ```{python} 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_aW ``` The outer conditional expectation is a regression of that pseudo-outcome on $(A,W)$, evaluated at $A=a'$ and averaged over the sample. ```{python} 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_10 ``` ### 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)$. ```{python} 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, } ) ``` ## Stochastic effects under an incremental propensity score intervention ### 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}]. $$ ### 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). $$ ### Data for this section ```{python} 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) ``` ### Estimate $b_1(a,m,w)$ ```{python} 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)) ``` ### Estimate $b_2(a,w)$ ```{python} 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)) ``` ### Estimate $g_\delta(a\mid W_i)$ First estimate the observed propensity score, then apply the incremental odds-shift formula. ```{python} 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_1W ``` ### 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). $$ ```{python} 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_0W ``` ### 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). $$ ```{python} 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_0W ``` ### Average the summands ```{python} 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)), } ) ``` ## 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. ```{python} #| eval: false 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.