Clustered data and learner diagnostics¶

Observations are often grouped: pupils in schools, workers in firms, repeated observations of the same customer. Outcomes within a group share unobserved shocks, so the observations are not independent and ordinary standard errors are too small. dml-rs implements cluster-robust DML for one-way and two-way clustering (Chiang, Kato, Ma and Sasaki, 2022) as in DoubleML: whole clusters are assigned to folds and the variance sums the score within clusters.

In [1]:

Copied!

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)

Pupils in schools¶

200 schools with 15 pupils each. School-level shocks affect both the treatment and the outcome.

In [2]:

Copied!





def schools(seed, n_schools=200, size=15, theta=0.5):
    rng = np.random.default_rng(seed)
    school = np.repeat(np.arange(n_schools), size)
    n = len(school)
    X = rng.normal(size=(n, 5))
    u, w = rng.normal(size=(2, n_schools))
    d = X[:, 0] + u[school] + rng.normal(size=n)
    y = theta * d + X[:, 1] + w[school] + rng.normal(size=n)
    return X, y, d, school

X, y, d, school = schools(0)
iid = dml_rs.PLR(dml_rs.Ols(), dml_rs.Ols()).fit(X, y, d)
clustered = dml_rs.PLR(dml_rs.Ols(), dml_rs.Ols()).fit(X, y, d, cluster=school)
pd.concat([iid.summary, clustered.summary], keys=["iid", "clustered by school"])
def schools(seed, n_schools=200, size=15, theta=0.5):
    rng = np.random.default_rng(seed)
    school = np.repeat(np.arange(n_schools), size)
    n = len(school)
    X = rng.normal(size=(n, 5))
    u, w = rng.normal(size=(2, n_schools))
    d = X[:, 0] + u[school] + rng.normal(size=n)
    y = theta * d + X[:, 1] + w[school] + rng.normal(size=n)
    return X, y, d, school

X, y, d, school = schools(0)
iid = dml_rs.PLR(dml_rs.Ols(), dml_rs.Ols()).fit(X, y, d)
clustered = dml_rs.PLR(dml_rs.Ols(), dml_rs.Ols()).fit(X, y, d, cluster=school)
pd.concat([iid.summary, clustered.summary], keys=["iid", "clustered by school"])

Out[2]:

		coef	std err	t	P>\|t\|	2.5 %	97.5 %
iid	d	0.4517	0.0177	25.5581	4.4633e-144	0.4171	0.4864
clustered by school	d	0.4562	0.0352	12.9602	2.0569e-38	0.3872	0.5251

The clustered standard error is about twice as large. A small simulation shows which one is right: over repeated samples, a 95% interval should contain the true effect 95% of the time.

In [3]:

Copied!





covered = {"iid": 0, "clustered": 0}
reps = 200
for r in range(reps):
    X, y, d, school = schools(1000 + r)
    for name, cl in [("iid", None), ("clustered", school)]:
        lo, hi = dml_rs.PLR(dml_rs.Ols(), dml_rs.Ols(), random_state=r).fit(X, y, d, cluster=cl).confint()[0]
        covered[name] += lo <= 0.5 <= hi
{name: f"{count / reps:.0%}" for name, count in covered.items()}
covered = {"iid": 0, "clustered": 0}
reps = 200
for r in range(reps):
    X, y, d, school = schools(1000 + r)
    for name, cl in [("iid", None), ("clustered", school)]:
        lo, hi = dml_rs.PLR(dml_rs.Ols(), dml_rs.Ols(), random_state=r).fit(X, y, d, cluster=cl).confint()[0]
        covered[name] += lo <= 0.5 <= hi
{name: f"{count / reps:.0%}" for name, count in covered.items()}

Out[3]:

{'iid': '62%', 'clustered': '94%'}

Two cluster variables, for example firms and years, give two-way clustering: pass both columns. Each variable's clusters are split into 5 folds, so cross-fitting uses the 5 × 5 = 25 combinations, and no firm or year appears in both a training and a test set.

In [4]:

Copied!





rng = np.random.default_rng(2)
firm, year = np.meshgrid(np.arange(40), np.arange(30), indexing="ij")
firm, year = firm.ravel(), year.ravel()
n = len(firm)
Xf = rng.normal(size=(n, 4))
firm_d, firm_y = rng.normal(size=(2, 40))   # firm shocks to treatment and outcome
year_d, year_y = rng.normal(size=(2, 30))   # year shocks
d2 = Xf[:, 0] + firm_d[firm] + year_d[year] + rng.normal(size=n)
y2 = 0.5 * d2 + Xf[:, 1] + firm_y[firm] + year_y[year] + rng.normal(size=n)
two_way = dml_rs.PLR(dml_rs.Ols(), dml_rs.Ols()).fit(Xf, y2, d2, cluster=np.column_stack([firm, year]))
print(two_way)
rng = np.random.default_rng(2)
firm, year = np.meshgrid(np.arange(40), np.arange(30), indexing="ij")
firm, year = firm.ravel(), year.ravel()
n = len(firm)
Xf = rng.normal(size=(n, 4))
firm_d, firm_y = rng.normal(size=(2, 40))   # firm shocks to treatment and outcome
year_d, year_y = rng.normal(size=(2, 30))   # year shocks
d2 = Xf[:, 0] + firm_d[firm] + year_d[year] + rng.normal(size=n)
y2 = 0.5 * d2 + Xf[:, 1] + firm_y[firm] + year_y[year] + rng.normal(size=n)
two_way = dml_rs.PLR(dml_rs.Ols(), dml_rs.Ols()).fit(Xf, y2, d2, cluster=np.column_stack([firm, year]))
print(two_way)

