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

dml_core/model/
mod.rs
