dml_core/model/
autodml.rs

1//! Automatic debiased machine learning (Chernozhukov, Newey and Singh, 2022):
2//! `θ = E[m(W; γ)]` for any linear functional `m` of the regression
3//! `γ(D, X) = E[Y | D, X]`, with the Riesz representer learned from the
4//! functional (see [`riesz`](crate::riesz)).
5
6use ndarray::{concatenate, Array1, Array2, Axis};
7
8use super::{
9    dml_model_api, CateSignal, LinearScoreModel, NuisanceTarget, ScoreEstimate, ScoreInput,
10};
11use crate::crossfit::{map_folds, Fold, SampleSplitting};
12use crate::data::gather;
13use crate::error::{DmlError, Result};
14use crate::learner::Regressor;
15use crate::riesz::{apply_terms, LinearFunctional, RieszLasso, RieszLearner, Term};
16use crate::score::LinearScore;
17use crate::sensitivity::{Nu2Parts, SensitivityParts};
18
19/// Auto-DML estimator of `θ = E[m(W; γ)]`.
20///
21/// `ml_g` learns `γ(D, X) = E[Y | D, X]` on `W = (D, X)` (the treatment in the
22/// first column), the Riesz learner (default [`RieszLasso`]) learns `α` from
23/// the functional, and the score is
24///
25/// ```text
26/// ψ = m(W; γ̂) + α̂(W)·(Y − γ̂(W)) − θ,
27/// ```
28///
29/// cross-fitted like every other model. Each treatment in the data gets the
30/// same functional applied to it (other treatments act as covariates).
31/// Sensitivity analysis uses `σ² = E[(Y − γ)²]` and `ν² = E[2m(W; α) − α²]`.
32///
33/// ```
34/// use dml_core::{AutoDml, AverageDerivative, DmlData, Ols, ShiftEffect};
35/// use ndarray::{Array1, Array2};
36///
37/// let n = 400;
38/// let x = Array2::from_shape_fn((n, 1), |(i, _)| ((i * 7) % 11) as f64 / 5.0);
39/// let d: Array1<f64> = (0..n).map(|i| x[[i, 0]] + ((i * 13) % 7) as f64 / 3.0).collect();
40/// let y: Array1<f64> = (0..n).map(|i| 0.5 * d[i] + x[[i, 0]].sin() + ((i * 17) % 5) as f64 / 4.0).collect();
41/// let data = DmlData::new(y, d, x)?;
42///
43/// // Average derivative E[∂γ/∂D]; here γ is linear in D with slope 0.5.
44/// let fit = AutoDml::new(Ols::new(), AverageDerivative::new()).fit(&data)?;
45/// assert!((fit.coef()[0] - 0.5).abs() < 0.1);
46///
47/// // Effect of raising every treatment by one unit.
48/// let shift = AutoDml::new(Ols::new(), ShiftEffect::new(1.0)).fit(&data)?;
49/// assert!((shift.coef()[0] - fit.coef()[0]).abs() < 1e-6);
50/// # Ok::<(), dml_core::DmlError>(())
51/// ```
52pub struct AutoDml {
53    ml_g: Box<dyn Regressor>,
54    riesz: Box<dyn RieszLearner>,
55    functional: Box<dyn LinearFunctional>,
56    splitting: SampleSplitting,
57}
58
59impl AutoDml {
60    pub fn new(
61        ml_g: impl Regressor + 'static,
62        functional: impl LinearFunctional + 'static,
63    ) -> Self {
64        Self::from_boxed(Box::new(ml_g), Box::new(functional))
65    }
66
67    /// Like [`AutoDml::new`] for boxed trait objects.
68    pub fn from_boxed(ml_g: Box<dyn Regressor>, functional: Box<dyn LinearFunctional>) -> Self {
69        Self {
70            ml_g,
71            riesz: Box::new(RieszLasso::default()),
72            functional,
73            splitting: SampleSplitting::default(),
74        }
75    }
76
77    /// Learner for the Riesz representer (default [`RieszLasso::default`]).
78    pub fn riesz(mut self, riesz: impl RieszLearner + 'static) -> Self {
79        self.riesz = Box::new(riesz);
80        self
81    }
82
83    pub fn riesz_boxed(mut self, riesz: Box<dyn RieszLearner>) -> Self {
84        self.riesz = riesz;
85        self
86    }
87
88    dml_model_api!();
89}
90
91/// Out-of-fold quantities of one fold, on its test rows.
92struct FoldResult {
93    g: Array1<f64>,
94    m_g: Array1<f64>,
95    alpha: Array1<f64>,
96    m_alpha: Array1<f64>,
97}
98
99fn rows_of(terms: &[Term], idx: &[usize]) -> Vec<Term> {
100    terms.iter().map(|t| t.rows(idx)).collect()
101}
102
103impl LinearScoreModel for AutoDml {
104    fn name(&self) -> String {
105        format!("Auto-DML ({})", self.functional.name())
106    }
107
108    fn estimate_score(&self, input: &ScoreInput<'_>, folds: &[Fold]) -> Result<ScoreEstimate> {
109        let ScoreInput { y, d, x, .. } = *input;
110        let n = y.len();
111        let w: Array2<f64> = concatenate![Axis(1), d.insert_axis(Axis(1)), x];
112        let terms = self.functional.terms(w.view())?;
113        if terms.is_empty() {
114            return Err(DmlError::InvalidArgument(
115                "the functional has no terms".into(),
116            ));
117        }
118        for t in &terms {
119            if t.w.dim() != w.dim() || t.weight.len() != n {
120                return Err(DmlError::DimensionMismatch("functional terms".into()));
121            }
122        }
123
124        let per_fold = map_folds(folds.len(), |k| -> Result<FoldResult> {
125            let fold = &folds[k];
126            let w_train = w.select(Axis(0), &fold.train);
127            let w_test = w.select(Axis(0), &fold.test);
128            let test_terms = rows_of(&terms, &fold.test);
129
130            let mut g = self.ml_g.fresh();
131            g.fit(w_train.view(), gather(y, &fold.train).view())?;
132            let g_test = g.predict(w_test.view())?;
133            let m_g = apply_terms(&test_terms, |wk| g.predict(wk))?;
134
135            let mut a = self.riesz.fresh();
136            a.fit(w_train.view(), &rows_of(&terms, &fold.train))?;
137            let alpha = a.predict(w_test.view())?;
138            let m_alpha = a.predict_functional(&test_terms)?;
139            Ok(FoldResult {
140                g: g_test,
141                m_g,
142                alpha,
143                m_alpha,
144            })
145        })?;
146
147        let mut g_hat = Array1::from_elem(n, f64::NAN);
148        let mut m_hat = Array1::from_elem(n, f64::NAN);
149        let mut alpha = Array1::from_elem(n, f64::NAN);
150        let mut m_alpha = Array1::from_elem(n, f64::NAN);
151        for (fold, r) in folds.iter().zip(per_fold) {
152            if r.g.len() != fold.test.len() || r.alpha.len() != fold.test.len() {
153                return Err(DmlError::DimensionMismatch(
154                    "learner returned wrong number of predictions".into(),
155                ));
156            }
157            for (j, &i) in fold.test.iter().enumerate() {
158                g_hat[i] = r.g[j];
159                m_hat[i] = r.m_g[j];
160                alpha[i] = r.alpha[j];
161                m_alpha[i] = r.m_alpha[j];
162            }
163        }
164        if [&g_hat, &m_hat, &alpha, &m_alpha]
165            .iter()
166            .any(|v| v.iter().any(|x| !x.is_finite()))
167        {
168            return Err(DmlError::InvalidArgument(
169                "folds do not cover every observation, or a learner returned a non-finite value"
170                    .into(),
171            ));
172        }
173
174        let resid: Array1<f64> = &y - &g_hat;
175        let psi_b: Array1<f64> = &m_hat + &(&alpha * &resid);
176        let parts = SensitivityParts {
177            outcome_residual: resid,
178            theta_factor: None,
179            riesz: alpha.clone(),
180            nu2: Nu2Parts::Elementwise(2.0 * &m_alpha - &alpha.mapv(|a| a * a)),
181        };
182        Ok(ScoreEstimate::new(
183            LinearScore::new(Array1::from_elem(n, -1.0), psi_b.clone())?,
184            vec![("ml_g", g_hat), ("riesz", alpha), ("functional", m_hat)],
185        )
186        .with_sensitivity(parts)
187        .with_targets(vec![NuisanceTarget::regression("ml_g", y.to_owned())])
188        // ψ_b is a per-unit signal of m(W; γ) (for the ATE the AIPW signal), so
189        // BLP / GATE inference project it on a basis, as for IRM.
190        .with_cate_signal(CateSignal {
191            signal: psi_b,
192            theta_factor: None,
193            basis_scale: None,
194            is_effect: true,
195        }))
196    }
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202    use crate::data::DmlData;
203    use crate::learner::Ols;
204    use crate::riesz::{Ate, AverageDerivative, PolicyEffect, RieszPenalty, ShiftEffect};
205    use rand::{Rng, SeedableRng};
206    use rand_chacha::ChaCha8Rng;
207    use rand_distr::StandardNormal;
208
209    /// `Y = θD + sin(X₁) + X₂ + ε`, `D = 0.5 X₁ + ν`: γ is linear in `D`, so the
210    /// average derivative and the per-unit shift effect are `θ`.
211    fn continuous(n: usize, seed: u64, theta: f64) -> DmlData {
212        let mut rng = ChaCha8Rng::seed_from_u64(seed);
213        let mut draw = || -> f64 { rng.sample(StandardNormal) };
214        let x = Array2::from_shape_fn((n, 2), |_| draw());
215        let d: Array1<f64> = (0..n).map(|i| 0.5 * x[[i, 0]] + draw()).collect();
216        let y: Array1<f64> = (0..n)
217            .map(|i| theta * d[i] + x[[i, 0]].sin() + x[[i, 1]] + draw())
218            .collect();
219        DmlData::new(y, d, x).unwrap()
220    }
221
222    #[test]
223    fn average_derivative_and_shift_recover_theta() {
224        let data = continuous(2000, 1, 0.7);
225        let fit = AutoDml::new(Ols::new(), AverageDerivative::new())
226            .fit(&data)
227            .unwrap();
228        assert!(
229            (fit.coef()[0] - 0.7).abs() < 4.0 * fit.se()[0],
230            "{}",
231            fit.coef()[0]
232        );
233        let shift = AutoDml::new(Ols::new(), ShiftEffect::new(0.5).per_unit(true))
234            .fit(&data)
235            .unwrap();
236        assert!((shift.coef()[0] - fit.coef()[0]).abs() < 1e-6);
237        // A policy that shifts everyone by 0.5 is the (unscaled) shift effect.
238        let policy = PolicyEffect::new(data.treatment(0).mapv(|d| d + 0.5));
239        let p = AutoDml::new(Ols::new(), policy).fit(&data).unwrap();
240        let s = AutoDml::new(Ols::new(), ShiftEffect::new(0.5))
241            .fit(&data)
242            .unwrap();
243        assert!((p.coef()[0] - s.coef()[0]).abs() < 1e-10);
244        assert!(fit.prediction("riesz", 0, 0).is_some());
245    }
246
247    #[test]
248    fn riesz_correction_removes_the_bias_of_a_misspecified_regression() {
249        // γ by OLS without the D·X₁ interaction is wrong; the learned Riesz
250        // representer corrects it. Y = D(1 + X₁) + X₁ + ε has average derivative 1.
251        let mut rng = ChaCha8Rng::seed_from_u64(7);
252        let n = 4000;
253        let mut draw = || -> f64 { rng.sample(StandardNormal) };
254        let x = Array2::from_shape_fn((n, 1), |_| draw());
255        let d: Array1<f64> = (0..n).map(|i| x[[i, 0]] + draw()).collect();
256        let y: Array1<f64> = (0..n)
257            .map(|i| d[i] * (1.0 + x[[i, 0]]) + x[[i, 0]] + draw())
258            .collect();
259        let data = DmlData::new(y, d, x).unwrap();
260        let fit = AutoDml::new(Ols::new(), AverageDerivative::new())
261            .fit(&data)
262            .unwrap();
263        assert!(
264            (fit.coef()[0] - 1.0).abs() < 4.0 * fit.se()[0],
265            "{} ± {}",
266            fit.coef()[0],
267            fit.se()[0]
268        );
269        let plug_in = fit.prediction("functional", 0, 0).unwrap().mean().unwrap();
270        assert!((plug_in - 1.0).abs() > (fit.coef()[0] - 1.0).abs());
271    }
272
273    #[test]
274    fn ate_is_close_to_irm() {
275        let mut rng = ChaCha8Rng::seed_from_u64(3);
276        let n = 3000;
277        let mut draw = || -> f64 { rng.sample(StandardNormal) };
278        let x = Array2::from_shape_fn((n, 2), |_| draw());
279        let mut u = ChaCha8Rng::seed_from_u64(4);
280        let d: Array1<f64> = (0..n)
281            .map(|i| {
282                let p = 1.0 / (1.0 + (-0.8 * x[[i, 0]]).exp());
283                f64::from(u.gen::<f64>() < p)
284            })
285            .collect();
286        let y: Array1<f64> = (0..n)
287            .map(|i| 0.5 * d[i] + x[[i, 0]] + 0.5 * x[[i, 1]] + draw())
288            .collect();
289        let data = DmlData::new(y, d, x).unwrap();
290        let auto = AutoDml::new(Ols::new(), Ate).fit(&data).unwrap();
291        let irm = crate::Irm::new(Ols::new(), crate::LogisticRegression::default())
292            .fit(&data)
293            .unwrap();
294        assert!((auto.coef()[0] - irm.coef()[0]).abs() < 0.5 * irm.se()[0]);
295        assert!((auto.coef()[0] - 0.5).abs() < 4.0 * auto.se()[0]);
296        let s = auto
297            .sensitivity_analysis(&crate::SensitivityParams::default())
298            .unwrap();
299        assert!(s.theta_lower[0] < auto.coef()[0]);
300    }
301
302    #[test]
303    fn fixed_penalty_and_errors() {
304        let data = continuous(300, 2, 0.5);
305        let fit = AutoDml::new(Ols::new(), AverageDerivative::new())
306            .riesz(RieszLasso::new().penalty(RieszPenalty::Fixed(0.0)))
307            .fit(&data)
308            .unwrap();
309        assert!(fit.coef()[0].is_finite());
310        assert!(AutoDml::new(Ols::new(), Ate).fit(&data).is_err());
311    }
312
313    #[test]
314    fn coverage_of_the_average_derivative() {
315        let reps = 120;
316        let mut hits = 0;
317        for r in 0..reps {
318            let data = continuous(400, 100 + r, 0.7);
319            let fit = AutoDml::new(Ols::new(), AverageDerivative::new())
320                .seed(r)
321                .fit(&data)
322                .unwrap();
323            let ci = fit.confint(0.95).unwrap();
324            if ci[[0, 0]] <= 0.7 && 0.7 <= ci[[0, 1]] {
325                hits += 1;
326            }
327        }
328        let cover = hits as f64 / reps as f64;
329        assert!(cover > 0.88, "coverage {cover}");
330    }
331}
dml_core/model/autodml.rs

dml_core/model/
autodml.rs
