solow_stats/mediation.rs
1//! Deterministic linear mediation effects.
2//!
3//! [`mediation_effects`] computes the point estimates of a causal mediation
4//! analysis for the linear-outcome / linear-mediator case with no interaction
5//! involving the mediator. These are exactly the deterministic limit of the
6//! reference `mediation.Mediation.fit` Monte-Carlo procedure: replacing each
7//! simulated parameter draw by its point estimate and each simulated potential
8//! mediator by its conditional mean (valid because the linear outcome predictor
9//! is affine in the mediator) collapses the average causal mediated effect
10//! (ACME), average direct effect (ADE), total effect and proportion mediated to
11//! closed-form quantities.
12//!
13//! The reference also reports simulation-based confidence intervals and
14//! p-values; those are RNG-driven and out of scope here (see crate notes).
15
16use ndarray::{Array1, Array2};
17use solow_regression::LinearModel;
18
19/// Deterministic point estimates from a linear mediation analysis.
20///
21/// Field names follow the reference `MediationResults`: `_ctrl` and `_tx` are
22/// the effects with the *other* treatment arm held at control (0) or treated
23/// (1); `_avg` averages the two. In the linear no-interaction case the control
24/// and treated variants coincide.
25#[derive(Debug, Clone, Copy)]
26pub struct MediationResults {
27 /// Average causal mediated (indirect) effect, control arm.
28 pub acme_ctrl: f64,
29 /// Average causal mediated (indirect) effect, treated arm.
30 pub acme_tx: f64,
31 /// Average direct effect, control arm.
32 pub ade_ctrl: f64,
33 /// Average direct effect, treated arm.
34 pub ade_tx: f64,
35 /// Total effect, `(ACME_ctrl + ACME_tx + ADE_ctrl + ADE_tx) / 2`.
36 pub total_effect: f64,
37 /// Proportion mediated, control arm.
38 pub prop_med_ctrl: f64,
39 /// Proportion mediated, treated arm.
40 pub prop_med_tx: f64,
41 /// Average proportion mediated.
42 pub prop_med_avg: f64,
43 /// Average causal mediated (indirect) effect.
44 pub acme_avg: f64,
45 /// Average direct effect.
46 pub ade_avg: f64,
47}
48
49/// Specification of a linear mediation analysis (no formulas).
50///
51/// `outcome_endog` / `outcome_exog` define the outcome regression `Y ~ X_o`
52/// (which must include both the mediator and the exposure as columns), and
53/// `mediator_endog` / `mediator_exog` the mediator regression `M ~ X_m`. The
54/// `*_pos` fields give the column positions of the exposure in each design and
55/// of the mediator in the outcome design, exactly as in the reference's
56/// positional API (`Mediation(outcome, mediator, [exp_pos_outcome,
57/// exp_pos_mediator], med_pos_outcome)`).
58#[derive(Debug, Clone)]
59pub struct Mediation {
60 /// Outcome model response vector.
61 pub outcome_endog: Array1<f64>,
62 /// Outcome model design matrix (includes exposure and mediator columns).
63 pub outcome_exog: Array2<f64>,
64 /// Mediator model response vector.
65 pub mediator_endog: Array1<f64>,
66 /// Mediator model design matrix (includes the exposure column).
67 pub mediator_exog: Array2<f64>,
68 /// Column position of the exposure in the outcome design.
69 pub exp_pos_outcome: usize,
70 /// Column position of the exposure in the mediator design.
71 pub exp_pos_mediator: usize,
72 /// Column position of the mediator in the outcome design.
73 pub med_pos_outcome: usize,
74}
75
76impl Mediation {
77 /// Mediator design with the exposure column set to `exposure`.
78 fn mediator_exog_at(&self, exposure: f64) -> Array2<f64> {
79 let mut m = self.mediator_exog.clone();
80 m.column_mut(self.exp_pos_mediator).fill(exposure);
81 m
82 }
83
84 /// Outcome design with exposure set to `exposure` and the mediator column
85 /// set to the per-observation `mediator` values.
86 fn outcome_exog_at(&self, exposure: f64, mediator: &Array1<f64>) -> Array2<f64> {
87 let mut o = self.outcome_exog.clone();
88 o.column_mut(self.exp_pos_outcome).fill(exposure);
89 o.column_mut(self.med_pos_outcome).assign(mediator);
90 o
91 }
92
93 /// Fit both linear models and return the deterministic mediation point
94 /// estimates. Mirrors the deterministic limit of the reference
95 /// `Mediation.fit` for linear outcome and mediator models.
96 pub fn fit(&self) -> MediationResults {
97 let beta_o = LinearModel::ols(self.outcome_endog.clone(), self.outcome_exog.clone())
98 .expect("outcome OLS")
99 .fit()
100 .expect("outcome fit")
101 .params;
102 let beta_m = LinearModel::ols(self.mediator_endog.clone(), self.mediator_exog.clone())
103 .expect("mediator OLS")
104 .fit()
105 .expect("mediator fit")
106 .params;
107
108 // potential_mediator[tm] = E[M | exposure = tm] (conditional mean).
109 let pm: [Array1<f64>; 2] = [
110 self.mediator_exog_at(0.0).dot(&beta_m),
111 self.mediator_exog_at(1.0).dot(&beta_m),
112 ];
113
114 // predicted_outcomes[tm][te] = E[Y | mediator = pm[tm], exposure = te].
115 let predict = |tm: usize, te: f64| self.outcome_exog_at(te, &pm[tm]).dot(&beta_o);
116 let po: [[Array1<f64>; 2]; 2] = [
117 [predict(0, 0.0), predict(0, 1.0)],
118 [predict(1, 0.0), predict(1, 1.0)],
119 ];
120
121 let mean = |a: &Array1<f64>| a.sum() / a.len() as f64;
122
123 // indirect_effects[t] = po[1][t] - po[0][t]; direct[t] = po[t][1] - po[t][0].
124 let acme_ctrl = mean(&(&po[1][0] - &po[0][0]));
125 let acme_tx = mean(&(&po[1][1] - &po[0][1]));
126 let ade_ctrl = mean(&(&po[0][1] - &po[0][0]));
127 let ade_tx = mean(&(&po[1][1] - &po[1][0]));
128
129 let total_effect = (acme_ctrl + acme_tx + ade_ctrl + ade_tx) / 2.0;
130 let prop_med_ctrl = acme_ctrl / total_effect;
131 let prop_med_tx = acme_tx / total_effect;
132 let prop_med_avg = (prop_med_ctrl + prop_med_tx) / 2.0;
133 let acme_avg = (acme_ctrl + acme_tx) / 2.0;
134 let ade_avg = (ade_ctrl + ade_tx) / 2.0;
135
136 MediationResults {
137 acme_ctrl,
138 acme_tx,
139 ade_ctrl,
140 ade_tx,
141 total_effect,
142 prop_med_ctrl,
143 prop_med_tx,
144 prop_med_avg,
145 acme_avg,
146 ade_avg,
147 }
148 }
149}
150
151#[cfg(test)]
152mod tests {
153 use super::*;
154 use ndarray::array;
155
156 #[test]
157 fn matches_baron_kenny() {
158 // M is NOT collinear with [const, T, Z] (it carries an extra component
159 // `noise`), so the outcome design is full rank and OLS recovers the
160 // exact coefficients from the noiseless outcome. Then ACME = a*b and
161 // ADE = c' exactly.
162 let t = array![0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 1.0, 0.0];
163 let z = array![10.0, 12.0, 14.0, 9.0, 11.0, 13.0, 8.0, 15.0, 7.0, 16.0];
164 let extra = array![0.3, -0.5, 0.1, 0.8, -0.2, 0.6, -0.7, 0.4, 0.2, -0.9];
165 let n = t.len();
166 let m: Array1<f64> =
167 Array1::from_iter((0..n).map(|i| 0.5 + 0.8 * t[i] + 0.05 * z[i] + extra[i]));
168 let y: Array1<f64> =
169 Array1::from_iter((0..n).map(|i| 1.0 + 1.5 * m[i] + 0.7 * t[i] + 0.02 * z[i]));
170
171 // outcome exog: [const, M, T, Z]; mediator exog: [const, T, Z].
172 let ones = Array1::ones(n);
173 let mut oe = Array2::zeros((n, 4));
174 oe.column_mut(0).assign(&ones);
175 oe.column_mut(1).assign(&m);
176 oe.column_mut(2).assign(&t);
177 oe.column_mut(3).assign(&z);
178 let mut me = Array2::zeros((n, 3));
179 me.column_mut(0).assign(&ones);
180 me.column_mut(1).assign(&t);
181 me.column_mut(2).assign(&z);
182
183 // Recover the fitted mediator T-coefficient (a) to verify the
184 // Baron-Kenny identity ACME = a * b with b = 1.5 (exactly recovered
185 // since Y is noiseless in [const, M, T, Z]).
186 let a_fit = LinearModel::ols(m.clone(), me.clone())
187 .unwrap()
188 .fit()
189 .unwrap()
190 .params[1];
191
192 let med = Mediation {
193 outcome_endog: y,
194 outcome_exog: oe,
195 mediator_endog: m,
196 mediator_exog: me,
197 exp_pos_outcome: 2,
198 exp_pos_mediator: 1,
199 med_pos_outcome: 1,
200 };
201 let r = med.fit();
202 // ACME = a * b with b = 1.5; ADE = c' = 0.7; control == treated (linear).
203 assert!(
204 (r.acme_avg - a_fit * 1.5).abs() < 1e-8,
205 "acme {}",
206 r.acme_avg
207 );
208 assert!((r.ade_avg - 0.7).abs() < 1e-8, "ade {}", r.ade_avg);
209 assert!((r.acme_ctrl - r.acme_tx).abs() < 1e-10);
210 assert!((r.total_effect - (r.acme_avg + r.ade_avg)).abs() < 1e-8);
211 }
212}