dml_core/model/
pliv.rs

1//! Partially linear IV regression (PLIV) with one instrument `Z`:
2//!
3//! ```text
4//! Y = θ·D + g(X) + ε,   E[ε | Z, X] = 0
5//! Z = m(X) + v,         E[v | X] = 0
6//! ```
7
8use ndarray::Array1;
9
10use super::{dml_model_api, LinearScoreModel, ScoreEstimate, ScoreInput};
11use crate::crossfit::{cross_fit_predict, Fold, SampleSplitting};
12use crate::data::DmlData;
13use crate::error::{DmlError, Result};
14use crate::learner::Regressor;
15use crate::score::LinearScore;
16
17/// PLIV estimator with the partialling-out score
18/// `ψ = (Y − ℓ(X) − θ(D − r(X)))·(Z − m(X))`.
19///
20/// `ml_l` learns `E[Y | X]`, `ml_m` learns `E[Z | X]`, `ml_r` learns `E[D | X]`.
21pub struct Pliv {
22    ml_l: Box<dyn Regressor>,
23    ml_m: Box<dyn Regressor>,
24    ml_r: Box<dyn Regressor>,
25    splitting: SampleSplitting,
26}
27
28impl Pliv {
29    pub fn new(
30        ml_l: impl Regressor + 'static,
31        ml_m: impl Regressor + 'static,
32        ml_r: impl Regressor + 'static,
33    ) -> Self {
34        Self::from_boxed(Box::new(ml_l), Box::new(ml_m), Box::new(ml_r))
35    }
36
37    pub fn from_boxed(
38        ml_l: Box<dyn Regressor>,
39        ml_m: Box<dyn Regressor>,
40        ml_r: Box<dyn Regressor>,
41    ) -> Self {
42        Self {
43            ml_l,
44            ml_m,
45            ml_r,
46            splitting: SampleSplitting::default(),
47        }
48    }
49
50    dml_model_api!();
51}
52
53impl LinearScoreModel for Pliv {
54    fn name(&self) -> String {
55        "PLIV (partialling out)".into()
56    }
57
58    fn validate(&self, data: &DmlData) -> Result<()> {
59        if data.z().is_none() {
60            return Err(DmlError::InvalidArgument(
61                "PLIV requires an instrument: use DmlData::with_instrument".into(),
62            ));
63        }
64        Ok(())
65    }
66
67    fn estimate_score(&self, input: &ScoreInput<'_>, folds: &[Fold]) -> Result<ScoreEstimate> {
68        let ScoreInput { y, d, x, z } = *input;
69        let z = z.ok_or_else(|| DmlError::InvalidArgument("PLIV requires an instrument".into()))?;
70        let l_hat = cross_fit_predict(self.ml_l.as_ref(), x, y, folds)?;
71        let m_hat = cross_fit_predict(self.ml_m.as_ref(), x, z, folds)?;
72        let r_hat = cross_fit_predict(self.ml_r.as_ref(), x, d, folds)?;
73        let u: Array1<f64> = &y - &l_hat;
74        let v: Array1<f64> = &z - &m_hat;
75        let w: Array1<f64> = &d - &r_hat;
76        Ok(ScoreEstimate {
77            score: LinearScore::new(-(&w * &v), &v * &u)?,
78            predictions: vec![("ml_l", l_hat), ("ml_m", m_hat), ("ml_r", r_hat)],
79        })
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86    use crate::learner::Ols;
87    use ndarray::array;
88
89    #[test]
90    fn requires_instrument() {
91        let data = DmlData::new(array![1.0, 2.0], array![0.0, 1.0], array![[1.0], [2.0]]).unwrap();
92        assert!(Pliv::new(Ols::new(), Ols::new(), Ols::new())
93            .fit(&data)
94            .is_err());
95    }
96}
dml_core/model/pliv.rs

dml_core/model/
pliv.rs
