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