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, Array2, ArrayView1, ArrayView2};
5
6use crate::crossfit::Fold;
7use crate::data::DmlData;
8use crate::error::{DmlError, Result};
9use crate::score::LinearScore;
10use crate::sensitivity::SensitivityParts;
11
12mod apo;
13pub(crate) mod did;
14mod iivm;
15pub(crate) mod irm;
16mod lplr;
17mod nonlinear;
18mod pliv;
19mod plpr;
20mod plr;
21mod quantile;
22mod rdflex;
23mod ssm;
24
25pub use apo::{Apo, Apos, AposFit};
26pub use did::{Did, DidCs, DidScore};
27pub use iivm::Iivm;
28pub use irm::{Irm, IrmScore};
29pub use lplr::{Lplr, LplrScore, TreatmentLearner};
30pub use pliv::{Pliv, PlivScore, PlivVariant};
31pub use plpr::{PanelApproach, Plpr, StaticPanel};
32pub use plr::{Plr, PlrScore};
33pub use quantile::{Cvar, Lpq, Pq, Qte, QteFit};
34pub use rdflex::{FirstStage, RdFlex, RdFlexFit, RdKernel};
35pub use ssm::{Ssm, SsmScore};
36
37/// The data a model sees when estimating the effect of one treatment.
38///
39/// `#[non_exhaustive]`: fields may be added in future versions, so code outside
40/// this crate reads the fields but cannot construct the struct.
41#[derive(Debug, Clone, Copy)]
42#[non_exhaustive]
43pub struct ScoreInput<'a> {
44    pub y: ArrayView1<'a, f64>,
45    /// The treatment being estimated.
46    pub d: ArrayView1<'a, f64>,
47    /// Covariates, including any other treatments.
48    pub x: ArrayView2<'a, f64>,
49    pub z: Option<ArrayView1<'a, f64>>,
50    /// All instruments (several for [`Pliv`]).
51    pub instruments: Option<ArrayView2<'a, f64>>,
52    /// 0/1 time-period indicator (repeated cross-sections).
53    pub t: Option<ArrayView1<'a, f64>>,
54    /// 0/1 selection indicator (sample selection models).
55    pub s: Option<ArrayView1<'a, f64>>,
56    /// Index of the sample-splitting repetition being estimated.
57    pub rep: usize,
58    /// Index of the treatment being estimated.
59    pub treatment: usize,
60    /// User-supplied nuisance predictions ([`DmlData::with_external_predictions`]).
61    pub external: Option<&'a ExternalPredictions>,
62}
63
64impl ScoreInput<'_> {
65    /// The user-supplied out-of-fold predictions of learner `name` for this
66    /// treatment and repetition if there are any, else `fit()`.
67    pub fn predictions_or(
68        &self,
69        name: &str,
70        fit: impl FnOnce() -> Result<Array1<f64>>,
71    ) -> Result<Array1<f64>> {
72        external_or(self.external_ref(), self.y.len(), name, fit)
73    }
74
75    pub(crate) fn external_ref(&self) -> Option<ExternalRef<'_>> {
76        self.external.map(|preds| ExternalRef {
77            preds,
78            treatment: self.treatment,
79            rep: self.rep,
80        })
81    }
82}
83
84/// External predictions of one treatment and repetition.
85#[derive(Debug, Clone, Copy)]
86pub(crate) struct ExternalRef<'a> {
87    pub preds: &'a ExternalPredictions,
88    pub treatment: usize,
89    pub rep: usize,
90}
91
92/// External predictions of learner `name` (`n` rows) if supplied, else `fit()`.
93pub(crate) fn external_or(
94    ext: Option<ExternalRef<'_>>,
95    n: usize,
96    name: &str,
97    fit: impl FnOnce() -> Result<Array1<f64>>,
98) -> Result<Array1<f64>> {
99    match ext.and_then(|e| e.preds.get(e.treatment, name).map(|p| (p, e.rep))) {
100        Some((p, rep)) => {
101            if p.nrows() != n || rep >= p.ncols() {
102                return Err(DmlError::DimensionMismatch(format!(
103                    "external predictions for {name} have shape {:?}; expected ({n}, n_rep)",
104                    p.dim()
105                )));
106            }
107            Ok(p.column(rep).to_owned())
108        }
109        None => fit(),
110    }
111}
112
113/// Out-of-fold nuisance predictions supplied by the user, as DoubleML's
114/// `external_predictions`: per treatment and learner name (e.g. `"ml_l"`,
115/// `"ml_m"`), one column per sample-splitting repetition. Learners with
116/// external predictions are not fitted, and have no out-of-sample loss.
117#[derive(Debug, Clone, Default)]
118pub struct ExternalPredictions {
119    preds: Vec<(usize, String, Array2<f64>)>,
120}
121
122impl ExternalPredictions {
123    pub fn new() -> Self {
124        Self::default()
125    }
126
127    /// Predictions of learner `name` for treatment `treatment` (0 for a single
128    /// treatment), shape `(n_obs, n_rep)`.
129    pub fn with(
130        mut self,
131        treatment: usize,
132        name: impl Into<String>,
133        predictions: Array2<f64>,
134    ) -> Self {
135        let name = name.into();
136        self.preds
137            .retain(|(t, n, _)| !(*t == treatment && *n == name));
138        self.preds.push((treatment, name, predictions));
139        self
140    }
141
142    /// One repetition: predictions as a vector.
143    pub fn with_vector(
144        self,
145        treatment: usize,
146        name: impl Into<String>,
147        predictions: Array1<f64>,
148    ) -> Self {
149        let n = predictions.len();
150        self.with(
151            treatment,
152            name,
153            predictions
154                .into_shape_with_order((n, 1))
155                .expect("1-D to (n, 1)"),
156        )
157    }
158
159    pub(crate) fn get(&self, treatment: usize, name: &str) -> Option<&Array2<f64>> {
160        self.preds
161            .iter()
162            .find(|(t, n, _)| *t == treatment && n == name)
163            .map(|(_, _, p)| p)
164    }
165
166    pub(crate) fn contains(&self, treatment: usize, name: &str) -> bool {
167        self.get(treatment, name).is_some()
168    }
169}
170
171/// Ingredients for best-linear-predictor inference on effect heterogeneity
172/// ([`DmlFit::cate`](crate::DmlFit::cate), [`DmlFit::gate`](crate::DmlFit::gate)):
173/// the BLP regresses `signal − θ·theta_factor` on the basis, with every row of
174/// the basis multiplied by `basis_scale` if given (partially linear models,
175/// where the signal is the outcome residual and the scale the treatment
176/// residual).
177#[derive(Debug, Clone)]
178pub struct CateSignal {
179    pub signal: Array1<f64>,
180    pub theta_factor: Option<Array1<f64>>,
181    pub basis_scale: Option<Array1<f64>>,
182    /// The signal is a (doubly robust) signal of each unit's treatment effect,
183    /// so a policy tree may be fitted on it ([`DmlFit::policy_tree`](crate::DmlFit::policy_tree)).
184    pub is_effect: bool,
185}
186
187/// Score elements plus the cross-fitted nuisance predictions behind them.
188///
189/// Build it with [`ScoreEstimate::new`] (the struct is `#[non_exhaustive]`, so
190/// fields can be added without breaking custom models).
191#[derive(Debug, Clone)]
192#[non_exhaustive]
193pub struct ScoreEstimate {
194    pub score: LinearScore,
195    /// `(name, out-of-fold predictions)` — e.g. `("ml_l", ℓ̂)`.
196    pub predictions: Vec<(Name, Array1<f64>)>,
197    /// Ingredients for [sensitivity analysis](crate::sensitivity), if the
198    /// model supports it.
199    pub sensitivity: Option<SensitivityParts>,
200    /// What each nuisance learner predicts, for out-of-sample diagnostics.
201    pub targets: Vec<NuisanceTarget>,
202    /// Signal for CATE / GATE inference, if the model supports it.
203    pub cate: Option<CateSignal>,
204}
205
206/// Name of a nuisance learner: a `&'static str` literal or an owned `String`
207/// (`Cow` holds either without copying literals).
208pub type Name = std::borrow::Cow<'static, str>;
209
210/// Loss used to evaluate a nuisance learner out of sample.
211#[derive(Debug, Clone, Copy, PartialEq, Eq)]
212pub enum LossKind {
213    /// Root mean squared error (regression learners).
214    Rmse,
215    /// Log loss with predictions clipped to `[1e-15, 1 − 1e-15]` (classifiers).
216    LogLoss,
217}
218
219/// The target a nuisance learner predicts; `NaN` marks rows where the
220/// nuisance does not apply (e.g. `g₀` on treated rows).
221#[derive(Debug, Clone)]
222pub struct NuisanceTarget {
223    pub name: Name,
224    pub target: Array1<f64>,
225    pub kind: LossKind,
226}
227
228impl NuisanceTarget {
229    pub fn regression(name: impl Into<Name>, target: Array1<f64>) -> Self {
230        Self {
231            name: name.into(),
232            target,
233            kind: LossKind::Rmse,
234        }
235    }
236
237    pub fn classification(name: impl Into<Name>, target: Array1<f64>) -> Self {
238        Self {
239            name: name.into(),
240            target,
241            kind: LossKind::LogLoss,
242        }
243    }
244
245    /// Keep the target only where `keep` is true (`NaN` elsewhere).
246    pub fn only_where(mut self, keep: impl Fn(usize) -> bool) -> Self {
247        for (i, v) in self.target.iter_mut().enumerate() {
248            if !keep(i) {
249                *v = f64::NAN;
250            }
251        }
252        self
253    }
254
255    /// Out-of-sample loss of `prediction` over the rows with a target.
256    pub fn loss(&self, prediction: &Array1<f64>) -> f64 {
257        let mut total = 0.0;
258        let mut count = 0usize;
259        for (&t, &p) in self.target.iter().zip(prediction) {
260            if t.is_nan() {
261                continue;
262            }
263            total += match self.kind {
264                LossKind::Rmse => (t - p).powi(2),
265                LossKind::LogLoss => {
266                    let p = p.clamp(1e-15, 1.0 - 1e-15);
267                    -(t * p.ln() + (1.0 - t) * (1.0 - p).ln())
268                }
269            };
270            count += 1;
271        }
272        if count == 0 {
273            return f64::NAN;
274        }
275        let mean = total / count as f64;
276        match self.kind {
277            LossKind::Rmse => mean.sqrt(),
278            LossKind::LogLoss => mean,
279        }
280    }
281}
282
283impl ScoreEstimate {
284    /// `predictions` pairs a name (a literal such as `"ml_l"`, or a
285    /// `String`) with the out-of-fold predictions.
286    pub fn new<N: Into<Name>>(score: LinearScore, predictions: Vec<(N, Array1<f64>)>) -> Self {
287        Self {
288            score,
289            predictions: predictions
290                .into_iter()
291                .map(|(n, p)| (n.into(), p))
292                .collect(),
293            sensitivity: None,
294            targets: Vec::new(),
295            cate: None,
296        }
297    }
298
299    /// Attach the signal for CATE / GATE inference.
300    pub fn with_cate_signal(mut self, cate: CateSignal) -> Self {
301        self.cate = Some(cate);
302        self
303    }
304
305    /// Attach the nuisance targets used for learner diagnostics.
306    pub fn with_targets(mut self, targets: Vec<NuisanceTarget>) -> Self {
307        self.targets = targets;
308        self
309    }
310
311    /// Attach the model's sensitivity ingredients.
312    pub fn with_sensitivity(mut self, parts: SensitivityParts) -> Self {
313        self.sensitivity = Some(parts);
314        self
315    }
316}
317
318/// A DML model whose orthogonal score is linear in θ.
319///
320/// Implement this to plug a new score into the shared machinery (repeated
321/// cross-fitting, aggregation, inference, bootstrap) via [`fit_model`](crate::fit_model).
322pub trait LinearScoreModel {
323    /// Label for summaries, e.g. `"PLR (partialling out)"`.
324    fn name(&self) -> String;
325
326    /// Reject data the model cannot handle (called once before fitting).
327    fn validate(&self, _data: &DmlData) -> Result<()> {
328        Ok(())
329    }
330
331    /// Cross-fit the nuisances on `folds` and return the score elements.
332    fn estimate_score(&self, input: &ScoreInput<'_>, folds: &[Fold]) -> Result<ScoreEstimate>;
333
334    /// Strata for the default sample splits (e.g. treatment levels), so each
335    /// fold holds about the same share of every stratum. `None` (the default)
336    /// draws plain K-fold splits.
337    fn strata(&self, _data: &DmlData) -> Option<Vec<usize>> {
338        None
339    }
340}
341
342/// Builder methods shared by every model: sample-splitting settings and `fit`.
343///
344/// (`macro_rules!` generates code at compile time; here it pastes the same
345/// methods into each model's `impl` block instead of repeating them.)
346macro_rules! dml_model_api {
347    () => {
348        /// Number of cross-fitting folds (default 5).
349        pub fn n_folds(mut self, n_folds: usize) -> Self {
350            self.splitting.n_folds = n_folds;
351            self
352        }
353
354        /// Number of independent sample splits (default 1). Estimates are
355        /// aggregated across repetitions by the median, as in DoubleML.
356        pub fn n_rep(mut self, n_rep: usize) -> Self {
357            self.splitting.n_rep = n_rep;
358            self
359        }
360
361        /// Seed for the sample splits (default 0).
362        pub fn seed(mut self, seed: u64) -> Self {
363            self.splitting.seed = seed;
364            self
365        }
366
367        /// Draw seeded sample splits and fit. With clustered data
368        /// ([`DmlData::with_clusters`](crate::DmlData::with_clusters)) whole clusters are assigned to folds
369        /// and standard errors are cluster-robust.
370        pub fn fit(&self, data: &$crate::DmlData) -> $crate::Result<$crate::DmlFit> {
371            match data.clusters() {
372                Some(clusters) => {
373                    let splits = self.splitting.draw_clustered(clusters)?;
374                    $crate::fit_model_clustered(self, data, &splits)
375                }
376                None => {
377                    let splits = match $crate::LinearScoreModel::strata(self, data) {
378                        Some(strata) => self.splitting.draw_stratified(&strata)?,
379                        None => self.splitting.draw(data.n_obs())?,
380                    };
381                    $crate::fit_model(self, data, &splits)
382                }
383            }
384        }
385
386        /// Fit clustered data with explicit cluster folds, one
387        /// `Vec<ClusterFold>` per repetition.
388        pub fn fit_with_cluster_folds(
389            &self,
390            data: &$crate::DmlData,
391            splits: &[Vec<$crate::ClusterFold>],
392        ) -> $crate::Result<$crate::DmlFit> {
393            $crate::fit_model_clustered(self, data, splits)
394        }
395
396        /// Sensitivity benchmark for the covariate columns `benchmarking_set`:
397        /// refits this model without them (with the same sample splitting, so
398        /// `long` must come from [`fit`](Self::fit) with the current settings)
399        /// and compares it with `long`.
400        pub fn sensitivity_benchmark(
401            &self,
402            data: &$crate::DmlData,
403            long: &$crate::DmlFit,
404            benchmarking_set: &[usize],
405        ) -> $crate::Result<$crate::SensitivityBenchmark> {
406            let short = self.fit(&data.without_covariates(benchmarking_set)?)?;
407            long.sensitivity_benchmark(&short)
408        }
409
410        /// Fit with caller-supplied splits, one `Vec<Fold>` per repetition
411        /// (e.g. to reproduce another implementation exactly).
412        pub fn fit_with_folds(
413            &self,
414            data: &$crate::DmlData,
415            splits: &[Vec<$crate::Fold>],
416        ) -> $crate::Result<$crate::DmlFit> {
417            $crate::fit_model(self, data, splits)
418        }
419    };
420}
421pub(crate) use dml_model_api;
dml_core/model/mod.rs

dml_core/model/
mod.rs
