Getting started: the partially linear model¶

This notebook estimates the effect of a continuous treatment \(D\) on an outcome \(Y\) when both depend on covariates \(X\) in unknown, nonlinear ways:

\[ Y = \theta D + g(X) + \varepsilon, \qquad D = m(X) + V . \]

Double machine learning (Chernozhukov et al., 2018) learns \(\ell(X) = E[Y \mid X]\) and \(m(X) = E[D \mid X]\) with any machine-learning method, then regresses the out-of-fold residuals \(Y - \hat\ell(X)\) on \(D - \hat m(X)\). Cross-fitting and the orthogonal score keep \(\hat\theta\) \(\sqrt n\)-consistent and approximately normal even though the nuisance learners converge slowly.

In [1]:

import warnings

import numpy as np
import pandas as pd

import dml_rs

warnings.filterwarnings("ignore", category=FutureWarning)  # scikit-learn deprecation notices
warnings.filterwarnings("ignore", message="IProgress not found")  # tqdm, imported by DoubleML
pd.set_option("display.precision", 4)

import warnings import numpy as np import pandas as pd import dml_rs warnings.filterwarnings("ignore", category=FutureWarning) # scikit-learn deprecation notices warnings.filterwarnings("ignore", message="IProgress not found") # tqdm, imported by DoubleML pd.set_option("display.precision", 4)

Simulated data¶

The true effect is \(\theta = 0.5\). Two covariates confound treatment and outcome nonlinearly: both jump when \(X_1\) crosses zero and both rise with \(X_2^2\).

In [2]:

rng = np.random.default_rng(42)
n, p = 2000, 20
X = rng.normal(size=(n, p))
m = (X[:, 0] > 0) + 0.5 * X[:, 1] ** 2 - 0.5        # E[D | X]
g = 2.0 * (X[:, 0] > 0) + X[:, 1] ** 2               # confounding in Y
d = m + rng.normal(size=n)
y = 0.5 * d + g + rng.normal(size=n)

rng = np.random.default_rng(42) n, p = 2000, 20 X = rng.normal(size=(n, p)) m = (X[:, 0] > 0) + 0.5 * X[:, 1] ** 2 - 0.5 # E[D | X] g = 2.0 * (X[:, 0] > 0) + X[:, 1] ** 2 # confounding in Y d = m + rng.normal(size=n) y = 0.5 * d + g + rng.normal(size=n)

Ignoring the confounders, or adjusting for them linearly, gives a biased answer:

In [3]:

def ols_slope(design, y):
    return np.linalg.lstsq(np.column_stack([np.ones(len(y)), design]), y, rcond=None)[0][1]

print(f"no adjustment:     {ols_slope(d, y):.3f}")
print(f"linear adjustment: {ols_slope(np.column_stack([d, X]), y):.3f}")

def ols_slope(design, y): return np.linalg.lstsq(np.column_stack([np.ones(len(y)), design]), y, rcond=None)[0][1] print(f"no adjustment: {ols_slope(d, y):.3f}") print(f"linear adjustment: {ols_slope(np.column_stack([d, X]), y):.3f}")

no adjustment:     1.350
linear adjustment: 1.227

Double machine learning¶

PLR takes a learner for \(\ell(X)\) (ml_l) and one for \(m(X)\) (ml_m). The built-in gradient boosting runs entirely in Rust; any scikit-learn estimator works too.

In [4]:

model = dml_rs.PLR(
    ml_l=dml_rs.GradientBoostingRegressor(n_estimators=300, learning_rate=0.05),
    ml_m=dml_rs.GradientBoostingRegressor(n_estimators=300, learning_rate=0.05),
    n_folds=5,
    random_state=1,
).fit(X, y, d)
model.summary

model = dml_rs.PLR( ml_l=dml_rs.GradientBoostingRegressor(n_estimators=300, learning_rate=0.05), ml_m=dml_rs.GradientBoostingRegressor(n_estimators=300, learning_rate=0.05), n_folds=5, random_state=1, ).fit(X, y, d) model.summary

Out[4]:

	coef	std err	t	P>|t|	2.5 %	97.5 %
d	0.4706	0.0222	21.1566	2.3980e-99	0.427	0.5142

The full report also shows the out-of-sample error of each nuisance learner:

In [5]:

print(model)

print(model)

PLR (partialling out)  |  n_obs = 2000, n_folds = 5, n_rep = 1
           coef    std err        t    P>|t|      2.5 %     97.5 %
d        0.4706     0.0222   21.157   0.0000     0.4270     0.5142
Nuisance learners (out-of-sample, mean over repetitions):
  ml_l         RMSE     1.1579
  ml_m         RMSE     1.0413

Which learner?¶

The quality of \(\hat\theta\) depends on how well the nuisances are learned. nuisance_loss_ reports each learner's out-of-fold RMSE, so learners can be compared on the same splits. Linear learners cannot capture the confounding here; tree ensembles can.

In [6]:

from sklearn.ensemble import RandomForestRegressor

learners = {
    "OLS": dml_rs.Ols(),
    "LassoCV": dml_rs.LassoCV(),
    "Random forest (scikit-learn)": RandomForestRegressor(n_estimators=200, min_samples_leaf=5, n_jobs=-1),
    "Gradient boosting": dml_rs.GradientBoostingRegressor(n_estimators=300, learning_rate=0.05),
}
rows = {}
for name, learner in learners.items():
    fit = dml_rs.PLR(learner, learner, random_state=1).fit(X, y, d)
    rows[name] = {
        "coef": fit.coef_[0],
        "std err": fit.se_[0],
        "RMSE l(X)": fit.nuisance_loss_["ml_l"][0, 0],
        "RMSE m(X)": fit.nuisance_loss_["ml_m"][0, 0],
    }
