dml_core/lib.rs
1//! # dml-core
2//!
3//! Double / debiased machine learning (Chernozhukov et al., 2018) in Rust.
4//!
5//! ## Models
6//!
7//! - [`Plr`]: partially linear regression (partialling-out or IV-type score)
8//! - [`Irm`]: interactive regression model for a binary treatment (ATE / ATTE)
9//! - [`Pliv`]: partially linear IV regression
10//! - [`Iivm`]: interactive IV model, local average treatment effect (LATE) with a binary instrument
11//! - [`Did`] / [`DidCs`]: doubly robust difference-in-differences for panel data and
12//! repeated cross-sections (Sant'Anna & Zhao, 2020)
13//! - [`AutoDml`]: automatic debiased ML of any mean of a linear functional of the regression
14//! (ATE, average derivatives, shift and policy effects, your own), with the Riesz
15//! representer learned from the functional (see [`riesz`])
16//! - [`StaggeredDid`]: DID with staggered adoption, group-time effects `ATT(g, t)` and their
17//! aggregation by cohort, calendar time or event time (Callaway & Sant'Anna, 2021)
18//!
19//! All of them have scores that are linear in θ and share one engine
20//! ([`fit_model`]): repeated cross-fitting, DML2 estimation, standard errors,
21//! confidence intervals, a multiplier [`Bootstrap`] for joint inference and
22//! multiple-testing adjustment. New scores plug in via [`LinearScoreModel`].
23//!
24//! ## Clustered data
25//!
26//! [`DmlData::with_clusters`] / [`DmlData::with_two_way_clusters`] make sample splitting
27//! cluster-aware and standard errors cluster-robust (Chiang et al., 2022).
28//!
29//! ## Sensitivity analysis
30//!
31//! [`DmlFit::sensitivity_analysis`] bounds the effect of omitted confounders (PLR, IRM, DID)
32//! and reports robustness values; `sensitivity_benchmark` calibrates the scenario against
33//! observed covariates.
34//!
35//! ## Heterogeneous effects
36//!
37//! - [`SLearner`], [`TLearner`], [`XLearner`], [`RLearner`], [`DrLearner`]: CATE meta-learners
38//! - [`Blp`]: best linear predictor of the CATE and group average effects (GATE) with
39//! robust inference, from [`DrLearner`] fits or [`DmlFit::cate`] / [`DmlFit::gate`]
40//!
41//! ## Building blocks
42//!
43//! - [`DmlData`]: outcome `y`, one or more treatments `d`, covariates `x`, optional instrument `z`
44//! - [`Regressor`] / [`Classifier`]: plug-in nuisance learners
45//! ([`Ols`], [`Lasso`], [`LassoCV`], [`LogisticRegression`], [`LogisticRegressionCV`],
46//! [`RandomForestRegressor`],
47//! [`RandomForestClassifier`], [`GradientBoostingRegressor`], [`GradientBoostingClassifier`]
48//! built in); models whose nuisance models share a learner argument (IRM, APO, IIVM, DID,
49//! SSM) also take a learner per nuisance model (`nuisance_regressor("ml_g1", ...)`, see
50//! [`NuisanceLearners`])
51//! - [`KFold`] / [`SampleSplitting`]: reproducible sample splitting
52//!
53//! ```
54//! use dml_core::{DmlData, Ols, Plr};
55//! use ndarray::{Array1, Array2};
56//!
57//! let n = 200;
58//! let x = Array2::from_shape_fn((n, 2), |(i, j)| ((i * 7 + j * 3) % 11) as f64);
59//! let d: Array1<f64> = (0..n).map(|i| x[[i, 0]] + ((i * 13) % 5) as f64).collect();
60//! let y: Array1<f64> = (0..n).map(|i| 0.5 * d[i] - x[[i, 1]] + ((i * 17) % 3) as f64).collect();
61//!
62//! let fit = Plr::new(Ols::new(), Ols::new()).n_folds(5).seed(42).fit(&DmlData::new(y, d, x)?)?;
63//! assert!((fit.coef()[0] - 0.5).abs() < 0.1);
64//! let ci = fit.confint(0.95)?; // shape (n_treat, 2)
65//! # let _ = ci;
66//! # Ok::<(), dml_core::DmlError>(())
67//! ```
68
69pub mod blp;
70pub mod bootstrap;
71pub mod cate;
72pub mod cluster;
73mod compat;
74pub mod crossfit;
75pub mod data;
76pub mod datasets;
77pub mod error;
78pub mod fit;
79pub mod learner;
80pub(crate) mod linalg;
81pub mod model;
82pub mod policy;
83pub mod propensity;
84pub mod rd;
85pub mod riesz;
86pub mod score;
87pub mod sensitivity;
88pub mod staggered;
89pub mod stats;
90
91pub use blp::Blp;
92pub use bootstrap::{Bootstrap, BootstrapMethod};
93pub use cate::{CateFit, DrLearner, RLearner, SLearner, TLearner, XLearner};
94pub use cluster::ClusterFold;
95pub use crossfit::{
96 cross_fit_predict, cross_fit_predict_proba, cross_fit_predict_proba_subset,
97 cross_fit_predict_subset, set_parallel_folds, Fold, KFold, SampleSplitting,
98};
99pub use data::DmlData;
100pub use error::{DmlError, Result};
101pub use fit::{fit_model, fit_model_clustered, DmlFit, PAdjust};
102pub use learner::{
103 Classifier, CvScoring, GradientBoostingClassifier, GradientBoostingRegressor,
104 GridSearchClassifier, GridSearchRegressor, Lasso, LassoCV, LogisticRegression,
105 LogisticRegressionCV, MaxFeatures, Ols, RandomForestClassifier, RandomForestRegressor,
106 Regressor,
107};
108pub use model::{
109 Apo, Apos, AposFit, AutoDml, CateSignal, Cvar, Did, DidCs, DidScore, ExternalPredictions,
110 FirstStage, Iivm, Irm, IrmScore, LinearScoreModel, LossKind, Lplr, LplrScore, Lpq, Name,
111 NuisanceLearners, NuisanceTarget, PanelApproach, Pliv, PlivScore, PlivVariant, Plpr, Plr,
112 PlrScore, Pq, Qte, QteFit, RdFlex, RdFlexFit, RdKernel, ScoreEstimate, ScoreInput, Ssm,
113 SsmScore, StaticPanel, TreatmentLearner,
114};
115pub use policy::{PolicyNode, PolicyTree};
116pub use propensity::{Calibration, IsotonicRegression};
117pub use riesz::{
118 Ate, AverageDerivative, CustomFunctional, Dictionary, LinearFunctional, PolicyEffect,
119 RieszLasso, RieszLearner, RieszPenalty, ShiftEffect, Term,
120};
121pub use score::LinearScore;
122pub use sensitivity::{
123 Nu2Parts, SensitivityBenchmark, SensitivityParams, SensitivityParts, SensitivityResult,
124};
125pub use staggered::{
126 Aggregation, ControlGroup, DidAggregation, GroupTime, GtCombinations, PanelData,
127 RepeatedCrossSections, StaggeredDid, StaggeredDidFit, StaggeredInput,
128};