Binary treatments: ATE, ATTE and effect heterogeneity¶

With a binary treatment the effect can vary across units. The interactive regression model (IRM) leaves the outcome fully flexible,

\[ Y = g_0(D, X) + \varepsilon, \qquad D \in \{0, 1\}, \qquad P(D = 1 \mid X) = m_0(X), \]

and targets the average treatment effect \(\text{ATE} = E[g_0(1, X) - g_0(0, X)]\) or the effect on the treated \(\text{ATTE} = E[g_0(1, X) - g_0(0, X) \mid D = 1]\) with the doubly robust (AIPW) score.

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¶

Treatment is more likely for units with high \(X_1\), and the effect \(\tau(x) = 1 + 0.5\,x_1\) grows with \(X_1\), so the treated benefit more than average: the ATTE exceeds the ATE.

In [2]:

rng = np.random.default_rng(7)
n, p = 4000, 10
X = rng.normal(size=(n, p))
propensity = 1 / (1 + np.exp(-(0.8 * X[:, 0] - 0.5 * X[:, 1])))
d = rng.binomial(1, propensity).astype(float)
tau = 1 + 0.5 * X[:, 0]
y = tau * d + X[:, 0] + 0.5 * X[:, 2] - 0.5 * X[:, 3] + rng.normal(size=n)

print(f"true ATE (sample):  {tau.mean():.3f}")
print(f"true ATTE (sample): {tau[d == 1].mean():.3f}")
print(f"naive difference in means: {y[d == 1].mean() - y[d == 0].mean():.3f}")

rng = np.random.default_rng(7) n, p = 4000, 10 X = rng.normal(size=(n, p)) propensity = 1 / (1 + np.exp(-(0.8 * X[:, 0] - 0.5 * X[:, 1]))) d = rng.binomial(1, propensity).astype(float) tau = 1 + 0.5 * X[:, 0] y = tau * d + X[:, 0] + 0.5 * X[:, 2] - 0.5 * X[:, 3] + rng.normal(size=n) print(f"true ATE (sample): {tau.mean():.3f}") print(f"true ATTE (sample): {tau[d == 1].mean():.3f}") print(f"naive difference in means: {y[d == 1].mean() - y[d == 0].mean():.3f}")

true ATE (sample):  0.985
true ATTE (sample): 1.167
naive difference in means: 1.915

ATE and ATTE¶

IRM takes an outcome regression ml_g (fitted separately for treated and controls) and a propensity classifier ml_m. Cross-validated Lasso and logistic regression are the classic choice; both are built in and pick their penalty as scikit-learn does.

In [3]:

ate = dml_rs.IRM(dml_rs.LassoCV(), dml_rs.LogisticRegressionCV(), score="ATE").fit(X, y, d)
atte = dml_rs.IRM(dml_rs.LassoCV(), dml_rs.LogisticRegressionCV(), score="ATTE").fit(X, y, d)
pd.concat([ate.summary, atte.summary], keys=["ATE", "ATTE"])

ate = dml_rs.IRM(dml_rs.LassoCV(), dml_rs.LogisticRegressionCV(), score="ATE").fit(X, y, d) atte = dml_rs.IRM(dml_rs.LassoCV(), dml_rs.LogisticRegressionCV(), score="ATTE").fit(X, y, d) pd.concat([ate.summary, atte.summary], keys=["ATE", "ATTE"])

Out[3]:

		coef	std err	t	P>|t|	2.5 %	97.5 %
ATE	d	0.9551	0.0373	25.5881	2.0682e-144	0.8820	1.0283
ATTE	d	1.1397	0.0424	26.9000	2.1972e-159	1.0566	1.2227

Group average treatment effects (GATE)¶

gate estimates the average effect within groups given as 0/1 dummy columns, with joint inference. Here: terciles of \(X_1\).

In [4]:

tercile = pd.qcut(X[:, 0], 3, labels=["low X1", "mid X1", "high X1"])
groups = pd.get_dummies(tercile).astype(float)
gate = ate.gate(groups)
table = gate.confint()
table["true"] = [tau[tercile == g].mean() for g in groups.columns]
table

tercile = pd.qcut(X[:, 0], 3, labels=["low X1", "mid X1", "high X1"]) groups = pd.get_dummies(tercile).astype(float) gate = ate.gate(groups) table = gate.confint() table["true"] = [tau[tercile == g].mean() for g in groups.columns] table

Out[4]:

	2.5 %	effect	97.5 %	true
low X1	0.2722	0.4041	0.5361	0.4426
mid X1	0.8732	0.9829	1.0925	0.9803
high X1	1.3482	1.4788	1.6094	1.5322

Best linear predictor of the CATE¶

cate projects the conditional effect \(\tau(x)\) on a basis of your choice. With the basis \((1, x_1)\) the coefficients recover \(\tau(x) = 1 + 0.5\,x_1\).

In [5]:

basis = pd.DataFrame({"intercept": 1.0, "X1": X[:, 0]})
ate.cate(basis).summary

basis = pd.DataFrame({"intercept": 1.0, "X1": X[:, 0]}) ate.cate(basis).summary

Out[5]:

	coef	std err	t	P>|t|	[0.025	0.975]
intercept	0.9695	0.0366	26.5017	9.2681e-155	0.8978	1.0412
X1	0.4756	0.0435	10.9355	7.7999e-28	0.3904	0.5609

Meta-learners¶

For a flexible estimate of \(\tau(x)\) itself, the S-, T-, X-, R- and DR-learners are available. The DR-learner regresses cross-fitted doubly robust pseudo-outcomes on \(X\):

In [6]:

dr = dml_rs.DRLearner(
    outcome_model=dml_rs.LassoCV(),
    propensity_model=dml_rs.LogisticRegressionCV(),
    effect_model=dml_rs.Ols(),
).fit(X, y, d)
tau_hat = dr.effect(X)
print(f"correlation of estimated and true effects: {np.corrcoef(tau_hat, tau)[0, 1]:.3f}")
print(f"average estimated effect: {dr.ate_:.3f}")

dr = dml_rs.DRLearner( outcome_model=dml_rs.LassoCV(), propensity_model=dml_rs.LogisticRegressionCV(), effect_model=dml_rs.Ols(), ).fit(X, y, d) tau_hat = dr.effect(X) print(f"correlation of estimated and true effects: {np.corrcoef(tau_hat, tau)[0, 1]:.3f}") print(f"average estimated effect: {dr.ate_:.3f}")

correlation of estimated and true effects: 0.973
average estimated effect: 0.955