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 rd;
78pub mod score;
79pub mod sensitivity;
80pub mod staggered;
81pub mod stats;
82
83pub use blp::Blp;
84pub use bootstrap::{Bootstrap, BootstrapMethod};
85pub use cate::{CateFit, DrLearner, RLearner, SLearner, TLearner, XLearner};
86pub use cluster::ClusterFold;
87pub use crossfit::{
88 cross_fit_predict, cross_fit_predict_proba, cross_fit_predict_proba_subset,
89 cross_fit_predict_subset, Fold, KFold, SampleSplitting,
90};
91pub use data::DmlData;
92pub use error::{DmlError, Result};
93pub use fit::{fit_model, fit_model_clustered, DmlFit, PAdjust};
94pub use learner::{
95 Classifier, CvScoring, GradientBoostingClassifier, GradientBoostingRegressor, Lasso, LassoCV,
96 LogisticRegression, LogisticRegressionCV, MaxFeatures, Ols, RandomForestClassifier,
97 RandomForestRegressor, Regressor,
98};
99pub use model::{
100 Apo, Apos, AposFit, CateSignal, Cvar, Did, DidCs, DidScore, FirstStage, Iivm, Irm, IrmScore,
101 LinearScoreModel, LossKind, Lplr, LplrScore, Lpq, NuisanceTarget, PanelApproach, Pliv, Plpr,
102 Plr, PlrScore, Pq, Qte, QteFit, RdFlex, RdFlexFit, RdKernel, ScoreEstimate, ScoreInput, Ssm,
103 SsmScore, StaticPanel, TreatmentLearner,
104};
105pub use policy::PolicyTree;
106pub use score::LinearScore;
107pub use sensitivity::{
108 Nu2Parts, SensitivityBenchmark, SensitivityParams, SensitivityParts, SensitivityResult,
109};
110pub use staggered::{
111 Aggregation, ControlGroup, DidAggregation, GroupTime, GtCombinations, PanelData,
112 RepeatedCrossSections, StaggeredDid, StaggeredDidFit, StaggeredInput,
113};