Rebuild the LaLonde sample
!pip install -q numpy pandas statsmodels scikit-learn matplotlib seaborn causaldata
import numpy as np, pandas as pd
import statsmodels.formula.api as smf
from sklearn.neighbors import NearestNeighbors
import matplotlib.pyplot as plt, seaborn as sns
from causaldata import nsw_mixtape, cps_mixtape
rng = np.random.default_rng(6)
sns.set_theme(style="whitegrid")
nsw = nsw_mixtape.load_pandas().data
obs = pd.concat([nsw[nsw.treat == 1], cps_mixtape.load_pandas().data], ignore_index=True)
obs["u74"] = (obs.re74 == 0).astype(int); obs["u75"] = (obs.re75 == 0).astype(int)
PS_FORMULA = ("treat ~ age + I(age**2) + I(age**3) + educ + I(educ**2) + marr + nodegree"
" + black + hisp + re74 + re75 + u74 + u75 + educ:re74")
obs["ps"] = smf.logit(PS_FORMULA, data=obs).fit(disp=0).predict(obs)
obs["lps"] = np.log(obs.ps / (1 - obs.ps)) # match on the log-odds: better spread
COVS = ["age", "educ", "black", "hisp", "marr", "nodegree", "re74", "re75", "u74", "u75"]
EXPERIMENTAL = 1_794
Nearest-neighbour propensity score matching
For each trainee, find the CPS respondent(s) with the closest propensity score, and compare earnings. The average of those differences estimates the ATT. We match with replacement: one good control can serve several trainees.
def psm_att(data, k=1, score="lps"):
t, c = data[data.treat == 1], data[data.treat == 0]
nn = NearestNeighbors(n_neighbors=k).fit(c[[score]])
_, idx = nn.kneighbors(t[[score]])
matched_y0 = c.re78.to_numpy()[idx].mean(axis=1)
return (t.re78.to_numpy() - matched_y0).mean(), c.iloc[idx.ravel()]
for k in (1, 5):
att, _ = psm_att(obs, k=k)
print(f"PSM, {k} neighbour(s): ATT = ${att:,.0f} (experimental ${EXPERIMENTAL:,})")
One neighbour gives low bias but high variance; more neighbours lower variance but use worse matches. This bias–variance trade-off returns with bandwidths in RDD (Module 10).
Did matching balance the covariates?
Matching is successful when the matched controls look like the treated on every covariate, not when the estimate looks "right". Compare standardised differences before and after.
def std_diff(t, c, cols):
return pd.Series({v: (t[v].mean() - c[v].mean()) / np.sqrt((t[v].var() + c[v].var()) / 2) for v in cols})
treated = obs[obs.treat == 1]
_, matched_controls = psm_att(obs, k=1)
bal = pd.DataFrame({
"before": std_diff(treated, obs[obs.treat == 0], COVS),
"after": std_diff(treated, matched_controls, COVS),
}).round(3)
print(bal)
bal.abs().plot.barh(figsize=(8, 4)); plt.axvline(0.1, ls="--", color="red")
plt.title("|Standardised difference| before and after matching"); plt.show()
Mahalanobis matching
Instead of collapsing X into one score, match directly on the covariates using the Mahalanobis distance, which accounts for their scales and correlations. It works well with few covariates; with many, good matches become rare (the curse of dimensionality), and the propensity score wins.
def mahalanobis_att(data, cols):
t, c = data[data.treat == 1], data[data.treat == 0]
VI = np.linalg.inv(np.cov(data[cols].to_numpy(dtype=float), rowvar=False))
nn = NearestNeighbors(n_neighbors=1, metric="mahalanobis", metric_params={"VI": VI}, algorithm="brute")
nn.fit(c[cols].to_numpy(dtype=float))
_, idx = nn.kneighbors(t[cols].to_numpy(dtype=float))
return (t.re78.to_numpy() - c.re78.to_numpy()[idx[:, 0]]).mean()
print(f"Mahalanobis ATT: ${mahalanobis_att(obs, COVS):,.0f}")
Uncertainty with the bootstrap
Matching estimators have awkward analytical variances. A practical answer is the bootstrap: resample units with replacement, redo the whole procedure (including the propensity score), and use the spread of the estimates. (Abadie and Imbens showed the naive bootstrap is not exactly valid for 1-NN matching, so treat this as an approximation; it works better for k > 1 and for weighting.)
def one_bootstrap(seed):
b = obs.sample(frac=1, replace=True, random_state=seed)
b = b.assign(ps=smf.logit(PS_FORMULA, data=b).fit(disp=0).predict(b))
b = b.assign(lps=np.log(b.ps / (1 - b.ps)))
return psm_att(b, k=5)[0]
boot = np.array([one_bootstrap(s) for s in range(100)]) # use 500+ in real work
est = psm_att(obs, k=5)[0]
print(f"PSM (k=5) ATT ${est:,.0f}, bootstrap SE ${boot.std():,.0f}, "
f"95% CI [${np.percentile(boot, 2.5):,.0f}, ${np.percentile(boot, 97.5):,.0f}]")
On the LaLonde data, careful matching lands near the experimental benchmark, while naive and linear regression comparisons were far off. Dehejia and Wahba's lesson: with overlap and a rich set of pre-treatment earnings, selection on observables can work. Smith and Todd's reply: results are sensitive to the sample and specification. Report several estimators, and run sensitivity analysis (Module 14).
Match exactly on black and nodegree (only compare trainees with CPS respondents in the same cell) and
by nearest propensity score within each cell. Does the estimate move?
Show a solution
parts = []
for (b, nd), cell in obs.groupby(["black", "nodegree"]):
if cell.treat.sum() > 0 and (cell.treat == 0).sum() > 0:
att, _ = psm_att(cell, k=1)
parts.append((att, cell.treat.sum()))
exact_att = sum(a * n for a, n in parts) / sum(n for _, n in parts)
print(f"Exact-on-(black, nodegree) + PSM: ${exact_att:,.0f}")
Checklist and recap
- I can implement PSM and Mahalanobis matching for the ATT.
- I judge matching by covariate balance, not by the estimate.
- I bootstrap the full procedure for uncertainty and compare with a benchmark when one exists.