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, Array2, ArrayView1, ArrayView2};
12
13use super::{
14    dml_model_api, nuisance_api, CateSignal, LinearScoreModel, NuisanceLearners, NuisanceTarget,
15    ScoreEstimate, ScoreInput,
16};
17use crate::compat::np_mean;
18use crate::crossfit::{cross_fit_predict_proba, cross_fit_predict_subset, Fold, SampleSplitting};
19use crate::data::DmlData;
20use crate::error::{DmlError, Result};
21use crate::learner::{Classifier, Regressor};
22use crate::propensity::{adjust_propensity, Calibration};
23use crate::score::LinearScore;
24use crate::sensitivity::{Nu2Parts, SensitivityParts};
25
26/// Target parameter of [`Irm`].
27#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
28pub enum IrmScore {
29    /// `E[g(1, X) − g(0, X)]`
30    #[default]
31    Ate,
32    /// `E[g(1, X) − g(0, X) | D = 1]`
33    Atte,
34}
35
36/// IRM estimator: `ml_g` learns `E[Y | D = d, X]` (fitted separately on the
37/// treated and control rows of each training fold), `ml_m` learns the
38/// propensity score `P(D = 1 | X)`.
39pub struct Irm {
40    nuisance: NuisanceLearners,
41    ml_g: Box<dyn Regressor>,
42    ml_m: Box<dyn Classifier>,
43    score: IrmScore,
44    clipping_threshold: f64,
45    calibration: Calibration,
46    normalize_ipw: bool,
47    weights: Option<ScoreWeights>,
48    splitting: SampleSplitting,
49}
50
51/// Observation weights for weighted average effects (DoubleML's `weights`):
52/// `weights` enter the target, `weights_bar` (one column per repetition; by
53/// default equal to `weights`) the inverse-probability correction.
54#[derive(Debug, Clone)]
55pub(crate) struct ScoreWeights {
56    pub weights: Array1<f64>,
57    pub weights_bar: Option<Array2<f64>>,
58}
59
60impl ScoreWeights {
61    pub(crate) fn check(&self, n: usize) -> Result<()> {
62        let ok = |w: &mut dyn Iterator<Item = f64>| {
63            let (mut sum, mut all_ok) = (0.0, true);
64            for v in w {
65                all_ok &= v.is_finite() && v >= 0.0;
66                sum += v;
67            }
68            all_ok && sum > 0.0
69        };
70        if self.weights.len() != n {
71            return Err(DmlError::DimensionMismatch(format!(
72                "weights have length {}, the data {n} rows",
73                self.weights.len()
74            )));
75        }
76        if !ok(&mut self.weights.iter().copied()) {
77            return Err(DmlError::InvalidArgument(
78                "weights must be finite, non-negative and not all zero".into(),
79            ));
80        }
81        if let Some(bar) = &self.weights_bar {
82            if bar.nrows() != n || bar.ncols() == 0 {
83                return Err(DmlError::DimensionMismatch(format!(
84                    "weights_bar must have shape ({n}, n_rep), got {:?}",
85                    bar.dim()
86                )));
87            }
88            if !ok(&mut bar.iter().copied()) {
89                return Err(DmlError::InvalidArgument(
90                    "weights_bar must be finite, non-negative and not all zero".into(),
91                ));
92            }
93        }
94        Ok(())
95    }
96
97    /// `(weights, weights_bar)` of repetition `rep` for an ATE-type score.
98    pub(crate) fn ate(
99        weights: Option<&ScoreWeights>,
100        n: usize,
101        rep: usize,
102    ) -> Result<(Array1<f64>, Array1<f64>)> {
103        let Some(w) = weights else {
104            return Ok((Array1::ones(n), Array1::ones(n)));
105        };
106        let bar = match &w.weights_bar {
107            None => w.weights.clone(),
108            Some(bar) if rep < bar.ncols() => bar.column(rep).to_owned(),
109            Some(bar) => {
110                return Err(DmlError::DimensionMismatch(format!(
111                    "weights_bar has {} columns, one per repetition is needed (repetition {rep})",
112                    bar.ncols()
113                )))
114            }
115        };
116        Ok((w.weights.clone(), bar))
117    }
118}
119
120impl Irm {
121    pub fn new(ml_g: impl Regressor + 'static, ml_m: impl Classifier + 'static) -> Self {
122        Self::from_boxed(Box::new(ml_g), Box::new(ml_m))
123    }
124
125    pub fn from_boxed(ml_g: Box<dyn Regressor>, ml_m: Box<dyn Classifier>) -> Self {
126        Self {
127            nuisance: NuisanceLearners::default(),
128            ml_g,
129            ml_m,
130            score: IrmScore::default(),
131            clipping_threshold: 1e-2,
132            calibration: Calibration::None,
133            normalize_ipw: false,
134            weights: None,
135            splitting: SampleSplitting::default(),
136        }
137    }
138
139    /// Weighted average effect, as DoubleML's `weights`: with the ATE score
140    /// `E[w·(g(1, X) − g(0, X))] / E[w]` (e.g. `w = 1{X ∈ G}` for the effect in
141    /// a subgroup); with the ATTE score `w` must be 0/1 and selects the
142    /// subgroup of the treated.
143    pub fn weights(mut self, weights: Array1<f64>) -> Self {
144        self.weights = Some(ScoreWeights {
145            weights,
146            weights_bar: None,
147        });
148        self
149    }
150
151    /// ATE weights with separate `weights_bar` for the inverse-probability
152    /// correction, one column per repetition (DoubleML's dict form).
153    pub fn weights_with_bar(mut self, weights: Array1<f64>, weights_bar: Array2<f64>) -> Self {
154        self.weights = Some(ScoreWeights {
155            weights,
156            weights_bar: Some(weights_bar),
157        });
158        self
159    }
160
161    pub fn score(mut self, score: IrmScore) -> Self {
162        self.score = score;
163        self
164    }
165
166    /// Propensity scores are clipped to `[t, 1 − t]` (default `t = 0.01`).
167    pub fn clipping_threshold(mut self, t: f64) -> Self {
168        self.clipping_threshold = t;
169        self
170    }
171    /// Calibrate the propensity scores before clipping (default: none), as
172    /// DoubleML's `ps_processor_config`.
173    pub fn calibration(mut self, calibration: Calibration) -> Self {
174        self.calibration = calibration;
175        self
176    }
177
178    /// Normalise the inverse-probability weights to mean one within each
179    /// treatment arm (Hájek-style), as DoubleML's `normalize_ipw=True`.
180    pub fn normalize_ipw(mut self, normalize: bool) -> Self {
181        self.normalize_ipw = normalize;
182        self
183    }
184
185    dml_model_api!();
186    nuisance_api!();
187}
188
189/// DoubleML's `_normalize_ipw`: rescale `m̂` so that `mean(D/m̂)` and
190/// `mean((1−D)/(1−m̂))` both become one.
191pub(crate) fn normalize_propensity(m: &Array1<f64>, d: &Array1<f64>) -> Array1<f64> {
192    // NumPy's pairwise means, so that the quantile models (whose preliminary
193    // estimate is sensitive to the last bit) reproduce DoubleML exactly.
194    let ratio1: Vec<f64> = d.iter().zip(m).map(|(&di, &mi)| di / mi).collect();
195    let ratio0: Vec<f64> = d
196        .iter()
197        .zip(m)
198        .map(|(&di, &mi)| (1.0 - di) / (1.0 - mi))
199        .collect();
200    let (mean1, mean0) = (np_mean(&ratio1), np_mean(&ratio0));
201    d.iter()
202        .zip(m)
203        .map(|(&di, &mi)| di * mi * mean1 + (1.0 - di) * (1.0 - (1.0 - mi) * mean0))
204        .collect()
205}
206
207impl LinearScoreModel for Irm {
208    fn name(&self) -> String {
209        match self.score {
210            IrmScore::Ate => "IRM (ATE)".into(),
211            IrmScore::Atte => "IRM (ATTE)".into(),
212        }
213    }
214
215    fn validate(&self, data: &DmlData) -> Result<()> {
216        self.nuisance.check("IRM", &["ml_g0", "ml_g1"], &["ml_m"])?;
217        if data.n_treat() != 1 || !data.is_binary_treatment(0) {
218            return Err(DmlError::InvalidArgument(
219                "IRM requires exactly one binary (0/1) treatment".into(),
220            ));
221        }
222        if !(self.clipping_threshold > 0.0 && self.clipping_threshold < 0.5) {
223            return Err(DmlError::InvalidArgument(
224                "clipping_threshold must be in (0, 0.5)".into(),
225            ));
226        }
227        if let Some(w) = &self.weights {
228            w.check(data.n_obs())?;
229            if self.score == IrmScore::Atte
230                && (w.weights_bar.is_some() || w.weights.iter().any(|&v| v != 0.0 && v != 1.0))
231            {
232                return Err(DmlError::InvalidArgument(
233                    "with the ATTE score, weights must be 0/1 (without weights_bar)".into(),
234                ));
235            }
236        }
237        Ok(())
238    }
239
240    fn estimate_score(&self, input: &ScoreInput<'_>, folds: &[Fold]) -> Result<ScoreEstimate> {
241        let ScoreInput { y, d, x, rep, .. } = *input;
242        let treated: Vec<bool> = d.iter().map(|&v| v == 1.0).collect();
243        let control: Vec<bool> = treated.iter().map(|t| !t).collect();
244        let nu = &self.nuisance;
245        let ml_g = self.ml_g.as_ref();
246        let g0 = input.predictions_or("ml_g0", || {
247            cross_fit_predict_subset(nu.regressor("ml_g0", ml_g), x, y, folds, &control)
248        })?;
249        let g1 = input.predictions_or("ml_g1", || {
250            cross_fit_predict_subset(nu.regressor("ml_g1", ml_g), x, y, folds, &treated)
251        })?;
252        let m_raw = input.predictions_or("ml_m", || {
253            cross_fit_predict_proba(nu.classifier("ml_m", self.ml_m.as_ref()), x, d, folds)
254        })?;
255        let m = adjust_propensity(m_raw, d, folds, self.clipping_threshold, self.calibration)?;
256        let d = d.to_owned();
257        let m_adj = if self.normalize_ipw {
258            normalize_propensity(&m, &d)
259        } else {
260            m.clone()
261        };
262        let ipw = ipw_residual(y, &d, &g0, &g1, &m_adj);
263
264        let n = d.len();
265        let (weights, weights_bar) = match self.score {
266            IrmScore::Ate => ScoreWeights::ate(self.weights.as_ref(), n, rep)?,
267            IrmScore::Atte => {
268                // The treated in the subgroup selected by the (0/1) weights.
269                let w = self
270                    .weights
271                    .as_ref()
272                    .map_or_else(|| Array1::ones(n), |w| w.weights.clone());
273                let subgroup = &w * &d;
274                let p = subgroup.mean().expect("n > 0");
275                (&subgroup / p, &m_adj * &w / p)
276            }
277        };
278        let psi_b = &weights * &(&g1 - &g0) + &weights_bar * &ipw;
279        let psi_a = -(&weights / weights.mean().expect("n > 0"));
280
281        // Sensitivity: outcome residual y − ĝ(D, X); Riesz representer and
282        // m(W, α) of the (weighted) average treatment effect.
283        let outcome_residual = &y - &(&d * &g1) - &(&(1.0 - &d) * &g0);
284        let m_alpha = &weights * &weights_bar * &(1.0 / &m_adj + 1.0 / &(1.0 - &m_adj));
285        let riesz = &weights_bar * &(&d / &m_adj - &(1.0 - &d) / &(1.0 - &m_adj));
286        let nu2 = 2.0 * &m_alpha - &riesz.mapv(|a| a * a);
287        let estimate = ScoreEstimate::new(
288            LinearScore::new(psi_a, psi_b.clone())?,
289            vec![("ml_g0", g0), ("ml_g1", g1), ("ml_m", m)],
290        )
291        .with_sensitivity(SensitivityParts {
292            outcome_residual,
293            theta_factor: None,
294            riesz,
295            nu2: Nu2Parts::Elementwise(nu2),
296        })
297        .with_targets(vec![
298            NuisanceTarget::regression("ml_g0", y.to_owned()).only_where(|i| d[i] == 0.0),
299            NuisanceTarget::regression("ml_g1", y.to_owned()).only_where(|i| d[i] == 1.0),
300            NuisanceTarget::classification("ml_m", d.clone()),
301        ]);
302        // CATE / GATE: the orthogonal signal ψ_b (ATE score only, as DoubleML).
303        Ok(match self.score {
304            IrmScore::Ate => estimate.with_cate_signal(CateSignal {
305                signal: psi_b,
306                theta_factor: None,
307                basis_scale: None,
308                is_effect: true,
309            }),
310            IrmScore::Atte => estimate,
311        })
312    }
313}
314
315/// Cross-fitted IRM nuisances: `ĝ₀` (trained on controls), `ĝ₁` (trained on
316/// treated) and the clipped propensity `m̂`.
317#[allow(clippy::too_many_arguments)]
318pub(crate) fn irm_nuisances(
319    ml_g: &dyn Regressor,
320    ml_m: &dyn Classifier,
321    x: ArrayView2<'_, f64>,
322    y: ArrayView1<'_, f64>,
323    d: ArrayView1<'_, f64>,
324    folds: &[Fold],
325    clipping_threshold: f64,
326    calibration: Calibration,
327) -> Result<(Array1<f64>, Array1<f64>, Array1<f64>)> {
328    let treated: Vec<bool> = d.iter().map(|&v| v == 1.0).collect();
329    let control: Vec<bool> = treated.iter().map(|t| !t).collect();
330    let g0 = cross_fit_predict_subset(ml_g, x, y, folds, &control)?;
331    let g1 = cross_fit_predict_subset(ml_g, x, y, folds, &treated)?;
332    let m_raw = cross_fit_predict_proba(ml_m, x, d, folds)?;
333    let m = adjust_propensity(m_raw, d, folds, clipping_threshold, calibration)?;
334    Ok((g0, g1, m))
335}
336
337/// Inverse-probability-weighted residual correction
338/// `D(Y − ĝ₁)/m̂ − (1 − D)(Y − ĝ₀)/(1 − m̂)`.
339pub(crate) fn ipw_residual(
340    y: ArrayView1<'_, f64>,
341    d: &Array1<f64>,
342    g0: &Array1<f64>,
343    g1: &Array1<f64>,
344    m: &Array1<f64>,
345) -> Array1<f64> {
346    let u0: Array1<f64> = &y - g0;
347    let u1: Array1<f64> = &y - g1;
348    d * &u1 / m - &(1.0 - d) * &u0 / &(1.0 - m)
349}
350
351#[cfg(test)]
352mod tests {
353    use super::*;
354    use crate::learner::{LogisticRegression, Ols};
355    use ndarray::{array, Array2};
356
357    #[test]
358    fn rejects_continuous_treatment_and_bad_clipping() {
359        let data = DmlData::new(
360            array![1.0, 2.0, 3.0],
361            array![0.0, 0.5, 1.0],
362            array![[1.0], [2.0], [3.0]],
363        )
364        .unwrap();
365        let irm = Irm::new(Ols::new(), LogisticRegression::default());
366        assert!(irm.fit(&data).is_err());
367        let binary =
368            DmlData::new(array![1.0, 2.0], array![0.0, 1.0], array![[1.0], [2.0]]).unwrap();
369        let bad = Irm::new(Ols::new(), LogisticRegression::default()).clipping_threshold(0.7);
370        assert!(bad.validate(&binary).is_err());
371    }
372
373    #[test]
374    fn normalized_weights_have_unit_means() {
375        let m = array![0.2, 0.5, 0.7, 0.4];
376        let d = array![1.0, 0.0, 1.0, 0.0];
377        let adj = normalize_propensity(&m, &d);
378        let n = 4.0;
379        let mean1: f64 = d.iter().zip(&adj).map(|(&a, &b)| a / b).sum::<f64>() / n;
380        let mean0: f64 = d
381            .iter()
382            .zip(&adj)
383            .map(|(&a, &b)| (1.0 - a) / (1.0 - b))
384            .sum::<f64>()
385            / n;
386        assert!((mean1 - 1.0).abs() < 1e-12 && (mean0 - 1.0).abs() < 1e-12);
387    }
388
389    #[test]
390    fn randomized_experiment_recovers_difference_in_means() {
391        // Constant propensity, no covariate signal: ATE ≈ treated − control mean.
392        let n = 200;
393        let x = Array2::from_shape_fn((n, 1), |(i, _)| ((i * 37) % 17) as f64);
394        let d: Array1<f64> = (0..n).map(|i| ((i * 7919) % 2) as f64).collect();
395        let y: Array1<f64> = (0..n).map(|i| 2.0 * d[i] + ((i * 31) % 7) as f64).collect();
396        let data = DmlData::new(y, d, x).unwrap();
397        let fit = Irm::new(Ols::new(), LogisticRegression::default())
398            .fit(&data)
399            .unwrap();
400        assert!((fit.coef()[0] - 2.0).abs() < 0.5, "ate = {}", fit.coef()[0]);
401    }
402}
dml_core/model/irm.rs

dml_core/model/
irm.rs
