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//!
11//! All three have scores that are linear in θ and share one engine
12//! ([`fit_model`]): repeated cross-fitting, DML2 estimation, standard errors,
13//! confidence intervals, a multiplier [`Bootstrap`] for joint inference and
14//! multiple-testing adjustment. New scores plug in via [`LinearScoreModel`].
15//!
16//! ## Heterogeneous effects
17//!
18//! - [`SLearner`], [`TLearner`], [`XLearner`], [`RLearner`], [`DrLearner`]: CATE meta-learners
19//! - [`Blp`]: best linear predictor of the CATE and group average effects (GATE) with
20//! robust inference, from [`DrLearner`] fits or [`DmlFit::cate`] / [`DmlFit::gate`]
21//!
22//! ## Building blocks
23//!
24//! - [`DmlData`]: outcome `y`, one or more treatments `d`, covariates `x`, optional instrument `z`
25//! - [`Regressor`] / [`Classifier`]: plug-in nuisance learners
26//! ([`Ols`], [`Lasso`], [`LogisticRegression`], [`RandomForestRegressor`],
27//! [`RandomForestClassifier`] built in)
28//! - [`KFold`] / [`SampleSplitting`]: reproducible sample splitting
29//!
30//! ```
31//! use dml_core::{DmlData, Ols, Plr};
32//! use ndarray::{Array1, Array2};
33//!
34//! let n = 200;
35//! let x = Array2::from_shape_fn((n, 2), |(i, j)| ((i * 7 + j * 3) % 11) as f64);
36//! let d: Array1<f64> = (0..n).map(|i| x[[i, 0]] + ((i * 13) % 5) as f64).collect();
37//! let y: Array1<f64> = (0..n).map(|i| 0.5 * d[i] - x[[i, 1]] + ((i * 17) % 3) as f64).collect();
38//!
39//! let fit = Plr::new(Ols::new(), Ols::new()).n_folds(5).seed(42).fit(&DmlData::new(y, d, x)?)?;
40//! assert!((fit.coef()[0] - 0.5).abs() < 0.1);
41//! let ci = fit.confint(0.95)?; // shape (n_treat, 2)
42//! # let _ = ci;
43//! # Ok::<(), dml_core::DmlError>(())
44//! ```
45
46pub mod blp;
47pub mod bootstrap;
48pub mod cate;
49pub mod crossfit;
50pub mod data;
51pub mod error;
52pub mod fit;
53pub mod learner;
54pub(crate) mod linalg;
55pub mod model;
56pub mod score;
57pub mod stats;
58
59pub use blp::Blp;
60pub use bootstrap::{Bootstrap, BootstrapMethod};
61pub use cate::{CateFit, DrLearner, RLearner, SLearner, TLearner, XLearner};
62pub use crossfit::{
63 cross_fit_predict, cross_fit_predict_proba, cross_fit_predict_subset, Fold, KFold,
64 SampleSplitting,
65};
66pub use data::DmlData;
67pub use error::{DmlError, Result};
68pub use fit::{fit_model, DmlFit, PAdjust};
69pub use learner::{
70 Classifier, Lasso, LogisticRegression, MaxFeatures, Ols, RandomForestClassifier,
71 RandomForestRegressor, Regressor,
72};
73pub use model::{Irm, IrmScore, LinearScoreModel, Pliv, Plr, PlrScore, ScoreEstimate, ScoreInput};
74pub use score::LinearScore;