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