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

dml_core/model/
iivm.rs
