Crate dml_core 

dml-core

Double / debiased machine learning (Chernozhukov et al., 2018) in Rust.

Models

• Plr: partially linear regression (partialling-out or IV-type score)
• Irm: interactive regression model for a binary treatment (ATE / ATTE)
• Pliv: partially linear IV regression
• Iivm: interactive IV model, local average treatment effect (LATE) with a binary instrument
• Did / DidCs: doubly robust difference-in-differences for panel data and repeated cross-sections (Sant’Anna & Zhao, 2020)
• StaggeredDid: DID with staggered adoption, group-time effects ATT(g, t) and their aggregation by cohort, calendar time or event time (Callaway & Sant’Anna, 2021)

All of them have scores that are linear in θ and share one engine (fit_model): repeated cross-fitting, DML2 estimation, standard errors, confidence intervals, a multiplier Bootstrap for joint inference and multiple-testing adjustment. New scores plug in via LinearScoreModel.

Heterogeneous effects

• SLearner, TLearner, XLearner, RLearner, DrLearner: CATE meta-learners
• Blp: best linear predictor of the CATE and group average effects (GATE) with robust inference, from DrLearner fits or DmlFit::cate / DmlFit::gate

Building blocks

• DmlData: outcome y, one or more treatments d, covariates x, optional instrument z
• Regressor / Classifier: plug-in nuisance learners (Ols, Lasso, LogisticRegression, RandomForestRegressor, RandomForestClassifier built in)
• KFold / SampleSplitting: reproducible sample splitting

use dml_core::{DmlData, Ols, Plr};
use ndarray::{Array1, Array2};

let n = 200;
let x = Array2::from_shape_fn((n, 2), |(i, j)| ((i * 7 + j * 3) % 11) as f64);
let d: Array1<f64> = (0..n).map(|i| x[[i, 0]] + ((i * 13) % 5) as f64).collect();
let y: Array1<f64> = (0..n).map(|i| 0.5 * d[i] - x[[i, 1]] + ((i * 17) % 3) as f64).collect();

let fit = Plr::new(Ols::new(), Ols::new()).n_folds(5).seed(42).fit(&DmlData::new(y, d, x)?)?;
assert!((fit.coef()[0] - 0.5).abs() < 0.1);
let ci = fit.confint(0.95)?; // shape (n_treat, 2)

Re-exports

pub use blp::Blp;

pub use bootstrap::Bootstrap;

pub use bootstrap::BootstrapMethod;

pub use cate::CateFit;

pub use cate::DrLearner;

pub use cate::RLearner;

pub use cate::SLearner;

pub use cate::TLearner;

pub use cate::XLearner;

pub use crossfit::cross_fit_predict;

pub use crossfit::cross_fit_predict_proba;

pub use crossfit::cross_fit_predict_proba_subset;

pub use crossfit::cross_fit_predict_subset;

pub use crossfit::Fold;

pub use crossfit::KFold;

pub use crossfit::SampleSplitting;

pub use data::DmlData;

pub use error::DmlError;

pub use error::Result;

pub use fit::fit_model;

pub use fit::DmlFit;

pub use fit::PAdjust;

pub use learner::Classifier;

pub use learner::Lasso;

pub use learner::LogisticRegression;

pub use learner::MaxFeatures;

pub use learner::Ols;

pub use learner::RandomForestClassifier;

pub use learner::RandomForestRegressor;

pub use learner::Regressor;

pub use model::Did;

pub use model::DidCs;

pub use model::DidScore;

pub use model::Iivm;

pub use model::Irm;

pub use model::IrmScore;

pub use model::LinearScoreModel;

pub use model::Pliv;

pub use model::Plr;

pub use model::PlrScore;

pub use model::ScoreEstimate;

pub use model::ScoreInput;

pub use score::LinearScore;

pub use staggered::Aggregation;

pub use staggered::ControlGroup;

pub use staggered::DidAggregation;

pub use staggered::GroupTime;

pub use staggered::GtCombinations;

pub use staggered::PanelData;

pub use staggered::RepeatedCrossSections;

pub use staggered::StaggeredDid;

pub use staggered::StaggeredDidFit;

pub use staggered::StaggeredInput;

Modules

blp

Best linear predictor (BLP) of a CATE and group average treatment effects (GATE), following Semenova & Chernozhukov (2021) and DoubleML.

bootstrap

Multiplier bootstrap for joint confidence intervals and Romano–Wolf multiple-testing adjustment.

cate

Meta-learners for the conditional average treatment effect τ(x) = E[Y(1) − Y(0) | X = x] of a binary treatment.

crossfit

data

error

fit

The shared DML driver and its result type, DmlFit.

learner

Nuisance learners: the Regressor / Classifier traits and the built-in implementations.

model

DML models with linear scores: Plr, Irm, Pliv, and the LinearScoreModel trait for adding your own.

score

Linear Neyman-orthogonal scores and the DML2 estimator.

staggered

Difference-in-differences with staggered adoption: units start treatment in different periods (cohorts), following Callaway & Sant’Anna (2021) with the doubly robust DML scores of DoubleML’s DoubleMLDIDMulti.

stats

Standard normal distribution helpers for Wald inference (no external deps).