dml_core/model/
mod.rs

1//! DML models with linear scores: [`Plr`], [`Irm`], [`Pliv`], and the
2//! [`LinearScoreModel`] trait for adding your own.
3
4use ndarray::{Array1, ArrayView1, ArrayView2};
5
6use crate::crossfit::Fold;
7use crate::data::DmlData;
8use crate::error::Result;
9use crate::score::LinearScore;
10
11pub(crate) mod irm;
12mod pliv;
13mod plr;
14
15pub use irm::{Irm, IrmScore};
16pub use pliv::Pliv;
17pub use plr::{Plr, PlrScore};
18
19/// The data a model sees when estimating the effect of one treatment.
20#[derive(Debug, Clone, Copy)]
21pub struct ScoreInput<'a> {
22    pub y: ArrayView1<'a, f64>,
23    /// The treatment being estimated.
24    pub d: ArrayView1<'a, f64>,
25    /// Covariates, including any other treatments.
26    pub x: ArrayView2<'a, f64>,
27    pub z: Option<ArrayView1<'a, f64>>,
28}
29
30/// Score elements plus the cross-fitted nuisance predictions behind them.
31#[derive(Debug, Clone)]
32pub struct ScoreEstimate {
33    pub score: LinearScore,
34    /// `(name, out-of-fold predictions)` — e.g. `("ml_l", ℓ̂)`.
35    pub predictions: Vec<(&'static str, Array1<f64>)>,
36}
37
38/// A DML model whose orthogonal score is linear in θ.
39///
40/// Implement this to plug a new score into the shared machinery (repeated
41/// cross-fitting, aggregation, inference, bootstrap) via [`fit_model`](crate::fit_model).
42pub trait LinearScoreModel {
43    /// Label for summaries, e.g. `"PLR (partialling out)"`.
44    fn name(&self) -> String;
45
46    /// Reject data the model cannot handle (called once before fitting).
47    fn validate(&self, _data: &DmlData) -> Result<()> {
48        Ok(())
49    }
50
51    /// Cross-fit the nuisances on `folds` and return the score elements.
52    fn estimate_score(&self, input: &ScoreInput<'_>, folds: &[Fold]) -> Result<ScoreEstimate>;
53}
54
55/// Builder methods shared by every model: sample-splitting settings and `fit`.
56///
57/// (`macro_rules!` generates code at compile time; here it pastes the same
58/// methods into each model's `impl` block instead of repeating them.)
59macro_rules! dml_model_api {
60    () => {
61        /// Number of cross-fitting folds (default 5).
62        pub fn n_folds(mut self, n_folds: usize) -> Self {
63            self.splitting.n_folds = n_folds;
64            self
65        }
66
67        /// Number of independent sample splits (default 1). Estimates are
68        /// aggregated across repetitions by the median, as in DoubleML.
69        pub fn n_rep(mut self, n_rep: usize) -> Self {
70            self.splitting.n_rep = n_rep;
71            self
72        }
73
74        /// Seed for the sample splits (default 0).
75        pub fn seed(mut self, seed: u64) -> Self {
76            self.splitting.seed = seed;
77            self
78        }
79
80        /// Draw seeded sample splits and fit.
81        pub fn fit(&self, data: &$crate::DmlData) -> $crate::Result<$crate::DmlFit> {
82            let splits = self.splitting.draw(data.n_obs())?;
83            $crate::fit_model(self, data, &splits)
84        }
85
86        /// Fit with caller-supplied splits, one `Vec<Fold>` per repetition
87        /// (e.g. to reproduce another implementation exactly).
88        pub fn fit_with_folds(
89            &self,
90            data: &$crate::DmlData,
91            splits: &[Vec<$crate::Fold>],
92        ) -> $crate::Result<$crate::DmlFit> {
93            $crate::fit_model(self, data, splits)
94        }
95    };
96}
97pub(crate) use dml_model_api;
dml_core/model/mod.rs

dml_core/model/
mod.rs
