dml-core 0.1.0

Double / debiased machine learning for causal inference in Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97

//! DML models with linear scores: [`Plr`], [`Irm`], [`Pliv`], and the
//! [`LinearScoreModel`] trait for adding your own.

use ndarray::{Array1, ArrayView1, ArrayView2};

use crate::crossfit::Fold;
use crate::data::DmlData;
use crate::error::Result;
use crate::score::LinearScore;

pub(crate) mod irm;
mod pliv;
mod plr;

pub use irm::{Irm, IrmScore};
pub use pliv::Pliv;
pub use plr::{Plr, PlrScore};

/// The data a model sees when estimating the effect of one treatment.
#[derive(Debug, Clone, Copy)]
pub struct ScoreInput<'a> {
    pub y: ArrayView1<'a, f64>,
    /// The treatment being estimated.
    pub d: ArrayView1<'a, f64>,
    /// Covariates, including any other treatments.
    pub x: ArrayView2<'a, f64>,
    pub z: Option<ArrayView1<'a, f64>>,
}

/// Score elements plus the cross-fitted nuisance predictions behind them.
#[derive(Debug, Clone)]
pub struct ScoreEstimate {
    pub score: LinearScore,
    /// `(name, out-of-fold predictions)` — e.g. `("ml_l", ℓ̂)`.
    pub predictions: Vec<(&'static str, Array1<f64>)>,
}

/// A DML model whose orthogonal score is linear in θ.
///
/// Implement this to plug a new score into the shared machinery (repeated
/// cross-fitting, aggregation, inference, bootstrap) via [`fit_model`](crate::fit_model).
pub trait LinearScoreModel {
    /// Label for summaries, e.g. `"PLR (partialling out)"`.
    fn name(&self) -> String;

    /// Reject data the model cannot handle (called once before fitting).
    fn validate(&self, _data: &DmlData) -> Result<()> {
        Ok(())
    }

    /// Cross-fit the nuisances on `folds` and return the score elements.
    fn estimate_score(&self, input: &ScoreInput<'_>, folds: &[Fold]) -> Result<ScoreEstimate>;
}

/// Builder methods shared by every model: sample-splitting settings and `fit`.
///
/// (`macro_rules!` generates code at compile time; here it pastes the same
/// methods into each model's `impl` block instead of repeating them.)
macro_rules! dml_model_api {
    () => {
        /// Number of cross-fitting folds (default 5).
        pub fn n_folds(mut self, n_folds: usize) -> Self {
            self.splitting.n_folds = n_folds;
            self
        }

        /// Number of independent sample splits (default 1). Estimates are
        /// aggregated across repetitions by the median, as in DoubleML.
        pub fn n_rep(mut self, n_rep: usize) -> Self {
            self.splitting.n_rep = n_rep;
            self
        }

        /// Seed for the sample splits (default 0).
        pub fn seed(mut self, seed: u64) -> Self {
            self.splitting.seed = seed;
            self
        }

        /// Draw seeded sample splits and fit.
        pub fn fit(&self, data: &$crate::DmlData) -> $crate::Result<$crate::DmlFit> {
            let splits = self.splitting.draw(data.n_obs())?;
            $crate::fit_model(self, data, &splits)
        }

        /// Fit with caller-supplied splits, one `Vec<Fold>` per repetition
        /// (e.g. to reproduce another implementation exactly).
        pub fn fit_with_folds(
            &self,
            data: &$crate::DmlData,
            splits: &[Vec<$crate::Fold>],
        ) -> $crate::Result<$crate::DmlFit> {
            $crate::fit_model(self, data, splits)
        }
    };
}
pub(crate) use dml_model_api;