PLR (partialling out)  |  n_obs = 1200, n_folds = 25, n_rep = 1, two-way clustered (30 clusters)
           coef    std err        t    P>|t|      2.5 %     97.5 %
d        0.4401     0.1057    4.163   0.0000     0.2329     0.6472
Nuisance learners (out-of-sample, mean over repetitions):
  ml_l         RMSE     1.8421
  ml_m         RMSE     1.8208

Learner diagnostics¶

How well do the nuisance learners predict? After a fit, nuisance_loss_ holds the out-of-fold RMSE of each regression learner and the log loss of each classifier, in DoubleML's layout (n_rep by n_treat). evaluate_learners scores the cross-fitted predictions with any metric.

In [5]:

Copied!





from sklearn.metrics import mean_absolute_error

rng = np.random.default_rng(4)
n = 3000
X = rng.normal(size=(n, 6))
d = rng.binomial(1, 1 / (1 + np.exp(-np.sin(2 * X[:, 0])))).astype(float)
y = d + np.cos(2 * X[:, 0]) + X[:, 1] ** 2 + rng.normal(size=n)

def mae(y_true, y_pred):
    keep = ~np.isnan(y_true)  # ml_g0 / ml_g1 only apply to controls / treated
    return mean_absolute_error(y_true[keep], y_pred[keep])

rows = {}
for name, (ml_g, ml_m) in {
    "Lasso / logistic (CV)": (dml_rs.LassoCV(), dml_rs.LogisticRegressionCV()),
    "Gradient boosting": (dml_rs.GradientBoostingRegressor(), dml_rs.GradientBoostingClassifier()),
}.items():
    fit = dml_rs.IRM(ml_g, ml_m).fit(X, y, d)
    losses = {f"{k} ({'log loss' if k == 'ml_m' else 'RMSE'})": v[0, 0] for k, v in fit.nuisance_loss_.items()}
    maes = {f"{k} (MAE)": v[0, 0] for k, v in fit.evaluate_learners(learners=["ml_g0", "ml_g1"], metric=mae).items()}
    rows[name] = {"ATE": fit.coef_[0], "std err": fit.se_[0], **losses, **maes}
pd.DataFrame(rows).T
from sklearn.metrics import mean_absolute_error

rng = np.random.default_rng(4)
n = 3000
X = rng.normal(size=(n, 6))
d = rng.binomial(1, 1 / (1 + np.exp(-np.sin(2 * X[:, 0])))).astype(float)
y = d + np.cos(2 * X[:, 0]) + X[:, 1] ** 2 + rng.normal(size=n)

def mae(y_true, y_pred):
    keep = ~np.isnan(y_true)  # ml_g0 / ml_g1 only apply to controls / treated
    return mean_absolute_error(y_true[keep], y_pred[keep])

rows = {}
for name, (ml_g, ml_m) in {
    "Lasso / logistic (CV)": (dml_rs.LassoCV(), dml_rs.LogisticRegressionCV()),
    "Gradient boosting": (dml_rs.GradientBoostingRegressor(), dml_rs.GradientBoostingClassifier()),
}.items():
    fit = dml_rs.IRM(ml_g, ml_m).fit(X, y, d)
    losses = {f"{k} ({'log loss' if k == 'ml_m' else 'RMSE'})": v[0, 0] for k, v in fit.nuisance_loss_.items()}
    maes = {f"{k} (MAE)": v[0, 0] for k, v in fit.evaluate_learners(learners=["ml_g0", "ml_g1"], metric=mae).items()}
    rows[name] = {"ATE": fit.coef_[0], "std err": fit.se_[0], **losses, **maes}
pd.DataFrame(rows).T

Out[5]:

	ATE	std err	ml_g0 (RMSE)	ml_g1 (RMSE)	ml_m (log loss)	ml_g0 (MAE)	ml_g1 (MAE)
Lasso / logistic (CV)	1.0044	0.0678	1.8459	1.7667	0.6848	1.3739	1.3484
Gradient boosting	0.9937	0.0427	1.0465	1.0714	0.6499	0.8271	0.8585

Both estimates are close to the true ATE of 1, but the flexible learners fit the nonlinear nuisances much better, and that buys precision: a clearly smaller standard error. Because the losses are computed on data each learner did not see, comparing them is a simple, honest way to choose learners.