dml_core/model/
pliv.rs

1//! Partially linear IV regression (PLIV) with one or several instruments `Z`:
2//!
3//! ```text
4//! Y = θ·D + g(X) + ε,   E[ε | Z, X] = 0
5//! ```
6//!
7//! As DoubleML's `DoubleMLPLIV` (Chernozhukov et al., 2018), in three
8//! variants ([`PlivVariant`]): partialling `X` out of `Y`, `D` and `Z`
9//! (default), partialling `Z` out of `D` (`E[D | X, Z]` as optimal
10//! instrument), or both.
11
12use ndarray::{concatenate, Array1, Array2, Axis};
13
14use super::{dml_model_api, LinearScoreModel, Name, NuisanceTarget, ScoreEstimate, ScoreInput};
15use crate::compat::np_mean;
16use crate::crossfit::{cross_fit_predict, Fold, SampleSplitting};
17use crate::data::{gather, DmlData};
18use crate::error::{DmlError, Result};
19use crate::learner::{Ols, Regressor};
20use crate::score::LinearScore;
21
22/// Score of [`Pliv`].
23#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
24pub enum PlivScore {
25    /// `ψ = (Y − ℓ(X) − θ(D − r(X)))·(Z − m(X))`.
26    #[default]
27    PartiallingOut,
28    /// `ψ = (Y − θD − g(X))·(Z − m(X))`; one instrument, [`PlivVariant::PartialX`] only.
29    IvType,
30}
31
32/// What [`Pliv`] partials out.
33#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
34pub enum PlivVariant {
35    /// Partial `X` out of `Y`, `D` and every instrument. With several
36    /// instruments the treatment residual is projected on the instrument
37    /// residuals (two-stage least squares). Learners `ml_l` (`E[Y | X]`),
38    /// `ml_m` (`E[Z | X]`, one fit per instrument), `ml_r` (`E[D | X]`).
39    #[default]
40    PartialX,
41    /// Use `E[D | X, Z]` (`ml_r`) as instrument; no other learner.
42    PartialZ,
43    /// Partial out both: `ml_m` learns `E[D | X, Z]`, `ml_r` the projection of
44    /// its in-sample predictions on `X`, `ml_l` learns `E[Y | X]`.
45    PartialXZ,
46}
47
48/// Partially linear IV regression. See [`PlivVariant`] and [`PlivScore`].
49pub struct Pliv {
50    ml_l: Option<Box<dyn Regressor>>,
51    ml_m: Option<Box<dyn Regressor>>,
52    ml_r: Box<dyn Regressor>,
53    ml_g: Option<Box<dyn Regressor>>,
54    score: PlivScore,
55    variant: PlivVariant,
56    splitting: SampleSplitting,
57}
58
59impl Pliv {
60    /// Partial `X` out ([`PlivVariant::PartialX`]): `ml_l` learns `E[Y | X]`,
61    /// `ml_m` `E[Z | X]`, `ml_r` `E[D | X]`.
62    pub fn new(
63        ml_l: impl Regressor + 'static,
64        ml_m: impl Regressor + 'static,
65        ml_r: impl Regressor + 'static,
66    ) -> Self {
67        Self::from_boxed(Box::new(ml_l), Box::new(ml_m), Box::new(ml_r))
68    }
69
70    pub fn from_boxed(
71        ml_l: Box<dyn Regressor>,
72        ml_m: Box<dyn Regressor>,
73        ml_r: Box<dyn Regressor>,
74    ) -> Self {
75        Self {
76            ml_l: Some(ml_l),
77            ml_m: Some(ml_m),
78            ml_r,
79            ml_g: None,
80            score: PlivScore::default(),
81            variant: PlivVariant::PartialX,
82            splitting: SampleSplitting::default(),
83        }
84    }
85
86    /// Partial `Z` out ([`PlivVariant::PartialZ`]): `ml_r` learns `E[D | X, Z]`.
87    pub fn partial_z(ml_r: impl Regressor + 'static) -> Self {
88        Self::partial_z_boxed(Box::new(ml_r))
89    }
90
91    pub fn partial_z_boxed(ml_r: Box<dyn Regressor>) -> Self {
92        Self {
93            ml_l: None,
94            ml_m: None,
95            ml_r,
96            ml_g: None,
97            score: PlivScore::default(),
98            variant: PlivVariant::PartialZ,
99            splitting: SampleSplitting::default(),
100        }
101    }
102
103    /// Partial out `X` and `Z` ([`PlivVariant::PartialXZ`]): `ml_l` learns
104    /// `E[Y | X]`, `ml_m` `E[D | X, Z]`, `ml_r` the projection of `ml_m`'s
105    /// in-sample predictions on `X`.
106    pub fn partial_xz(
107        ml_l: impl Regressor + 'static,
108        ml_m: impl Regressor + 'static,
109        ml_r: impl Regressor + 'static,
110    ) -> Self {
111        Self::from_boxed(Box::new(ml_l), Box::new(ml_m), Box::new(ml_r))
112            .variant(PlivVariant::PartialXZ)
113    }
114
115    /// Choose the variant (learners as for [`Pliv::new`]; [`PlivVariant::PartialZ`]
116    /// uses only `ml_r`, learning `E[D | X, Z]`).
117    pub fn variant(mut self, variant: PlivVariant) -> Self {
118        self.variant = variant;
119        self
120    }
121
122    pub fn score(mut self, score: PlivScore) -> Self {
123        self.score = score;
124        self
125    }
126
127    /// Learner for `E[Y − θD | X]` with the IV-type score (default: a copy of `ml_l`).
128    pub fn ml_g(self, ml_g: impl Regressor + 'static) -> Self {
129        self.ml_g_boxed(Box::new(ml_g))
130    }
131
132    pub fn ml_g_boxed(mut self, ml_g: Box<dyn Regressor>) -> Self {
133        self.ml_g = Some(ml_g);
134        self
135    }
136
137    dml_model_api!();
138
139    fn learner(&self, l: &Option<Box<dyn Regressor>>, name: &str) -> Result<Box<dyn Regressor>> {
140        l.as_ref().map(|l| l.fresh()).ok_or_else(|| {
141            DmlError::InvalidArgument(format!("this PLIV variant needs a learner {name}"))
142        })
143    }
144}
145
146impl LinearScoreModel for Pliv {
147    fn name(&self) -> String {
148        let variant = match self.variant {
149            PlivVariant::PartialX => "",
150            PlivVariant::PartialZ => ", partial Z",
151            PlivVariant::PartialXZ => ", partial XZ",
152        };
153        let score = match self.score {
154            PlivScore::PartiallingOut => "partialling out",
155            PlivScore::IvType => "IV-type",
156        };
157        format!("PLIV ({score}{variant})")
158    }
159
160    fn validate(&self, data: &DmlData) -> Result<()> {
161        let k = data.n_instruments();
162        if k == 0 {
163            return Err(DmlError::InvalidArgument(
164                "PLIV requires an instrument: use DmlData::with_instrument".into(),
165            ));
166        }
167        if self.score == PlivScore::IvType && (k > 1 || self.variant != PlivVariant::PartialX) {
168            return Err(DmlError::InvalidArgument(
169                "the IV-type score needs one instrument and the partial X variant".into(),
170            ));
171        }
172        Ok(())
173    }
174
175    fn estimate_score(&self, input: &ScoreInput<'_>, folds: &[Fold]) -> Result<ScoreEstimate> {
176        let ScoreInput { y, d, x, .. } = *input;
177        let z = input
178            .instruments
179            .ok_or_else(|| DmlError::InvalidArgument("PLIV requires an instrument".into()))?;
180        match self.variant {
181            PlivVariant::PartialX => {
182                let ml_l = self.learner(&self.ml_l, "ml_l")?;
183                let ml_m = self.learner(&self.ml_m, "ml_m")?;
184                let l_hat = input
185                    .predictions_or("ml_l", || cross_fit_predict(ml_l.as_ref(), x, y, folds))?;
186                let r_hat = input.predictions_or("ml_r", || {
187                    cross_fit_predict(self.ml_r.as_ref(), x, d, folds)
188                })?;
189                let u: Array1<f64> = &y - &l_hat;
190                let w: Array1<f64> = &d - &r_hat;
191                let mut predictions: Vec<(Name, Array1<f64>)> = vec![("ml_l".into(), l_hat)];
192                let mut targets = vec![NuisanceTarget::regression("ml_l", y.to_owned())];
193                if z.ncols() == 1 {
194                    let z = z.column(0);
195                    let m_hat = input
196                        .predictions_or("ml_m", || cross_fit_predict(ml_m.as_ref(), x, z, folds))?;
197                    let v: Array1<f64> = &z - &m_hat;
198                    predictions.push(("ml_m".into(), m_hat));
199                    targets.push(NuisanceTarget::regression("ml_m", z.to_owned()));
200                    let score = match self.score {
201                        PlivScore::PartiallingOut => LinearScore::new(-(&w * &v), &v * &u)?,
202                        PlivScore::IvType => {
203                            // Initial θ from the partialling-out score.
204                            let psi_a: Vec<f64> = w.iter().zip(&v).map(|(w, v)| -(w * v)).collect();
205                            let psi_b: Vec<f64> = v.iter().zip(&u).map(|(v, u)| v * u).collect();
206                            let theta_init = -np_mean(&psi_b) / np_mean(&psi_a);
207                            let target: Array1<f64> = &y - &(theta_init * &d);
208                            let ml_g = match &self.ml_g {
209                                Some(g) => g.fresh(),
210                                None => ml_l.fresh(),
211                            };
212                            let g_hat = input.predictions_or("ml_g", || {
213                                cross_fit_predict(ml_g.as_ref(), x, target.view(), folds)
214                            })?;
215                            let score = LinearScore::new(-(&v * &d), &v * &(&y - &g_hat))?;
216                            predictions.push(("ml_g".into(), g_hat));
217                            targets.push(NuisanceTarget::regression("ml_g", target));
218                            score
219                        }
220                    };
221                    predictions.push(("ml_r".into(), r_hat.clone()));
222                    targets.push(NuisanceTarget::regression("ml_r", d.to_owned()));
223                    return Ok(ScoreEstimate::new(score, predictions).with_targets(targets));
224                }
225                // Several instruments: residualise each, then project the
226                // treatment residual on them (OLS with intercept).
227                let mut v = Array2::zeros(z.dim());
228                for (j, col) in z.columns().into_iter().enumerate() {
229                    let name = format!("ml_m_z{}", j + 1);
230                    let m_hat = input.predictions_or(&name, || {
231                        cross_fit_predict(ml_m.fresh().as_ref(), x, col, folds)
232                    })?;
233                    v.column_mut(j).assign(&(&col - &m_hat));
234                    targets.push(NuisanceTarget::regression(name.clone(), col.to_owned()));
235                    predictions.push((name.into(), m_hat));
236                }
237                let mut proj = Ols::new();
238                proj.fit(v.view(), w.view())?;
239                let r_tilde = proj.predict(v.view())?;
240                predictions.push(("ml_r".into(), r_hat.clone()));
241                targets.push(NuisanceTarget::regression("ml_r", d.to_owned()));
242                Ok(ScoreEstimate::new(
243                    LinearScore::new(-(&w * &r_tilde), &r_tilde * &u)?,
244                    predictions,
245                )
246                .with_targets(targets))
247            }
248            PlivVariant::PartialZ => {
249                let xz = concatenate![Axis(1), x, z];
250                let r_hat = cross_fit_predict(self.ml_r.as_ref(), xz.view(), d, folds)?;
251                Ok(ScoreEstimate::new(
252                    LinearScore::new(-(&r_hat * &d), &r_hat * &y)?,
253                    vec![("ml_r", r_hat)],
254                )
255                .with_targets(vec![NuisanceTarget::regression("ml_r", d.to_owned())]))
256            }
257            PlivVariant::PartialXZ => {
258                let ml_l = self.learner(&self.ml_l, "ml_l")?;
259                let ml_m = self.learner(&self.ml_m, "ml_m")?;
260                let xz = concatenate![Axis(1), x, z];
261                let l_hat = cross_fit_predict(ml_l.as_ref(), x, y, folds)?;
262                // ml_m on (X, Z), predicting both its training rows (the
263                // targets of ml_r) and the test rows; ml_r on X.
264                let n = y.len();
265                let fold_fit = |fold: &Fold| -> Result<(Array1<f64>, Array1<f64>)> {
266                    let mut m = ml_m.fresh();
267                    let x_train = xz.select(Axis(0), &fold.train);
268                    m.fit(x_train.view(), gather(d, &fold.train).view())?;
269                    let train_pred = m.predict(x_train.view())?;
270                    let test_pred = m.predict(xz.select(Axis(0), &fold.test).view())?;
271                    let mut r = self.ml_r.fresh();
272                    r.fit(x.select(Axis(0), &fold.train).view(), train_pred.view())?;
273                    let tilde = r.predict(x.select(Axis(0), &fold.test).view())?;
274                    Ok((test_pred, tilde))
275                };
276                let per_fold: Vec<(Array1<f64>, Array1<f64>)> =
277                    folds.iter().map(fold_fit).collect::<Result<_>>()?;
278                let mut m_hat = Array1::from_elem(n, f64::NAN);
279                let mut m_tilde = Array1::from_elem(n, f64::NAN);
280                for (fold, (pred, tilde)) in folds.iter().zip(&per_fold) {
281                    for (k, &i) in fold.test.iter().enumerate() {
282                        m_hat[i] = pred[k];
283                        m_tilde[i] = tilde[k];
284                    }
285                }
286                let u: Array1<f64> = &y - &l_hat;
287                let w: Array1<f64> = &d - &m_tilde;
288                let instrument: Array1<f64> = &m_hat - &m_tilde;
289                Ok(ScoreEstimate::new(
290                    LinearScore::new(-(&w * &instrument), &instrument * &u)?,
291                    vec![("ml_l", l_hat), ("ml_m", m_hat), ("ml_r", m_tilde)],
292                )
293                .with_targets(vec![
294                    NuisanceTarget::regression("ml_l", y.to_owned()),
295                    NuisanceTarget::regression("ml_m", d.to_owned()),
296                ]))
297            }
298        }
299    }
300}
301
302#[cfg(test)]
303mod tests {
304    use super::*;
305    use crate::learner::Ols;
306    use ndarray::array;
307
308    #[test]
309    fn requires_instrument() {
310        let data = DmlData::new(array![1.0, 2.0], array![0.0, 1.0], array![[1.0], [2.0]]).unwrap();
311        assert!(Pliv::new(Ols::new(), Ols::new(), Ols::new())
312            .fit(&data)
313            .is_err());
314    }
315}
dml_core/model/pliv.rs

dml_core/model/
pliv.rs
