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;
68pub mod crossfit;
69pub mod data;
70pub mod error;
71pub mod fit;
72pub mod learner;
73pub(crate) mod linalg;
74pub mod model;
75pub mod policy;
76pub mod score;
77pub mod sensitivity;
78pub mod staggered;
79pub mod stats;
80
81pub use blp::Blp;
82pub use bootstrap::{Bootstrap, BootstrapMethod};
83pub use cate::{CateFit, DrLearner, RLearner, SLearner, TLearner, XLearner};
84pub use cluster::ClusterFold;
85pub use crossfit::{
86 cross_fit_predict, cross_fit_predict_proba, cross_fit_predict_proba_subset,
87 cross_fit_predict_subset, Fold, KFold, SampleSplitting,
88};
89pub use data::DmlData;
90pub use error::{DmlError, Result};
91pub use fit::{fit_model, fit_model_clustered, DmlFit, PAdjust};
92pub use learner::{
93 Classifier, CvScoring, GradientBoostingClassifier, GradientBoostingRegressor, Lasso, LassoCV,
94 LogisticRegression, LogisticRegressionCV, MaxFeatures, Ols, RandomForestClassifier,
95 RandomForestRegressor, Regressor,
96};
97pub use model::{
98 Apo, Apos, AposFit, CateSignal, Did, DidCs, DidScore, Iivm, Irm, IrmScore, LinearScoreModel,
99 LossKind, NuisanceTarget, Pliv, Plr, PlrScore, ScoreEstimate, ScoreInput,
100};
101pub use policy::PolicyTree;
102pub use score::LinearScore;
103pub use sensitivity::{
104 Nu2Parts, SensitivityBenchmark, SensitivityParams, SensitivityParts, SensitivityResult,
105};
106pub use staggered::{
107 Aggregation, ControlGroup, DidAggregation, GroupTime, GtCombinations, PanelData,
108 RepeatedCrossSections, StaggeredDid, StaggeredDidFit, StaggeredInput,
109};