dml_core/model/
irm.rs

1//! Interactive regression model (IRM) for a binary treatment:
2//!
3//! ```text
4//! Y = g(D, X) + U,   E[U | X, D] = 0
5//! D = m(X) + V,      E[V | X] = 0,   m(X) = P(D = 1 | X)
6//! ```
7//!
8//! Targets the average treatment effect (ATE) or the average treatment effect
9//! on the treated (ATTE) with the doubly robust (AIPW) score.
10
11use ndarray::{Array1, ArrayView1, ArrayView2};
12
13use super::{dml_model_api, LinearScoreModel, NuisanceTarget, ScoreEstimate, ScoreInput};
14use crate::crossfit::{cross_fit_predict_proba, cross_fit_predict_subset, Fold, SampleSplitting};
15use crate::data::DmlData;
16use crate::error::{DmlError, Result};
17use crate::learner::{Classifier, Regressor};
18use crate::score::LinearScore;
19use crate::sensitivity::{Nu2Parts, SensitivityParts};
20
21/// Target parameter of [`Irm`].
22#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
23pub enum IrmScore {
24    /// `E[g(1, X) − g(0, X)]`
25    #[default]
26    Ate,
27    /// `E[g(1, X) − g(0, X) | D = 1]`
28    Atte,
29}
30
31/// IRM estimator: `ml_g` learns `E[Y | D = d, X]` (fitted separately on the
32/// treated and control rows of each training fold), `ml_m` learns the
33/// propensity score `P(D = 1 | X)`.
34pub struct Irm {
35    ml_g: Box<dyn Regressor>,
36    ml_m: Box<dyn Classifier>,
37    score: IrmScore,
38    clipping_threshold: f64,
39    normalize_ipw: bool,
40    splitting: SampleSplitting,
41}
42
43impl Irm {
44    pub fn new(ml_g: impl Regressor + 'static, ml_m: impl Classifier + 'static) -> Self {
45        Self::from_boxed(Box::new(ml_g), Box::new(ml_m))
46    }
47
48    pub fn from_boxed(ml_g: Box<dyn Regressor>, ml_m: Box<dyn Classifier>) -> Self {
49        Self {
50            ml_g,
51            ml_m,
52            score: IrmScore::default(),
53            clipping_threshold: 1e-2,
54            normalize_ipw: false,
55            splitting: SampleSplitting::default(),
56        }
57    }
58
59    pub fn score(mut self, score: IrmScore) -> Self {
60        self.score = score;
61        self
62    }
63
64    /// Propensity scores are clipped to `[t, 1 − t]` (default `t = 0.01`).
65    pub fn clipping_threshold(mut self, t: f64) -> Self {
66        self.clipping_threshold = t;
67        self
68    }
69
70    /// Normalise the inverse-probability weights to mean one within each
71    /// treatment arm (Hájek-style), as DoubleML's `normalize_ipw=True`.
72    pub fn normalize_ipw(mut self, normalize: bool) -> Self {
73        self.normalize_ipw = normalize;
74        self
75    }
76
77    dml_model_api!();
78}
79
80/// DoubleML's `_normalize_ipw`: rescale `m̂` so that `mean(D/m̂)` and
81/// `mean((1−D)/(1−m̂))` both become one.
82pub(crate) fn normalize_propensity(m: &Array1<f64>, d: &Array1<f64>) -> Array1<f64> {
83    let n = m.len() as f64;
84    let mean1 = d.iter().zip(m).map(|(&di, &mi)| di / mi).sum::<f64>() / n;
85    let mean0 = d
86        .iter()
87        .zip(m)
88        .map(|(&di, &mi)| (1.0 - di) / (1.0 - mi))
89        .sum::<f64>()
90        / n;
91    d.iter()
92        .zip(m)
93        .map(|(&di, &mi)| di * mi * mean1 + (1.0 - di) * (1.0 - (1.0 - mi) * mean0))
94        .collect()
95}
96
97impl LinearScoreModel for Irm {
98    fn name(&self) -> String {
99        match self.score {
100            IrmScore::Ate => "IRM (ATE)".into(),
101            IrmScore::Atte => "IRM (ATTE)".into(),
102        }
103    }
104
105    fn validate(&self, data: &DmlData) -> Result<()> {
106        if data.n_treat() != 1 || !data.is_binary_treatment(0) {
107            return Err(DmlError::InvalidArgument(
108                "IRM requires exactly one binary (0/1) treatment".into(),
109            ));
110        }
111        if !(self.clipping_threshold > 0.0 && self.clipping_threshold < 0.5) {
112            return Err(DmlError::InvalidArgument(
113                "clipping_threshold must be in (0, 0.5)".into(),
114            ));
115        }
116        Ok(())
117    }
118
119    fn estimate_score(&self, input: &ScoreInput<'_>, folds: &[Fold]) -> Result<ScoreEstimate> {
120        let ScoreInput { y, d, x, .. } = *input;
121        let (g0, g1, m) = irm_nuisances(
122            self.ml_g.as_ref(),
123            self.ml_m.as_ref(),
124            x,
125            y,
126            d,
127            folds,
128            self.clipping_threshold,
129        )?;
130        let d = d.to_owned();
131        let m_adj = if self.normalize_ipw {
132            normalize_propensity(&m, &d)
133        } else {
134            m.clone()
135        };
136        let ipw = ipw_residual(y, &d, &g0, &g1, &m_adj);
137
138        let n = d.len();
139        let (weights, weights_bar) = match self.score {
140            IrmScore::Ate => (Array1::ones(n), Array1::ones(n)),
141            IrmScore::Atte => {
142                let p = d.mean().expect("n > 0");
143                (&d / p, &m_adj / p)
144            }
145        };
146        let psi_b = &weights * &(&g1 - &g0) + &weights_bar * &ipw;
147        let psi_a = -(&weights / weights.mean().expect("n > 0"));
148
149        // Sensitivity: outcome residual y − ĝ(D, X); Riesz representer and
150        // m(W, α) of the (weighted) average treatment effect.
151        let outcome_residual = &y - &(&d * &g1) - &(&(1.0 - &d) * &g0);
152        let m_alpha = &weights * &weights_bar * &(1.0 / &m_adj + 1.0 / &(1.0 - &m_adj));
153        let riesz = &weights_bar * &(&d / &m_adj - &(1.0 - &d) / &(1.0 - &m_adj));
154        let nu2 = 2.0 * &m_alpha - &riesz.mapv(|a| a * a);
155        Ok(ScoreEstimate::new(
156            LinearScore::new(psi_a, psi_b)?,
157            vec![("ml_g0", g0), ("ml_g1", g1), ("ml_m", m)],
158        )
159        .with_sensitivity(SensitivityParts {
160            outcome_residual,
161            theta_factor: None,
162            riesz,
163            nu2: Nu2Parts::Elementwise(nu2),
164        })
165        .with_targets(vec![
166            NuisanceTarget::regression("ml_g0", y.to_owned()).only_where(|i| d[i] == 0.0),
167            NuisanceTarget::regression("ml_g1", y.to_owned()).only_where(|i| d[i] == 1.0),
168            NuisanceTarget::classification("ml_m", d.clone()),
169        ]))
170    }
171}
172
173/// Cross-fitted IRM nuisances: `ĝ₀` (trained on controls), `ĝ₁` (trained on
174/// treated) and the clipped propensity `m̂`.
175pub(crate) fn irm_nuisances(
176    ml_g: &dyn Regressor,
177    ml_m: &dyn Classifier,
178    x: ArrayView2<'_, f64>,
179    y: ArrayView1<'_, f64>,
180    d: ArrayView1<'_, f64>,
181    folds: &[Fold],
182    clipping_threshold: f64,
183) -> Result<(Array1<f64>, Array1<f64>, Array1<f64>)> {
184    let treated: Vec<bool> = d.iter().map(|&v| v == 1.0).collect();
185    let control: Vec<bool> = treated.iter().map(|t| !t).collect();
186    let g0 = cross_fit_predict_subset(ml_g, x, y, folds, &control)?;
187    let g1 = cross_fit_predict_subset(ml_g, x, y, folds, &treated)?;
188    let t = clipping_threshold;
189    let m = cross_fit_predict_proba(ml_m, x, d, folds)?.mapv(|p| p.clamp(t, 1.0 - t));
190    Ok((g0, g1, m))
191}
192
193/// Inverse-probability-weighted residual correction
194/// `D(Y − ĝ₁)/m̂ − (1 − D)(Y − ĝ₀)/(1 − m̂)`.
195pub(crate) fn ipw_residual(
196    y: ArrayView1<'_, f64>,
197    d: &Array1<f64>,
198    g0: &Array1<f64>,
199    g1: &Array1<f64>,
200    m: &Array1<f64>,
201) -> Array1<f64> {
202    let u0: Array1<f64> = &y - g0;
203    let u1: Array1<f64> = &y - g1;
204    d * &u1 / m - &(1.0 - d) * &u0 / &(1.0 - m)
205}
206
207#[cfg(test)]
208mod tests {
209    use super::*;
210    use crate::learner::{LogisticRegression, Ols};
211    use ndarray::{array, Array2};
212
213    #[test]
214    fn rejects_continuous_treatment_and_bad_clipping() {
215        let data = DmlData::new(
216            array![1.0, 2.0, 3.0],
217            array![0.0, 0.5, 1.0],
218            array![[1.0], [2.0], [3.0]],
219        )
220        .unwrap();
221        let irm = Irm::new(Ols::new(), LogisticRegression::default());
222        assert!(irm.fit(&data).is_err());
223        let binary =
224            DmlData::new(array![1.0, 2.0], array![0.0, 1.0], array![[1.0], [2.0]]).unwrap();
225        let bad = Irm::new(Ols::new(), LogisticRegression::default()).clipping_threshold(0.7);
226        assert!(bad.validate(&binary).is_err());
227    }
228
229    #[test]
230    fn normalized_weights_have_unit_means() {
231        let m = array![0.2, 0.5, 0.7, 0.4];
232        let d = array![1.0, 0.0, 1.0, 0.0];
233        let adj = normalize_propensity(&m, &d);
234        let n = 4.0;
235        let mean1: f64 = d.iter().zip(&adj).map(|(&a, &b)| a / b).sum::<f64>() / n;
236        let mean0: f64 = d
237            .iter()
238            .zip(&adj)
239            .map(|(&a, &b)| (1.0 - a) / (1.0 - b))
240            .sum::<f64>()
241            / n;
242        assert!((mean1 - 1.0).abs() < 1e-12 && (mean0 - 1.0).abs() < 1e-12);
243    }
244
245    #[test]
246    fn randomized_experiment_recovers_difference_in_means() {
247        // Constant propensity, no covariate signal: ATE ≈ treated − control mean.
248        let n = 200;
249        let x = Array2::from_shape_fn((n, 1), |(i, _)| ((i * 37) % 17) as f64);
250        let d: Array1<f64> = (0..n).map(|i| ((i * 7919) % 2) as f64).collect();
251        let y: Array1<f64> = (0..n).map(|i| 2.0 * d[i] + ((i * 31) % 7) as f64).collect();
252        let data = DmlData::new(y, d, x).unwrap();
253        let fit = Irm::new(Ols::new(), LogisticRegression::default())
254            .fit(&data)
255            .unwrap();
256        assert!((fit.coef()[0] - 2.0).abs() < 0.5, "ate = {}", fit.coef()[0]);
257    }
258}
dml_core/model/irm.rs

dml_core/model/
irm.rs
