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

dml_core/model/
mod.rs
