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); models whose nuisance models share a learner argument (IRM, APO, IIVM, DID,
46//! SSM) also take a learner per nuisance model (`nuisance_regressor("ml_g1", ...)`, see
47//! [`NuisanceLearners`])
48//! - [`KFold`] / [`SampleSplitting`]: reproducible sample splitting
49//!
50//! ```
51//! use dml_core::{DmlData, Ols, Plr};
52//! use ndarray::{Array1, Array2};
53//!
54//! let n = 200;
55//! let x = Array2::from_shape_fn((n, 2), |(i, j)| ((i * 7 + j * 3) % 11) as f64);
56//! let d: Array1<f64> = (0..n).map(|i| x[[i, 0]] + ((i * 13) % 5) as f64).collect();
57//! let y: Array1<f64> = (0..n).map(|i| 0.5 * d[i] - x[[i, 1]] + ((i * 17) % 3) as f64).collect();
58//!
59//! let fit = Plr::new(Ols::new(), Ols::new()).n_folds(5).seed(42).fit(&DmlData::new(y, d, x)?)?;
60//! assert!((fit.coef()[0] - 0.5).abs() < 0.1);
61//! let ci = fit.confint(0.95)?; // shape (n_treat, 2)
62//! # let _ = ci;
63//! # Ok::<(), dml_core::DmlError>(())
64//! ```
65
66pub mod blp;
67pub mod bootstrap;
68pub mod cate;
69pub mod cluster;
70mod compat;
71pub mod crossfit;
72pub mod data;
73pub mod datasets;
74pub mod error;
75pub mod fit;
76pub mod learner;
77pub(crate) mod linalg;
78pub mod model;
79pub mod policy;
80pub mod propensity;
81pub mod rd;
82pub mod score;
83pub mod sensitivity;
84pub mod staggered;
85pub mod stats;
86
87pub use blp::Blp;
88pub use bootstrap::{Bootstrap, BootstrapMethod};
89pub use cate::{CateFit, DrLearner, RLearner, SLearner, TLearner, XLearner};
90pub use cluster::ClusterFold;
91pub use crossfit::{
92 cross_fit_predict, cross_fit_predict_proba, cross_fit_predict_proba_subset,
93 cross_fit_predict_subset, set_parallel_folds, Fold, KFold, SampleSplitting,
94};
95pub use data::DmlData;
96pub use error::{DmlError, Result};
97pub use fit::{fit_model, fit_model_clustered, DmlFit, PAdjust};
98pub use learner::{
99 Classifier, CvScoring, GradientBoostingClassifier, GradientBoostingRegressor,
100 GridSearchClassifier, GridSearchRegressor, Lasso, LassoCV, LogisticRegression,
101 LogisticRegressionCV, MaxFeatures, Ols, RandomForestClassifier, RandomForestRegressor,
102 Regressor,
103};
104pub use model::{
105 Apo, Apos, AposFit, CateSignal, Cvar, Did, DidCs, DidScore, ExternalPredictions, FirstStage,
106 Iivm, Irm, IrmScore, LinearScoreModel, LossKind, Lplr, LplrScore, Lpq, Name, NuisanceLearners,
107 NuisanceTarget, PanelApproach, Pliv, PlivScore, PlivVariant, Plpr, Plr, PlrScore, Pq, Qte,
108 QteFit, RdFlex, RdFlexFit, RdKernel, ScoreEstimate, ScoreInput, Ssm, SsmScore, StaticPanel,
109 TreatmentLearner,
110};
111pub use policy::{PolicyNode, PolicyTree};
112pub use propensity::{Calibration, IsotonicRegression};
113pub use score::LinearScore;
114pub use sensitivity::{
115 Nu2Parts, SensitivityBenchmark, SensitivityParams, SensitivityParts, SensitivityResult,
116};
117pub use staggered::{
118 Aggregation, ControlGroup, DidAggregation, GroupTime, GtCombinations, PanelData,
119 RepeatedCrossSections, StaggeredDid, StaggeredDidFit, StaggeredInput,
120};