pd.DataFrame(rows).T

from sklearn.ensemble import RandomForestRegressor learners = { "OLS": dml_rs.Ols(), "LassoCV": dml_rs.LassoCV(), "Random forest (scikit-learn)": RandomForestRegressor(n_estimators=200, min_samples_leaf=5, n_jobs=-1), "Gradient boosting": dml_rs.GradientBoostingRegressor(n_estimators=300, learning_rate=0.05), } rows = {} for name, learner in learners.items(): fit = dml_rs.PLR(learner, learner, random_state=1).fit(X, y, d) rows[name] = { "coef": fit.coef_[0], "std err": fit.se_[0], "RMSE l(X)": fit.nuisance_loss_["ml_l"][0, 0], "RMSE m(X)": fit.nuisance_loss_["ml_m"][0, 0], } pd.DataFrame(rows).T

Out[6]:

	coef	std err	RMSE l(X)	RMSE m(X)
OLS	1.2234	0.0406	2.2353	1.2975
LassoCV	1.2290	0.0408	2.2320	1.2926
Random forest (scikit-learn)	0.5030	0.0242	1.1973	1.0495
Gradient boosting	0.4706	0.0222	1.1579	1.0413

Repeated sample splitting¶

The estimate depends a little on the random split into folds. With n_rep repetitions, the reported coefficient is the median over repetitions and the standard error accounts for the variation between them (as in DoubleML).

In [7]:

repeated = dml_rs.PLR(
    dml_rs.GradientBoostingRegressor(n_estimators=300, learning_rate=0.05),
    dml_rs.GradientBoostingRegressor(n_estimators=300, learning_rate=0.05),
    n_rep=5,
    random_state=1,
).fit(X, y, d)
print("per repetition:", np.round(repeated.all_coef_[0], 4))
repeated.summary

repeated = dml_rs.PLR( dml_rs.GradientBoostingRegressor(n_estimators=300, learning_rate=0.05), dml_rs.GradientBoostingRegressor(n_estimators=300, learning_rate=0.05), n_rep=5, random_state=1, ).fit(X, y, d) print("per repetition:", np.round(repeated.all_coef_[0], 4)) repeated.summary

per repetition: [0.4706 0.4768 0.4859 0.4683 0.4804]

Out[7]:

	coef	std err	t	P>|t|	2.5 %	97.5 %
d	0.4768	0.0227	20.9593	1.5449e-97	0.4322	0.5214

Same result as DoubleML, faster¶

On the same sample splits, dml-rs gives DoubleML's answer. Here DoubleML uses scikit-learn's GradientBoostingRegressor and dml-rs its own implementation, which grows the same trees. When several splits fit a small node equally well, scikit-learn breaks the tie by a random feature order (set by its random_state), so its own estimate moves slightly from seed to seed; dml-rs agrees to that precision. (With linear learners such as LassoCV the two agree to ~1e-15.)

In [8]:

import time

import doubleml as dml
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.model_selection import KFold

folds = list(KFold(5, shuffle=True, random_state=0).split(X))

rows = {}
for seed in [0, 1]:
    start = time.perf_counter()
    ref = dml.DoubleMLPLR(
        dml.DoubleMLData.from_arrays(X, y, d),
        GradientBoostingRegressor(n_estimators=300, learning_rate=0.05, random_state=seed),
        GradientBoostingRegressor(n_estimators=300, learning_rate=0.05, random_state=seed),
        n_folds=5,
        draw_sample_splitting=False,
    )
    ref.set_sample_splitting([folds]).fit()
    rows[f"DoubleML + scikit-learn (random_state={seed})"] = [ref.coef[0], ref.se[0], time.perf_counter() - start]

start = time.perf_counter()
ours = dml_rs.PLR(
    dml_rs.GradientBoostingRegressor(n_estimators=300, learning_rate=0.05),
    dml_rs.GradientBoostingRegressor(n_estimators=300, learning_rate=0.05),
).fit(X, y, d, folds=folds)
rows["dml-rs"] = [ours.coef_[0], ours.se_[0], time.perf_counter() - start]

pd.DataFrame(rows, index=["coef", "std err", "seconds"]).T

import time import doubleml as dml from sklearn.ensemble import GradientBoostingRegressor from sklearn.model_selection import KFold folds = list(KFold(5, shuffle=True, random_state=0).split(X)) rows = {} for seed in [0, 1]: start = time.perf_counter() ref = dml.DoubleMLPLR( dml.DoubleMLData.from_arrays(X, y, d), GradientBoostingRegressor(n_estimators=300, learning_rate=0.05, random_state=seed), GradientBoostingRegressor(n_estimators=300, learning_rate=0.05, random_state=seed), n_folds=5, draw_sample_splitting=False, ) ref.set_sample_splitting([folds]).fit() rows[f"DoubleML + scikit-learn (random_state={seed})"] = [ref.coef[0], ref.se[0], time.perf_counter() - start] start = time.perf_counter() ours = dml_rs.PLR( dml_rs.GradientBoostingRegressor(n_estimators=300, learning_rate=0.05), dml_rs.GradientBoostingRegressor(n_estimators=300, learning_rate=0.05), ).fit(X, y, d, folds=folds) rows["dml-rs"] = [ours.coef_[0], ours.se_[0], time.perf_counter() - start] pd.DataFrame(rows, index=["coef", "std err", "seconds"]).T

Out[8]:

	coef	std err	seconds
DoubleML + scikit-learn (random_state=0)	0.4797	0.0219	22.4528
DoubleML + scikit-learn (random_state=1)	0.4815	0.0219	22.3588
dml-rs	0.4809	0.0219	2.5784