dml_core/model/
iivm.rs

1//! Interactive IV model (IIVM) for a binary treatment `D` and a binary
2//! instrument `Z`, identifying the local average treatment effect (LATE) for
3//! compliers:
4//!
5//! ```text
6//! Y = ℓ(D, X) + ζ,       E[ζ | Z, X] = 0
7//! D = r(Z, X) + ν,       E[ν | Z, X] = 0
8//! Z = m(X) + V,          E[V | X] = 0
9//! θ = (E[g(1,X) − g(0,X)]) / (E[r(1,X) − r(0,X)])
10//! ```
11//!
12//! where `g(z, X) = E[Y | Z = z, X]` and `r(z, X) = P(D = 1 | Z = z, X)`.
13
14use ndarray::Array1;
15
16use super::irm::{ipw_residual, normalize_propensity};
17use super::{
18    dml_model_api, nuisance_api, LinearScoreModel, NuisanceLearners, NuisanceTarget, ScoreEstimate,
19    ScoreInput,
20};
21use crate::crossfit::{
22    cross_fit_predict_proba, cross_fit_predict_proba_subset, cross_fit_predict_subset, Fold,
23    SampleSplitting,
24};
25use crate::data::{is_binary, DmlData};
26use crate::error::{DmlError, Result};
27use crate::learner::{Classifier, Regressor};
28use crate::propensity::{adjust_propensity, Calibration};
29use crate::score::LinearScore;
30
31/// IIVM estimator of the LATE: `ml_g` learns `E[Y | Z, X]` (fitted separately
32/// on the `Z = 0` and `Z = 1` rows), `ml_m` the instrument propensity
33/// `P(Z = 1 | X)`, and `ml_r` the treatment probability `P(D = 1 | Z, X)`
34/// (again per instrument arm).
35pub struct Iivm {
36    nuisance: NuisanceLearners,
37    ml_g: Box<dyn Regressor>,
38    ml_m: Box<dyn Classifier>,
39    ml_r: Box<dyn Classifier>,
40    always_takers: bool,
41    never_takers: bool,
42    clipping_threshold: f64,
43    calibration: Calibration,
44    normalize_ipw: bool,
45    splitting: SampleSplitting,
46}
47
48impl Iivm {
49    pub fn new(
50        ml_g: impl Regressor + 'static,
51        ml_m: impl Classifier + 'static,
52        ml_r: impl Classifier + 'static,
53    ) -> Self {
54        Self::from_boxed(Box::new(ml_g), Box::new(ml_m), Box::new(ml_r))
55    }
56
57    pub fn from_boxed(
58        ml_g: Box<dyn Regressor>,
59        ml_m: Box<dyn Classifier>,
60        ml_r: Box<dyn Classifier>,
61    ) -> Self {
62        Self {
63            nuisance: NuisanceLearners::default(),
64            ml_g,
65            ml_m,
66            ml_r,
67            always_takers: true,
68            never_takers: true,
69            clipping_threshold: 1e-2,
70            calibration: Calibration::None,
71            normalize_ipw: false,
72            splitting: SampleSplitting::default(),
73        }
74    }
75
76    /// Whether always-takers (`D = 1` despite `Z = 0`) and never-takers
77    /// (`D = 0` despite `Z = 1`) can exist. Without always-takers
78    /// `r(0, X) = 0` is imposed instead of estimated; without never-takers
79    /// `r(1, X) = 1`. With one-sided non-compliance, switch the impossible
80    /// group off: its classifier would otherwise see only one class.
81    pub fn subgroups(mut self, always_takers: bool, never_takers: bool) -> Self {
82        self.always_takers = always_takers;
83        self.never_takers = never_takers;
84        self
85    }
86
87    /// Instrument propensities are clipped to `[t, 1 − t]` (default 0.01).
88    pub fn clipping_threshold(mut self, t: f64) -> Self {
89        self.clipping_threshold = t;
90        self
91    }
92    /// Calibrate the propensity scores before clipping (default: none), as
93    /// DoubleML's `ps_processor_config`.
94    pub fn calibration(mut self, calibration: Calibration) -> Self {
95        self.calibration = calibration;
96        self
97    }
98
99    /// Normalise the inverse-probability weights within each instrument arm.
100    ///
101    /// Note: DoubleML 0.11 normalises with the treatment `D` here; this
102    /// implementation uses the instrument `Z`, whose propensity the weights
103    /// invert. Results therefore differ from DoubleML when this is enabled.
104    /// Reported upstream as
105    /// [DoubleML#411](https://github.com/DoubleML/doubleml-for-py/issues/411).
106    pub fn normalize_ipw(mut self, normalize: bool) -> Self {
107        self.normalize_ipw = normalize;
108        self
109    }
110
111    dml_model_api!();
112    nuisance_api!();
113}
114
115impl LinearScoreModel for Iivm {
116    fn name(&self) -> String {
117        "IIVM (LATE)".into()
118    }
119
120    fn validate(&self, data: &DmlData) -> Result<()> {
121        self.nuisance
122            .check("IIVM", &["ml_g0", "ml_g1"], &["ml_m", "ml_r0", "ml_r1"])?;
123        if data.n_treat() != 1 || !data.is_binary_treatment(0) {
124            return Err(DmlError::InvalidArgument(
125                "IIVM requires exactly one binary (0/1) treatment".into(),
126            ));
127        }
128        if data.n_instruments() > 1 {
129            return Err(DmlError::InvalidArgument(
130                "IIVM needs exactly one instrument".into(),
131            ));
132        }
133        match data.z() {
134            Some(z) if is_binary(z) => {}
135            _ => {
136                return Err(DmlError::InvalidArgument(
137                    "IIVM requires a binary (0/1) instrument: use DmlData::with_instrument".into(),
138                ))
139            }
140        }
141        if !(self.clipping_threshold > 0.0 && self.clipping_threshold < 0.5) {
142            return Err(DmlError::InvalidArgument(
143                "clipping_threshold must be in (0, 0.5)".into(),
144            ));
145        }
146        Ok(())
147    }
148
149    fn estimate_score(&self, input: &ScoreInput<'_>, folds: &[Fold]) -> Result<ScoreEstimate> {
150        let ScoreInput { y, d, x, z, .. } = *input;
151        let z = z.ok_or_else(|| DmlError::InvalidArgument("IIVM requires an instrument".into()))?;
152        let z1: Vec<bool> = z.iter().map(|&v| v == 1.0).collect();
153        let z0: Vec<bool> = z1.iter().map(|b| !b).collect();
154        let n = y.len();
155
156        let nu = &self.nuisance;
157        let (ml_g, ml_r) = (self.ml_g.as_ref(), self.ml_r.as_ref());
158        let g0 = input.predictions_or("ml_g0", || {
159            cross_fit_predict_subset(nu.regressor("ml_g0", ml_g), x, y, folds, &z0)
160        })?;
161        let g1 = input.predictions_or("ml_g1", || {
162            cross_fit_predict_subset(nu.regressor("ml_g1", ml_g), x, y, folds, &z1)
163        })?;
164        let m_raw = input.predictions_or("ml_m", || {
165            cross_fit_predict_proba(nu.classifier("ml_m", self.ml_m.as_ref()), x, z, folds)
166        })?;
167        let m = adjust_propensity(m_raw, z, folds, self.clipping_threshold, self.calibration)?;
168        let r0 = if self.always_takers {
169            input.predictions_or("ml_r0", || {
170                cross_fit_predict_proba_subset(nu.classifier("ml_r0", ml_r), x, d, folds, &z0)
171            })?
172        } else {
173            Array1::zeros(n)
174        };
175        let r1 = if self.never_takers {
176            input.predictions_or("ml_r1", || {
177                cross_fit_predict_proba_subset(nu.classifier("ml_r1", ml_r), x, d, folds, &z1)
178            })?
179        } else {
180            Array1::ones(n)
181        };
182
183        let z = z.to_owned();
184        let m_adj = if self.normalize_ipw {
185            normalize_propensity(&m, &z)
186        } else {
187            m.clone()
188        };
189        // Reduced form (effect of Z on Y) over first stage (effect of Z on D).
190        let psi_b = &g1 - &g0 + &ipw_residual(y, &z, &g0, &g1, &m_adj);
191        let first_stage = &r1 - &r0 + &ipw_residual(d, &z, &r0, &r1, &m_adj);
192        // Without a subgroup, r is fixed (0 or 1) and its "target" is that constant.
193        let r0_target = if self.always_takers {
194            d.to_owned()
195        } else {
196            Array1::zeros(n)
197        };
198        let r1_target = if self.never_takers {
199            d.to_owned()
200        } else {
201            Array1::ones(n)
202        };
203        Ok(ScoreEstimate::new(
204            LinearScore::new(-first_stage, psi_b)?,
205            vec![
206                ("ml_g0", g0),
207                ("ml_g1", g1),
208                ("ml_m", m),
209                ("ml_r0", r0),
210                ("ml_r1", r1),
211            ],
212        )
213        .with_targets(vec![
214            NuisanceTarget::regression("ml_g0", y.to_owned()).only_where(|i| z[i] == 0.0),
215            NuisanceTarget::regression("ml_g1", y.to_owned()).only_where(|i| z[i] == 1.0),
216            NuisanceTarget::classification("ml_m", z.clone()),
217            NuisanceTarget::classification("ml_r0", r0_target).only_where(|i| z[i] == 0.0),
218            NuisanceTarget::classification("ml_r1", r1_target).only_where(|i| z[i] == 1.0),
219        ]))
220    }
221}
222
223#[cfg(test)]
224mod tests {
225    use super::*;
226    use crate::learner::{LogisticRegression, Ols};
227    use ndarray::array;
228
229    #[test]
230    fn requires_binary_treatment_and_instrument() {
231        let x = array![[1.0], [2.0], [3.0], [4.0]];
232        let base = DmlData::new(
233            array![1.0, 2.0, 3.0, 4.0],
234            array![0.0, 1.0, 0.0, 1.0],
235            x.clone(),
236        )
237        .unwrap();
238        let iivm = Iivm::new(
239            Ols::new(),
240            LogisticRegression::default(),
241            LogisticRegression::default(),
242        );
243        assert!(iivm.validate(&base).is_err()); // no instrument
244        let continuous_z = base
245            .clone()
246            .with_instrument(array![0.5, 1.0, 0.0, 1.0])
247            .unwrap();
248        assert!(iivm.validate(&continuous_z).is_err());
249        let ok = base.with_instrument(array![0.0, 1.0, 0.0, 1.0]).unwrap();
250        assert!(iivm.validate(&ok).is_ok());
251    }
252}
dml_core/model/iivm.rs

dml_core/model/
iivm.rs
