import matplotlib
matplotlib.use("Agg") # headless: render to buffers, never to a screen
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import statsmodels.api as sm
from scipy.spatial import cKDTree
rng = np.random.default_rng(42) # one generator, pinned, drives the whole notebook
plt.rcParams["figure.figsize"] = (9, 5)
plt.rcParams["axes.grid"] = True
plt.rcParams["grid.alpha"] = 0.3
print("pandas", pd.__version__)
print("statsmodels", sm.__version__)
# We deliberately use statsmodels.Logit for the propensity model: logit is the convention
# because its coefficients are interpretable as log-odds of enrolling, and the fit is transparent.
# (scikit-learn's LogisticRegression would also work and gives the same scores up to its default
# L2 penalty; we keep the propensity model unpenalized and explicit here.)nb3.2 — Propensity-Score Matching with Balance Diagnostics
Week 3, Chapter 3.2. Companion notebook.
Maya wants the causal effect of a free financial-literacy program on loan approval. The trap is that the people who choose to enroll already look better to the lender — savvier, higher-credit, higher-income — so the obvious comparison (enrollee approval rate minus non-enrollee approval rate) mixes the program’s effect with a selection bias. This notebook walks the whole selection-on-observables pipeline on a simulated version of Maya’s question, where we engineer a known true effect so we can grade every estimator against the truth:
- Build the world from a known data-generating process (DGP): three confounders —
credit,income,age— drive both enrollment D and approval Y. Compute the true ATT, the naive difference, and confirm their gap is exactly the selection bias. - Estimate the propensity score \hat e(\mathbf{X})=\Pr(D=1\mid\mathbf{X}) by logistic regression (
statsmodels.Logit). - Match: 1:1 nearest-neighbor on the score, with replacement, under a caliper.
- Balance diagnostics: standardized mean differences (SMDs) before vs. after, with a Love plot; check the common-support / overlap with propensity histograms by group.
- Estimate the ATT on the matched sample and compare to the truth and the naive number.
- Your Turn: add an unobserved confounder (motivation) the score cannot see, and watch a perfectly-balanced matching deliver a confidently wrong answer — CIA failing in front of you.
The one rule to carry through the whole notebook: you tune the design by checking balance, never by watching the outcome. The outcome Y enters only after the matching recipe is frozen.
1. Build Maya’s world from a known DGP
The result in one sentence. We write down a true program effect, then let three confounders push savvier people to both enroll and get approved anyway — manufacturing the exact selection bias we will spend the rest of the notebook removing.
We build N=4{,}000 applicants. Three standardized confounders — creditworthiness, income, age — enter a selection index that drives enrollment D through a logistic link:
\Pr(D_i=1\mid\mathbf{X}_i)=\frac{1}{1+\exp(-(\,-0.2 + 0.9\,\text{credit}_i + 0.6\,\text{income}_i + 0.3\,\text{age}_i))}.
The same confounders drive an approval index. Each applicant has two potential outcomes, Y_i(0) (approved if not enrolled) and Y_i(1) (approved if enrolled); enrollment shifts the latent index up by a known amount \tau before thresholding at zero:
y^{(0)}_i = -0.1 + 0.3\,\text{credit}_i + 0.5\,\text{income}_i + 0.2\,\text{age}_i + u_i,\qquad Y_i(0)=\mathbb{1}\{y^{(0)}_i>0\},\quad Y_i(1)=\mathbb{1}\{y^{(0)}_i+\tau>0\}.
We observe only one outcome per person — the fundamental problem of causal inference — Y_i = D_i\,Y_i(1) + (1-D_i)\,Y_i(0). But because we built both potential outcomes, we know the truth and can grade ourselves. The target is the average treatment effect on the treated (ATT), \tau_{\text{ATT}}=\mathbb{E}[Y_i(1)-Y_i(0)\mid D_i=1].
N = 4000
credit = rng.normal(0, 1, N) # standardized creditworthiness
income = rng.normal(0, 1, N)
age = rng.normal(0, 1, N)
# Selection: savvier (higher credit/income/age) applicants enroll more. This IS the bias source.
sel_index = -0.2 + 0.9 * credit + 0.6 * income + 0.3 * age
D = (rng.uniform(size=N) < 1.0 / (1.0 + np.exp(-sel_index))).astype(int)
# Potential outcomes on a latent approval index; treatment shifts the index up by tau.
tau = 0.8
y0_index = -0.1 + 0.3 * credit + 0.5 * income + 0.2 * age + rng.normal(0, 0.3, N)
appr0 = (y0_index > 0).astype(int) # Y(0): approved if NOT enrolled
appr1 = (y0_index + tau > 0).astype(int) # Y(1): approved if enrolled
Y = np.where(D == 1, appr1, appr0) # we observe only one per person
df = pd.DataFrame({"D": D, "Y": Y, "credit": credit, "income": income, "age": age})
print(f"N = {N}, enrolled (D=1): {int(D.sum())}, not enrolled (D=0): {int((1 - D).sum())}")
df.head()1.1 The truth, the naive lie, and the selection bias between them
The naive difference decomposes exactly (Chapter 3.1):
\underbrace{\mathbb{E}[Y\mid D=1]-\mathbb{E}[Y\mid D=0]}_{\text{naive}} = \underbrace{\tau_{\text{ATT}}}_{\text{what we want}} + \underbrace{\big(\mathbb{E}[Y(0)\mid D=1]-\mathbb{E}[Y(0)\mid D=0]\big)}_{\text{selection bias}}.
Because we simulated both potential outcomes, we can compute all three pieces directly and confirm the identity holds to the last decimal. The selection bias is the gap in the untreated potential outcome Y(0) between enrollees and non-enrollees: enrollees would have been approved at higher rates even without the course, because they are the higher-credit, higher-income applicants.
att_true = (appr1[D == 1] - appr0[D == 1]).mean() # truth, known by construction
naive = Y[D == 1].mean() - Y[D == 0].mean() # the contaminated comparison
sel_bias = appr0[D == 1].mean() - appr0[D == 0].mean() # gap in Y(0) between the groups
print(f"True ATT tau_ATT : {att_true:+.4f} (chapter target ~0.322)")
print(f"Naive diff E[Y|D=1] - E[Y|D=0] : {naive:+.4f} (chapter target ~0.629)")
print(f"Selection bias E[Y0|D=1]-E[Y0|D=0] : {sel_bias:+.4f} (chapter target ~0.307)")
print()
print(f"Decomposition check ATT + bias : {att_true + sel_bias:+.4f}")
print(f" naive : {naive:+.4f}")
# The exact decomposition must hold by construction.
assert abs((att_true + sel_bias) - naive) < 1e-9, "naive = ATT + selection bias must hold exactly"
# The naive number is roughly double the truth: a textbook lie.
assert naive > 1.7 * att_true, "naive should be badly inflated by selection bias"
print("\nCHECK PASSED: naive = ATT + selection bias exactly; the naive estimate is ~2x the truth.")2. Estimate the propensity score with logistic regression
The result in one sentence. The propensity score e(\mathbf{X})=\Pr(D=1\mid\mathbf{X}) is a single number that summarizes all of \mathbf{X} for the purpose of predicting treatment; Rosenbaum–Rubin proved that matching on this one scalar is enough to recover the effect (under CIA).
We fit a logit of D on credit, income, age by maximum likelihood and read off the fitted probabilities \hat e(\mathbf{X}_i). Two disciplines, different from Week 2 model-fitting:
- The propensity model is judged by the balance it produces, not by how well it predicts D. A model that predicts D near-perfectly is a disaster — \hat e near 0/1 for everyone means no comparable units survive (an overlap catastrophe).
- The outcome Y never appears in this regression. The score is a function of D and \mathbf{X} alone. That firewall is what keeps the procedure honest.
X = df[["credit", "income", "age"]].to_numpy()
Xc = sm.add_constant(X) # add intercept column
logit_res = sm.Logit(df["D"].to_numpy(), Xc).fit(disp=0)
df["ps"] = logit_res.predict(Xc) # fitted propensity scores e_hat(X)
print(logit_res.summary().tables[1])
print()
print("Estimated log-odds coefficients (const, credit, income, age):")
print(np.round(logit_res.params, 3), " <- recover the DGP signs (-0.2, 0.9, 0.6, 0.3)")
print(f"\nPropensity score range: [{df['ps'].min():.3f}, {df['ps'].max():.3f}]")
print("Scores stay well inside (0, 1): no overlap catastrophe, comparable units exist on both sides.")
# Sanity: the logit recovers the selection coefficients (sign and rough magnitude).
assert logit_res.params[1] > 0.5, "credit coefficient should be clearly positive"
assert df["ps"].between(0.01, 0.99).mean() > 0.9, "scores should mostly avoid the 0/1 extremes"
print("\nCHECK PASSED: logit recovers selection signs; scores avoid the overlap-killing extremes.")3. Nearest-neighbor caliper matching on the score
The result in one sentence. For each treated applicant, find the untreated applicant with the closest \hat e, use that control’s outcome as the counterfactual, and average the within-pair differences — that average is the ATT estimate.
We do 1:1 nearest-neighbor matching with replacement (a single excellent control may serve many treated units) under a caliper: a maximum allowed distance on the score. A common rule of thumb (Austin 2011) sets the caliper at 0.2 times the standard deviation of the score. If a treated unit’s nearest neighbor lies outside the caliper, we drop it rather than force a bad match — the caliper is overlap enforced unit by unit. We use a KD-tree for fast nearest-neighbor lookup.
Why ATT and why with replacement? ATT only asks “what would the enrollees have looked like without the program,” so it only needs counterfactuals for the treated, and there are plenty of untreated units to draw from. Matching every treated unit to its genuinely-closest control (reusing controls freely) minimizes match distance — lower bias — at the cost of correlated matched pairs (a standard-error subtlety we flag but do not chase here).
treated = df[df["D"] == 1].copy()
control = df[df["D"] == 0].copy()
caliper = 0.2 * df["ps"].std()
print(f"Caliper = 0.2 * SD(ps) = {caliper:.4f} on the score scale")
# KD-tree over control propensity scores; query nearest control for each treated unit.
tree = cKDTree(control[["ps"]].to_numpy())
dist, j = tree.query(treated[["ps"]].to_numpy(), k=1) # distance + index of nearest control
dist = dist.ravel(); j = j.ravel()
keep = dist <= caliper # drop treated units with no close neighbor
matched_treated = treated[keep].copy()
matched_control = control.iloc[j[keep]].copy()
n_drop = int((~keep).sum())
print(f"Treated units: {len(treated)} | matched: {int(keep.sum())} | dropped by caliper: {n_drop}")
print(f"Distinct controls reused (with replacement): {matched_control.index.nunique()}")
print(f"Mean within-caliper match distance on the score: {dist[keep].mean():.5f}")4. Balance diagnostics — the audit that earns the estimate
The result in one sentence. Before believing any number, you must verify the matching made the treated and control groups actually comparable on \mathbf{X} — and you measure that with the standardized mean difference (SMD), before vs. after matching, never a t-test.
For covariate X_j with treated/control means \bar X_{j,1},\bar X_{j,0} and pooled SD s_j,
\text{SMD}_j=\frac{\bar X_{j,1}-\bar X_{j,0}}{s_j}.
It is the gap in means in standard-deviation units — scale-free, so credit and dollars sit on the same footing. The convention (Austin 2011, Stuart 2010) is |\text{SMD}_j|<0.1 = acceptable balance. We do not use a t-test: a t-test conflates the size of the imbalance with the sample size, and matching changes the sample size — so a “significant” imbalance can vanish just by dropping units, with no real improvement. The SMD measures the imbalance itself.
def smd(treated_vals, control_vals):
# Standardized mean difference using the pooled (treated+control) SD.
t = np.asarray(treated_vals, dtype=float)
c = np.asarray(control_vals, dtype=float)
pooled_sd = np.sqrt((t.var(ddof=1) + c.var(ddof=1)) / 2.0)
return (t.mean() - c.mean()) / pooled_sd
covs = ["credit", "income", "age", "ps"]
rows = []
for col in covs:
before = smd(treated[col].to_numpy(), control[col].to_numpy())
after = smd(matched_treated[col].to_numpy(), matched_control[col].to_numpy())
rows.append({"covariate": col, "SMD_before": before, "SMD_after": after,
"|after| < 0.1": abs(after) < 0.1})
balance = pd.DataFrame(rows)
print(balance.to_string(index=False, float_format=lambda v: f"{v:+.3f}"))
# Every covariate must collapse from large imbalance to under the 0.1 threshold.
assert (balance["SMD_before"].abs() > 0.2).any(), "some covariate should be badly imbalanced BEFORE"
assert (balance["SMD_after"].abs() < 0.1).all(), "every covariate should balance AFTER matching"
print("\nCHECK PASSED: every SMD collapses from large (pre) to under 0.1 (post). This table is the"
" evidence the estimate deserves to be believed.")4.1 The Love plot
The Love plot (named for Thomas Love) draws each covariate’s absolute SMD before and after matching on a single horizontal axis, with a dashed line at the 0.1 threshold. A good matching shows every “after” dot collapsing leftward, under the line. This single picture is what a referee reads first in a matching paper — not the regression output.
order = balance.reindex(balance["SMD_before"].abs().sort_values().index)
ypos = np.arange(len(order))
fig, ax = plt.subplots(figsize=(8, 4.5))
ax.scatter(order["SMD_before"].abs(), ypos, s=80, color="crimson",
label="before matching", zorder=3)
ax.scatter(order["SMD_after"].abs(), ypos, s=80, color="green",
label="after matching", zorder=3)
for y, b, a in zip(ypos, order["SMD_before"].abs(), order["SMD_after"].abs()):
ax.plot([a, b], [y, y], color="gray", lw=1, zorder=1) # connector showing the collapse
ax.axvline(0.1, color="black", ls="--", lw=1, label="0.1 balance threshold")
ax.set_yticks(ypos)
ax.set_yticklabels(order["covariate"])
ax.set_xlabel("absolute standardized mean difference |SMD|")
ax.set_title("Love plot: covariate imbalance collapses below 0.1 after matching")
ax.legend(loc="lower right")
fig.tight_layout()
print("Every covariate's |SMD| jumps left across the 0.1 line: balance achieved on all observables.")4.2 Common support / overlap
Balance is one half of the license to compare; overlap is the other. The matching could only work because, at most propensity values, both enrollees and non-enrollees exist. We check this directly by plotting the distribution of \hat e for each group. Where the two histograms overlap is the region of common support — the band where a treated unit can find a genuine control. Mass for one group with no counterpart from the other (especially treated mass near \hat e\to 1) signals an overlap problem; that is exactly the region the caliper polices.
fig, ax = plt.subplots()
bins = np.linspace(0, 1, 31)
ax.hist(treated["ps"], bins=bins, alpha=0.55, color="crimson", density=True,
label=f"enrolled (D=1), n={len(treated)}")
ax.hist(control["ps"], bins=bins, alpha=0.55, color="steelblue", density=True,
label=f"not enrolled (D=0), n={len(control)}")
ax.set_xlabel(r"estimated propensity score $\hat e(\mathbf{X})$")
ax.set_ylabel("density")
ax.set_title("Common support: enrollees skew high, but the groups overlap across most of the range")
ax.legend()
fig.tight_layout()
# Quantify overlap: share of treated whose score lies within the control score range.
lo, hi = control["ps"].min(), control["ps"].max()
on_support = treated["ps"].between(lo, hi).mean()
print(f"Control score range: [{lo:.3f}, {hi:.3f}]")
print(f"Share of treated inside the control score range (common support): {on_support:.3f}")
print("Enrollees pile up at higher scores (the selection), but enough controls live there to match.")5. Estimate the ATT on the matched sample
Now — and only now, with the design frozen and balance verified — we let the outcome Y enter. The ATT estimate is the mean within-pair outcome difference on the matched sample:
\hat\tau_{\text{ATT}}=\frac{1}{n_m}\sum_{i\in\text{matched treated}}\big(Y_i - Y_{j(i)}\big),
where j(i) is the matched control for treated unit i. We compare it to the truth and to the naive lie. If matching did its job, \hat\tau_{\text{ATT}} lands near the true ATT and far from the naive number — because it compared each enrollee to a non-enrollee who was just as likely to have enrolled, and therefore just as likely to have been approved anyway. The leftover difference is the program.
pair_diff = matched_treated["Y"].to_numpy() - matched_control["Y"].to_numpy()
att_psm = pair_diff.mean()
print(f"Naive difference : {naive:+.4f} (badly inflated)")
print(f"PSM matched ATT estimate: {att_psm:+.4f} (chapter target ~0.311)")
print(f"True ATT : {att_true:+.4f}")
print()
print(f"|PSM - truth| : {abs(att_psm - att_true):.4f}")
print(f"|naive - truth| : {abs(naive - att_true):.4f}")
# The matched estimate must be much closer to the truth than the naive comparison.
assert abs(att_psm - att_true) < abs(naive - att_true), "PSM must beat naive"
assert abs(att_psm - att_true) < 0.05, "PSM should land within ~0.05 of the true ATT"
print("\nCHECK PASSED: PSM ATT is within ~0.05 of the truth and far closer than the naive estimate."
"\nMatching on the single number e_hat(X) stripped out the selection bias.")5.1 The three numbers, side by side
One picture closes the loop: the naive bar overshoots wildly, the PSM bar sits right on the true-ATT line. Selection bias was the distance between them, and matching collapsed it.
fig, ax = plt.subplots(figsize=(7, 4.5))
labels = ["naive\n(contaminated)", "PSM matched\n(de-confounded)", "true ATT\n(known)"]
vals = [naive, att_psm, att_true]
colors = ["crimson", "green", "black"]
ax.bar(labels, vals, color=colors, alpha=0.8)
ax.axhline(att_true, color="black", ls=":", lw=1.5, label="true ATT")
for i, v in enumerate(vals):
ax.text(i, v + 0.01, f"{v:.3f}", ha="center", fontweight="bold")
ax.set_ylabel("estimated effect on approval probability")
ax.set_title("Naive doubles the truth; PSM recovers it")
ax.legend()
fig.tight_layout()
print("The naive bar is the lie; the PSM bar lands on the dotted true-ATT line.")Your Turn — Add an UNOBSERVED confounder and watch CIA fail
Everything above rested on the Conditional Independence Assumption (CIA): that \mathbf{X} contains all the confounders. CIA is fundamentally untestable — a beautiful balance table is silent about variables you never measured. Here we break it on purpose.
Introduce motivation: the drive to improve one’s finances that makes someone both seek out the free course (raises D) and clean up their application in ways the lender rewards (raises Y). Motivation is a real confounder, but it is not in \mathbf{X} — the propensity model never sees it. By the omitted-variable logic of Chapter 2.5, with \beta_2>0 (motivation raises approval) and \delta_1>0 (motivation raises enrollment), the program effect is biased upward.
Watch what happens: we rebuild the world with motivation driving both D and Y, estimate the score on credit, income, age only, match, and check balance. The observed covariates will balance beautifully (every SMD under 0.1) — and the ATT estimate will still be confidently wrong, because the matched enrollees are more motivated than their matched controls, and the balance table cannot see it.
rng_yt = np.random.default_rng(42) # fresh pinned generator for the "Your Turn" world
Nyt = 4000
credit_u = rng_yt.normal(0, 1, Nyt)
income_u = rng_yt.normal(0, 1, Nyt)
age_u = rng_yt.normal(0, 1, Nyt)
motivation = rng_yt.normal(0, 1, Nyt) # UNOBSERVED: drives both D and Y, never measured
sel_u = -0.2 + 0.9 * credit_u + 0.6 * income_u + 0.3 * age_u + 0.8 * motivation
D_u = (rng_yt.uniform(size=Nyt) < 1.0 / (1.0 + np.exp(-sel_u))).astype(int)
y0_u = (-0.1 + 0.3 * credit_u + 0.5 * income_u + 0.2 * age_u
+ 0.6 * motivation + rng_yt.normal(0, 0.3, Nyt))
appr0_u = (y0_u > 0).astype(int)
appr1_u = (y0_u + tau > 0).astype(int)
Y_u = np.where(D_u == 1, appr1_u, appr0_u)
att_true_u = (appr1_u[D_u == 1] - appr0_u[D_u == 1]).mean()
naive_u = Y_u[D_u == 1].mean() - Y_u[D_u == 0].mean()
dfu = pd.DataFrame({"D": D_u, "Y": Y_u, "credit": credit_u, "income": income_u,
"age": age_u, "motivation": motivation})
# Propensity score WITHOUT motivation -- we cannot see it.
Xu = sm.add_constant(dfu[["credit", "income", "age"]].to_numpy())
dfu["ps"] = sm.Logit(dfu["D"].to_numpy(), Xu).fit(disp=0).predict(Xu)
tu = dfu[dfu["D"] == 1].copy()
cu = dfu[dfu["D"] == 0].copy()
cal_u = 0.2 * dfu["ps"].std()
treeu = cKDTree(cu[["ps"]].to_numpy())
du, ju = treeu.query(tu[["ps"]].to_numpy(), k=1)
du = du.ravel(); ju = ju.ravel()
keepu = du <= cal_u
mtu = tu[keepu].copy()
mcu = cu.iloc[ju[keepu]].copy()
att_psm_u = (mtu["Y"].to_numpy() - mcu["Y"].to_numpy()).mean()
print(f"True ATT (with motivation) : {att_true_u:+.4f}")
print(f"Naive difference : {naive_u:+.4f}")
print(f"PSM ATT (motivation HIDDEN) : {att_psm_u:+.4f} <- BIASED, well above the truth")
print(f"PSM bias = PSM - truth : {att_psm_u - att_true_u:+.4f}")# Balance on the OBSERVED covariates: looks perfect. Balance on the HIDDEN one: still broken.
print("Post-matching SMDs:")
for col in ["credit", "income", "age", "ps"]:
print(f" {col:10s} (observed) : {smd(mtu[col].to_numpy(), mcu[col].to_numpy()):+.3f}")
hidden_smd = smd(mtu["motivation"].to_numpy(), mcu["motivation"].to_numpy())
print(f" {'motivation':10s} (HIDDEN) : {hidden_smd:+.3f} <- still badly imbalanced!")
# The lesson, asserted: observables balance, the estimate is still biased, the hidden gap remains.
obs_balanced = all(abs(smd(mtu[c].to_numpy(), mcu[c].to_numpy())) < 0.1
for c in ["credit", "income", "age"])
assert obs_balanced, "observed covariates should balance"
assert abs(att_psm_u - att_true_u) > 0.08, "PSM should be clearly biased when motivation is hidden"
assert hidden_smd > 0.3, "the hidden confounder should remain imbalanced after matching"
print("\nCHECK PASSED: every OBSERVED covariate is balanced (|SMD|<0.1), yet the ATT is biased upward"
"\nbecause the UNOBSERVED confounder stays imbalanced. This is CIA failing -- and no balance"
"\ntable can ever detect it. Matching balanced what you measured; it is helpless about what you"
"\ndid not.")What you just saw
- Selection on observables works when CIA holds. Matching on the single scalar \hat e(\mathbf{X}) collapsed every covariate SMD from large to under 0.1 and recovered the true ATT (~0.32), while the naive comparison (~0.63) was roughly double the truth.
- Balance, not the outcome, is what you tune. The outcome Y entered only after the matching recipe was frozen — the firewall that prevents specification-searching on your answer.
- CIA is untestable, and matching is still selection-on-observables. The moment a confounder (motivation) sits outside \mathbf{X}, a perfectly balanced matching delivers a confidently wrong answer. An omitted confounder is the omitted-variable bias \beta_2\delta_1 of Week 2 wearing a potential-outcomes costume, and no balance table can see it.
Try next. (1) Tighten the caliper to 0.05\cdot\text{SD}(\hat e) and loosen it to 1.0 — watch how many treated units get dropped and how the ATT and balance move (the bias–variance knob). (2) Switch to matching without replacement (mark each used control and skip it) and see balance degrade as the control pool is spent. (3) Add a useless covariate (pure noise) to the propensity model and confirm it changes little — then add a mismeasured confounder and watch balance suffer. The discipline never changes: make a defensible choice, check balance, and report how much the answer moves as you flip the knobs.