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 external_get(ext, n, name)? {
100        Some(p) => Ok(p),
101        None => fit(),
102    }
103}
104
105/// External predictions of learner `name` (`n` rows), if supplied.
106pub(crate) fn external_get(
107    ext: Option<ExternalRef<'_>>,
108    n: usize,
109    name: &str,
110) -> Result<Option<Array1<f64>>> {
111    match ext.and_then(|e| e.preds.get(e.treatment, name).map(|p| (p, e.rep))) {
112        Some((p, rep)) => {
113            if p.nrows() != n || rep >= p.ncols() {
114                return Err(DmlError::DimensionMismatch(format!(
115                    "external predictions for {name} have shape {:?}; expected ({n}, n_rep)",
116                    p.dim()
117                )));
118            }
119            Ok(Some(p.column(rep).to_owned()))
120        }
121        None => Ok(None),
122    }
123}
124
125/// Out-of-fold nuisance predictions supplied by the user, as DoubleML's
126/// `external_predictions`: per treatment and learner name (e.g. `"ml_l"`,
127/// `"ml_m"`), one column per sample-splitting repetition. Learners with
128/// external predictions are not fitted, and have no out-of-sample loss.
129#[derive(Debug, Clone, Default)]
130pub struct ExternalPredictions {
131    preds: Vec<(usize, String, Array2<f64>)>,
132}
133
134impl ExternalPredictions {
135    pub fn new() -> Self {
136        Self::default()
137    }
138
139    /// Predictions of learner `name` for treatment `treatment` (0 for a single
140    /// treatment), shape `(n_obs, n_rep)`.
141    pub fn with(
142        mut self,
143        treatment: usize,
144        name: impl Into<String>,
145        predictions: Array2<f64>,
146    ) -> Self {
147        let name = name.into();
148        self.preds
149            .retain(|(t, n, _)| !(*t == treatment && *n == name));
150        self.preds.push((treatment, name, predictions));
151        self
152    }
153
154    /// One repetition: predictions as a vector.
155    pub fn with_vector(
156        self,
157        treatment: usize,
158        name: impl Into<String>,
159        predictions: Array1<f64>,
160    ) -> Self {
161        let n = predictions.len();
162        self.with(
163            treatment,
164            name,
165            predictions
166                .into_shape_with_order((n, 1))
167                .expect("1-D to (n, 1)"),
168        )
169    }
170
171    pub(crate) fn get(&self, treatment: usize, name: &str) -> Option<&Array2<f64>> {
172        self.preds
173            .iter()
174            .find(|(t, n, _)| *t == treatment && n == name)
175            .map(|(_, _, p)| p)
176    }
177
178    pub(crate) fn contains(&self, treatment: usize, name: &str) -> bool {
179        self.get(treatment, name).is_some()
180    }
181}
182
183/// Ingredients for best-linear-predictor inference on effect heterogeneity
184/// ([`DmlFit::cate`](crate::DmlFit::cate), [`DmlFit::gate`](crate::DmlFit::gate)):
185/// the BLP regresses `signal − θ·theta_factor` on the basis, with every row of
186/// the basis multiplied by `basis_scale` if given (partially linear models,
187/// where the signal is the outcome residual and the scale the treatment
188/// residual).
189#[derive(Debug, Clone)]
190pub struct CateSignal {
191    pub signal: Array1<f64>,
192    pub theta_factor: Option<Array1<f64>>,
193    pub basis_scale: Option<Array1<f64>>,
194    /// The signal is a (doubly robust) signal of each unit's treatment effect,
195    /// so a policy tree may be fitted on it ([`DmlFit::policy_tree`](crate::DmlFit::policy_tree)).
196    pub is_effect: bool,
197}
198
199/// Score elements plus the cross-fitted nuisance predictions behind them.
200///
201/// Build it with [`ScoreEstimate::new`] (the struct is `#[non_exhaustive]`, so
202/// fields can be added without breaking custom models).
203#[derive(Debug, Clone)]
204#[non_exhaustive]
205pub struct ScoreEstimate {
206    pub score: LinearScore,
207    /// `(name, out-of-fold predictions)` — e.g. `("ml_l", ℓ̂)`.
208    pub predictions: Vec<(Name, Array1<f64>)>,
209    /// Ingredients for [sensitivity analysis](crate::sensitivity), if the
210    /// model supports it.
211    pub sensitivity: Option<SensitivityParts>,
212    /// What each nuisance learner predicts, for out-of-sample diagnostics.
213    pub targets: Vec<NuisanceTarget>,
214    /// Signal for CATE / GATE inference, if the model supports it.
215    pub cate: Option<CateSignal>,
216}
217
218/// Name of a nuisance learner: a `&'static str` literal or an owned `String`
219/// (`Cow` holds either without copying literals).
220pub type Name = std::borrow::Cow<'static, str>;
221
222/// Loss used to evaluate a nuisance learner out of sample.
223#[derive(Debug, Clone, Copy, PartialEq, Eq)]
224pub enum LossKind {
225    /// Root mean squared error (regression learners).
226    Rmse,
227    /// Log loss with predictions clipped to `[1e-15, 1 − 1e-15]` (classifiers).
228    LogLoss,
229}
230
231/// The target a nuisance learner predicts; `NaN` marks rows where the
232/// nuisance does not apply (e.g. `g₀` on treated rows).
233#[derive(Debug, Clone)]
234pub struct NuisanceTarget {
235    pub name: Name,
236    pub target: Array1<f64>,
237    pub kind: LossKind,
238}
239
240impl NuisanceTarget {
241    pub fn regression(name: impl Into<Name>, target: Array1<f64>) -> Self {
242        Self {
243            name: name.into(),
244            target,
245            kind: LossKind::Rmse,
246        }
247    }
248
249    pub fn classification(name: impl Into<Name>, target: Array1<f64>) -> Self {
250        Self {
251            name: name.into(),
252            target,
253            kind: LossKind::LogLoss,
254        }
255    }
256
257    /// Keep the target only where `keep` is true (`NaN` elsewhere).
258    pub fn only_where(mut self, keep: impl Fn(usize) -> bool) -> Self {
259        for (i, v) in self.target.iter_mut().enumerate() {
260            if !keep(i) {
261                *v = f64::NAN;
262            }
263        }
264        self
265    }
266
267    /// Out-of-sample loss of `prediction` over the rows with a target.
268    pub fn loss(&self, prediction: &Array1<f64>) -> f64 {
269        let mut total = 0.0;
270        let mut count = 0usize;
271        for (&t, &p) in self.target.iter().zip(prediction) {
272            if t.is_nan() {
273                continue;
274            }
275            total += match self.kind {
276                LossKind::Rmse => (t - p).powi(2),
277                LossKind::LogLoss => {
278                    let p = p.clamp(1e-15, 1.0 - 1e-15);
279                    -(t * p.ln() + (1.0 - t) * (1.0 - p).ln())
280                }
281            };
282            count += 1;
283        }
284        if count == 0 {
285            return f64::NAN;
286        }
287        let mean = total / count as f64;
288        match self.kind {
289            LossKind::Rmse => mean.sqrt(),
290            LossKind::LogLoss => mean,
291        }
292    }
293}
294
295impl ScoreEstimate {
296    /// `predictions` pairs a name (a literal such as `"ml_l"`, or a
297    /// `String`) with the out-of-fold predictions.
298    pub fn new<N: Into<Name>>(score: LinearScore, predictions: Vec<(N, Array1<f64>)>) -> Self {
299        Self {
300            score,
301            predictions: predictions
302                .into_iter()
303                .map(|(n, p)| (n.into(), p))
304                .collect(),
305            sensitivity: None,
306            targets: Vec::new(),
307            cate: None,
308        }
309    }
310
311    /// Attach the signal for CATE / GATE inference.
312    pub fn with_cate_signal(mut self, cate: CateSignal) -> Self {
313        self.cate = Some(cate);
314        self
315    }
316
317    /// Attach the nuisance targets used for learner diagnostics.
318    pub fn with_targets(mut self, targets: Vec<NuisanceTarget>) -> Self {
319        self.targets = targets;
320        self
321    }
322
323    /// Attach the model's sensitivity ingredients.
324    pub fn with_sensitivity(mut self, parts: SensitivityParts) -> Self {
325        self.sensitivity = Some(parts);
326        self
327    }
328}
329
330/// A DML model whose orthogonal score is linear in θ.
331///
332/// Implement this to plug a new score into the shared machinery (repeated
333/// cross-fitting, aggregation, inference, bootstrap) via [`fit_model`](crate::fit_model).
334pub trait LinearScoreModel {
335    /// Label for summaries, e.g. `"PLR (partialling out)"`.
336    fn name(&self) -> String;
337
338    /// Reject data the model cannot handle (called once before fitting).
339    fn validate(&self, _data: &DmlData) -> Result<()> {
340        Ok(())
341    }
342
343    /// Cross-fit the nuisances on `folds` and return the score elements.
344    fn estimate_score(&self, input: &ScoreInput<'_>, folds: &[Fold]) -> Result<ScoreEstimate>;
345
346    /// Strata for the default sample splits (e.g. treatment levels), so each
347    /// fold holds about the same share of every stratum. `None` (the default)
348    /// draws plain K-fold splits.
349    fn strata(&self, _data: &DmlData) -> Option<Vec<usize>> {
350        None
351    }
352}
353
354/// Builder methods shared by every model: sample-splitting settings and `fit`.
355///
356/// (`macro_rules!` generates code at compile time; here it pastes the same
357/// methods into each model's `impl` block instead of repeating them.)
358macro_rules! dml_model_api {
359    () => {
360        /// Number of cross-fitting folds (default 5).
361        pub fn n_folds(mut self, n_folds: usize) -> Self {
362            self.splitting.n_folds = n_folds;
363            self
364        }
365
366        /// Number of independent sample splits (default 1). Estimates are
367        /// aggregated across repetitions by the median, as in DoubleML.
368        pub fn n_rep(mut self, n_rep: usize) -> Self {
369            self.splitting.n_rep = n_rep;
370            self
371        }
372
373        /// Seed for the sample splits (default 0).
374        pub fn seed(mut self, seed: u64) -> Self {
375            self.splitting.seed = seed;
376            self
377        }
378
379        /// Draw seeded sample splits and fit. With clustered data
380        /// ([`DmlData::with_clusters`](crate::DmlData::with_clusters)) whole clusters are assigned to folds
381        /// and standard errors are cluster-robust.
382        pub fn fit(&self, data: &$crate::DmlData) -> $crate::Result<$crate::DmlFit> {
383            match data.clusters() {
384                Some(clusters) => {
385                    let splits = self.splitting.draw_clustered(clusters)?;
386                    $crate::fit_model_clustered(self, data, &splits)
387                }
388                None => {
389                    let splits = match $crate::LinearScoreModel::strata(self, data) {
390                        Some(strata) => self.splitting.draw_stratified(&strata)?,
391                        None => self.splitting.draw(data.n_obs())?,
392                    };
393                    $crate::fit_model(self, data, &splits)
394                }
395            }
396        }
397
398        /// Fit clustered data with explicit cluster folds, one
399        /// `Vec<ClusterFold>` per repetition.
400        pub fn fit_with_cluster_folds(
401            &self,
402            data: &$crate::DmlData,
403            splits: &[Vec<$crate::ClusterFold>],
404        ) -> $crate::Result<$crate::DmlFit> {
405            $crate::fit_model_clustered(self, data, splits)
406        }
407
408        /// Sensitivity benchmark for the covariate columns `benchmarking_set`:
409        /// refits this model without them (with the same sample splitting, so
410        /// `long` must come from [`fit`](Self::fit) with the current settings)
411        /// and compares it with `long`.
412        pub fn sensitivity_benchmark(
413            &self,
414            data: &$crate::DmlData,
415            long: &$crate::DmlFit,
416            benchmarking_set: &[usize],
417        ) -> $crate::Result<$crate::SensitivityBenchmark> {
418            let short = self.fit(&data.without_covariates(benchmarking_set)?)?;
419            long.sensitivity_benchmark(&short)
420        }
421
422        /// Fit with caller-supplied splits, one `Vec<Fold>` per repetition
423        /// (e.g. to reproduce another implementation exactly).
424        pub fn fit_with_folds(
425            &self,
426            data: &$crate::DmlData,
427            splits: &[Vec<$crate::Fold>],
428        ) -> $crate::Result<$crate::DmlFit> {
429            $crate::fit_model(self, data, splits)
430        }
431    };
432}
433pub(crate) use dml_model_api;
434
435/// Learners for individual nuisance functions that replace a model's default
436/// learner for them, e.g. different settings for `ml_g0` and `ml_g1` in
437/// [`Irm`] (DoubleML's per-nuisance parameters, `set_ml_nuisance_params`).
438/// Set through the models' `nuisance_regressor` / `nuisance_classifier`.
439#[derive(Default)]
440pub struct NuisanceLearners {
441    regressors: Vec<(String, Box<dyn crate::learner::Regressor>)>,
442    classifiers: Vec<(String, Box<dyn crate::learner::Classifier>)>,
443}
444
445impl NuisanceLearners {
446    pub(crate) fn set_regressor(
447        &mut self,
448        name: &str,
449        learner: Box<dyn crate::learner::Regressor>,
450    ) {
451        self.regressors.retain(|(n, _)| n != name);
452        self.regressors.push((name.to_string(), learner));
453    }
454
455    pub(crate) fn set_classifier(
456        &mut self,
457        name: &str,
458        learner: Box<dyn crate::learner::Classifier>,
459    ) {
460        self.classifiers.retain(|(n, _)| n != name);
461        self.classifiers.push((name.to_string(), learner));
462    }
463
464    /// The learner for nuisance `name`, or `default`.
465    pub(crate) fn regressor<'a>(
466        &'a self,
467        name: &str,
468        default: &'a dyn crate::learner::Regressor,
469    ) -> &'a dyn crate::learner::Regressor {
470        self.regressors
471            .iter()
472            .find(|(n, _)| n == name)
473            .map_or(default, |(_, l)| l.as_ref())
474    }
475
476    /// The learner for nuisance `name`, or `default`.
477    pub(crate) fn classifier<'a>(
478        &'a self,
479        name: &str,
480        default: &'a dyn crate::learner::Classifier,
481    ) -> &'a dyn crate::learner::Classifier {
482        self.classifiers
483            .iter()
484            .find(|(n, _)| n == name)
485            .map_or(default, |(_, l)| l.as_ref())
486    }
487
488    /// Error unless every overridden name is one of the model's regression
489    /// (`regressors`) or classification (`classifiers`) nuisances.
490    pub(crate) fn check(
491        &self,
492        model: &str,
493        regressors: &[&str],
494        classifiers: &[&str],
495    ) -> Result<()> {
496        for (name, valid) in self
497            .regressors
498            .iter()
499            .map(|(n, _)| (n, regressors))
500            .chain(self.classifiers.iter().map(|(n, _)| (n, classifiers)))
501        {
502            if !valid.contains(&name.as_str()) {
503                return Err(DmlError::InvalidArgument(format!(
504                    "{model} has no nuisance {name:?} for this learner type (valid: {})",
505                    valid.join(", ")
506                )));
507            }
508        }
509        Ok(())
510    }
511}
512
513/// Builder methods for per-nuisance learners (a field `nuisance:
514/// NuisanceLearners`).
515macro_rules! nuisance_api {
516    () => {
517        /// Use `learner` for the regression nuisance `name` (DoubleML's name,
518        /// e.g. `"ml_g1"`) instead of the model's default learner for it.
519        pub fn nuisance_regressor(
520            self,
521            name: &str,
522            learner: impl $crate::learner::Regressor + 'static,
523        ) -> Self {
524            self.nuisance_regressor_boxed(name, Box::new(learner))
525        }
526
527        /// As [`nuisance_regressor`](Self::nuisance_regressor), boxed.
528        pub fn nuisance_regressor_boxed(
529            mut self,
530            name: &str,
531            learner: Box<dyn $crate::learner::Regressor>,
532        ) -> Self {
533            self.nuisance.set_regressor(name, learner);
534            self
535        }
536
537        /// Use `learner` for the classification nuisance `name` instead of the
538        /// model's default learner for it.
539        pub fn nuisance_classifier(
540            self,
541            name: &str,
542            learner: impl $crate::learner::Classifier + 'static,
543        ) -> Self {
544            self.nuisance_classifier_boxed(name, Box::new(learner))
545        }
546
547        /// As [`nuisance_classifier`](Self::nuisance_classifier), boxed.
548        pub fn nuisance_classifier_boxed(
549            mut self,
550            name: &str,
551            learner: Box<dyn $crate::learner::Classifier>,
552        ) -> Self {
553            self.nuisance.set_classifier(name, learner);
554            self
555        }
556    };
557}
558pub(crate) use nuisance_api;
dml_core/model/mod.rs

dml_core/model/
mod.rs
