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

dml_core/model/
irm.rs
