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.

Clustered data

DmlData::with_clusters / DmlData::with_two_way_clusters make sample splitting cluster-aware and standard errors cluster-robust (Chiang et al., 2022).

Sensitivity analysis

DmlFit::sensitivity_analysis bounds the effect of omitted confounders (PLR, IRM, DID) and reports robustness values; sensitivity_benchmark calibrates the scenario against observed covariates.

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, LassoCV, LogisticRegression, LogisticRegressionCV, RandomForestRegressor, RandomForestClassifier, GradientBoostingRegressor, GradientBoostingClassifier built in); models whose nuisance models share a learner argument (IRM, APO, IIVM, DID, SSM) also take a learner per nuisance model (nuisance_regressor("ml_g1", ...), see NuisanceLearners)
• 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 cluster::ClusterFold;

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::set_parallel_folds;

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::fit_model_clustered;

pub use fit::DmlFit;

pub use fit::PAdjust;

pub use learner::Classifier;

pub use learner::CvScoring;

pub use learner::GradientBoostingClassifier;

pub use learner::GradientBoostingRegressor;

pub use learner::GridSearchClassifier;

pub use learner::GridSearchRegressor;

pub use learner::Lasso;

pub use learner::LassoCV;

pub use learner::LogisticRegression;

pub use learner::LogisticRegressionCV;

pub use learner::MaxFeatures;

pub use learner::Ols;

pub use learner::RandomForestClassifier;

pub use learner::RandomForestRegressor;

pub use learner::Regressor;

pub use model::Apo;

pub use model::Apos;

pub use model::AposFit;

pub use model::CateSignal;

pub use model::Cvar;

pub use model::Did;

pub use model::DidCs;

pub use model::DidScore;

pub use model::ExternalPredictions;

pub use model::FirstStage;

pub use model::Iivm;

pub use model::Irm;

pub use model::IrmScore;

pub use model::LinearScoreModel;

pub use model::LossKind;

pub use model::Lplr;

pub use model::LplrScore;

pub use model::Lpq;

pub use model::Name;

pub use model::NuisanceLearners;

pub use model::NuisanceTarget;

pub use model::PanelApproach;

pub use model::Pliv;

pub use model::PlivScore;

pub use model::PlivVariant;

pub use model::Plpr;

pub use model::Plr;

pub use model::PlrScore;

pub use model::Pq;

pub use model::Qte;

pub use model::QteFit;

pub use model::RdFlex;

pub use model::RdFlexFit;

pub use model::RdKernel;

pub use model::ScoreEstimate;

pub use model::ScoreInput;

pub use model::Ssm;

pub use model::SsmScore;

pub use model::StaticPanel;

pub use model::TreatmentLearner;

pub use policy::PolicyNode;

pub use policy::PolicyTree;

pub use propensity::Calibration;

pub use propensity::IsotonicRegression;

pub use score::LinearScore;

pub use sensitivity::Nu2Parts;

pub use sensitivity::SensitivityBenchmark;

pub use sensitivity::SensitivityParams;

pub use sensitivity::SensitivityParts;

pub use sensitivity::SensitivityResult;

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.

cluster

Cluster-robust DML (Chiang, Kato, Ma & Sasaki, 2022), as DoubleML: one-way or two-way clustered observations.

crossfit

data

datasets

Simulated data sets: the data generating processes of DoubleML’s make_* functions, for examples, tests and simulation studies.

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.

policy

Policy trees: a shallow decision tree that recommends treatment where the estimated effect is positive, as DoubleML’s DoubleMLPolicyTree (Athey and Wager, 2021).

propensity

Propensity score processing, as DoubleML’s PSProcessor: optional isotonic calibration of the cross-fitted propensity scores against the treatment (or instrument), followed by clipping to [t, 1 − t].

rd

Local polynomial regression discontinuity (RD) estimation: a port of the rdrobust package 2.1.1 (Calonico, Cattaneo and Titiunik, 2014; Calonico, Cattaneo and Farrell, 2020), used by RdFlex.

score

Linear Neyman-orthogonal scores and the DML2 estimator.

sensitivity

Sensitivity analysis for omitted confounders (Chernozhukov, Cinelli, Newey, Sharma & Syrgkanis, 2022, “Long Story Short”), as in DoubleML.

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).