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

dml_core/model/
mod.rs
