Learners

The nuisance functions are ordinary prediction problems. dml-rs accepts two kinds of learners:

• Built-in learners, implemented in Rust. Cross-fitting with them never touches Python, so folds run in parallel and the GIL is released.
• Any scikit-learn compatible estimator (fit / predict, and predict_proba for classifiers). It is cloned for each fold and called from Rust.

Built-in learners

Regression	Classification	Matches scikit-learn's
Ols(ridge=0.0)		LinearRegression / Ridge(alpha=ridge)
Lasso(alpha)		Lasso(alpha)
LassoCV()		LassoCV()
	LogisticRegression(C)	LogisticRegression(C) (L2)
	LogisticRegressionCV()	LogisticRegressionCV(scoring="neg_log_loss")
RandomForestRegressor(...)	RandomForestClassifier(...)	RandomForest* (same parameters)
GradientBoostingRegressor(...)	GradientBoostingClassifier(...)	GradientBoosting* (same parameters)

Parameter names follow scikit-learn, so switching between the two is a matter of changing the import.

• Linear learners minimise the same objectives as scikit-learn with the same algorithms and stopping rules (coordinate descent for the Lasso, Newton's method for logistic regression), and agree with it to numerical precision.
• Cross-validated learners use scikit-learn's grids, fold schemes and selection rules: LassoCV tries 100 penalties on a log scale from the smallest penalty giving an all-zero solution down to 1/1000 of it, on 5 contiguous folds, and keeps the one with the lowest mean test error; LogisticRegressionCV tries 10 values of C from 1e-4 to 1e4 on stratified folds, scored by log loss (scikit-learn's default from version 1.11; scoring="accuracy" is available). Both refit on all the data with the chosen penalty, and the fitted attributes (alpha_, mse_path_, C_, scores_, ...) use scikit-learn's names.
• Gradient boosting grows the same trees as scikit-learn. When several splits fit a small node equally well, scikit-learn breaks the tie by a random feature order set by its random_state; dml-rs breaks it deterministically. Predictions on the training data are identical; on new data they agree up to the variation scikit-learn itself shows across seeds.
• Random forests follow scikit-learn's parameters and bootstrap scheme with their own random number generator, so they are statistically equivalent rather than identical.

scikit-learn estimators

from sklearn.ensemble import HistGradientBoostingRegressor, HistGradientBoostingClassifier

model = dml_rs.IRM(HistGradientBoostingRegressor(), HistGradientBoostingClassifier()).fit(X, y, d)

Pipelines and tuned search objects (GridSearchCV, ...) work too. Python learners are slower than built-in ones because each fold calls back into Python, but the DML algebra still runs in Rust.

Tuning

GridSearchRegressor and GridSearchClassifier choose the best of several candidate learners by 5-fold cross-validation (mean squared error, or log loss with stratified folds) and refit it. Candidates can be built-in or scikit-learn learners. As a nuisance learner the search runs inside every cross-fitting fold on that fold's training data, which is DoubleML's tune(..., tune_on_folds=True); with scikit-learn candidates it reproduces GridSearchCV with scoring="neg_mean_squared_error" exactly. LassoCV and LogisticRegressionCV tune their penalty the same way, faster.

PythonRust

ml = dml_rs.GridSearchRegressor([dml_rs.RandomForestRegressor(max_depth=d) for d in (3, 5, 8)])
model = dml_rs.PLR(ml, ml).fit(X, y, d)

let grid = || {
    let candidates: Vec<Box<dyn Regressor>> = [3, 5, 8]
        .iter()
        .map(|&depth| Box::new(RandomForestRegressor::new().max_depth(Some(depth))) as Box<dyn Regressor>)
        .collect();
    GridSearchRegressor::new(candidates)
};
let fit = Plr::new(grid(), grid()).fit(&data)?;

Comparing learners

Each fit reports the out-of-fold loss of every nuisance learner, on data the learner did not see: RMSE for regressions and log loss for classifications, in DoubleML's layout (n_rep by n_treat). Comparing learners on these losses is a simple, honest way to choose them.

model.nuisance_loss_
# {'ml_g0': array([[1.05]]), 'ml_g1': array([[1.07]]), 'ml_m': array([[0.65]])}

from sklearn.metrics import mean_absolute_error

def mae(y_true, y_pred):
    keep = ~np.isnan(y_true)  # some learners apply only to part of the sample
    return mean_absolute_error(y_true[keep], y_pred[keep])

model.evaluate_learners(metric=mae)

See the clustered data and diagnostics example.

External predictions

Nuisance predictions computed elsewhere (a model trained on another system, or an expensive model fitted once and reused) can replace any learner, as DoubleML's external_predictions. They must be out-of-fold predictions for the same sample splits, one column per repetition. The learners they replace are not fitted and have no out-of-sample loss. Supported by PLR, IRM, APO, PLIV, IIVM, DID, DIDCS, PLPR, PQ, LPQ and LPLR, with DoubleML's learner names (ml_l, ml_m, ml_g0, ...). LPQ takes all five of its learners or none. LPLR's double cross-fitted ml_M and ml_a also need their inner predictions, ml_M_inner_k and ml_a_inner_k for every outer fold k (a fit reports them in predictions_).

PythonRust

# {treatment: {learner: array (n_obs, n_rep)}}, or {learner: array} for one treatment
model = dml_rs.PLR(ml_l, ml_m).fit(X, y, d, folds=folds, external_predictions={"ml_m": m_hat})

let external = ExternalPredictions::new().with(0, "ml_m", m_hat); // (n_obs, n_rep)
let fit = Plr::new(ml_l, ml_m).fit_with_folds(&data.with_external_predictions(external), &folds)?;

Your own learner in Rust

Implement the Regressor (or Classifier) trait: fit, predict (or predict_proba) and fresh, which returns an unfitted copy with the same settings for the next fold. fit_weighted is optional and needed only by the R-learner. Learners must be Send + Sync so that folds can be fitted on parallel threads.

use dml_core::{Regressor, Result};
use ndarray::{Array1, ArrayView1, ArrayView2};

#[derive(Clone)]
struct Mean(f64);

impl Regressor for Mean {
    fn fit(&mut self, _x: ArrayView2<'_, f64>, y: ArrayView1<'_, f64>) -> Result<()> {
        self.0 = y.mean().unwrap_or(0.0);
        Ok(())
    }
    fn predict(&self, x: ArrayView2<'_, f64>) -> Result<Array1<f64>> {
        Ok(Array1::from_elem(x.nrows(), self.0))
    }
    fn fresh(&self) -> Box<dyn Regressor> {
        Box::new(Mean(0.0))
    }
}