Causal Inference
Run top to bottom. Uses pandas, statsmodels, and scikit-learn. The data is simulated with a known true effect (120 dollars) so we can see which estimate is right.
import numpy as np, pandas as pd
import matplotlib.pyplot as plt
import statsmodels.formula.api as smf
from sklearn.linear_model import LogisticRegression
plt.rcParams.update({"figure.dpi":110,"axes.grid":True,"grid.alpha":0.25,"font.size":11})
FU, BL, GR, RD = "#a21caf", "#2563eb", "#16a34a", "#dc2626"
TRUE = 120 # the causal effect built into this simulated data
BASE = "https://raw.githubusercontent.com/johnfisher-ai/Statistics-Data-Science-AI-Visual-Book/main/data/"
fn = "causal-inference--loyalty-program.xlsx"
try:
df = pd.read_excel("../../data/" + fn)
except FileNotFoundError:
df = pd.read_excel(BASE + fn)
print(df.shape); print(df.head())
(2000, 5) customer_id prior_spend tenure_months enrolled spend 0 C200000 202.39 42 0 294.92 1 C200001 295.18 39 1 631.32 2 C200002 285.73 7 1 564.11 3 C200003 164.28 14 0 327.24 4 C200004 179.14 48 0 373.30
DEMO 1 · Explore first: the naive comparison and its confounding¶
Enrolled customers really do spend more. But look at who enrolled: they already spent more and had been around longer, before the program existed. That imbalance is the confounding that makes the naive difference a lie.
bal = df.groupby("enrolled")[["prior_spend", "tenure_months", "spend"]].mean().round(1)
print(bal)
naive = df.loc[df.enrolled==1, "spend"].mean() - df.loc[df.enrolled==0, "spend"].mean()
print(f"\nNAIVE difference in mean spend = ${naive:.1f} (true effect is ${TRUE})")
fig, ax = plt.subplots(1, 2, figsize=(11, 3.8))
ax[0].bar(["not enrolled", "enrolled"], bal["spend"], color=[BL, FU])
ax[0].set_title(f"Outcome: this year's spend"); ax[0].set_ylabel("mean spend")
ax[1].bar(["not enrolled", "enrolled"], bal["prior_spend"], color=[BL, FU])
ax[1].set_title("Confounder: last year's spend"); ax[1].set_ylabel("mean prior spend")
plt.tight_layout(); plt.show()
prior_spend tenure_months spend enrolled 0 181.0 28.4 387.1 1 243.7 36.9 582.8 NAIVE difference in mean spend = $195.8 (true effect is $120)
The picture: a causal DAG and the backdoor¶
Before adjusting, draw why the naive number is wrong. In a causal diagram an arrow means 'directly causes'. The confounders prior_spend and tenure_months point into BOTH enrolled and spend, opening a backdoor path enrolled <- prior_spend -> spend that carries non-causal association. Adjusting for the confounders blocks that path, leaving only the causal enrolled -> spend arrow.
import math
from matplotlib.patches import FancyBboxPatch, FancyArrowPatch
FIGW, FIGH, XR, YR = 8.8, 5.2, 10, 7.4
fig, ax = plt.subplots(figsize=(FIGW, FIGH)); ax.set_xlim(0, XR); ax.set_ylim(0, YR); ax.axis("off")
def angle(x0, y0, x1, y1): # on-screen slope in degrees (corrects for aspect ratio)
a = math.degrees(math.atan2((y1-y0)*(FIGH/YR), (x1-x0)*(FIGW/XR)))
return a + 180 if a < -90 else (a - 180 if a > 90 else a)
def add_node(cx, cy, label, fc, ec, w=2.2, h=1.0):
p = FancyBboxPatch((cx-w/2, cy-h/2), w, h, boxstyle="round,pad=0.08,rounding_size=0.15",
fc=fc, ec=ec, lw=1.8, zorder=3)
ax.add_patch(p)
ax.text(cx, cy, label, ha="center", va="center", fontsize=10.5, fontweight="bold", color=ec, zorder=4)
return (p, (cx, cy))
ps = add_node(3.0, 5.6, "prior_spend", "#fdf4ff", "#a21caf") # confounder
tn = add_node(7.0, 5.6, "tenure_months", "#fdf4ff", "#a21caf") # confounder
tr = add_node(1.8, 2.0, "enrolled\n(treatment)", "#dbeafe", "#2563eb")
ou = add_node(8.2, 2.0, "spend\n(outcome)", "#dcfce7", "#16a34a")
# patchA/patchB clip each arrow to the node borders, so arrows actually TOUCH the boxes
def edge(a, b, color, lw=2.0, mut=15, z=2):
ax.add_patch(FancyArrowPatch(a[1], b[1], patchA=a[0], patchB=b[0], shrinkA=1, shrinkB=1,
arrowstyle="-|>", mutation_scale=mut, color=color, lw=lw, zorder=z))
for cf in (ps, tn): # backdoor arrows: confounder -> BOTH (enter via the top edge)
edge(cf, tr, "#a21caf"); edge(cf, ou, "#a21caf")
edge(tr, ou, "#15803d", lw=3.5, mut=22, z=6) # causal effect: thick, enters spend via its left edge
ax.text(5.0, 1.05, "causal effect we want (+120 dollars)", ha="center", fontsize=10, color="#15803d", fontweight="bold")
# "backdoor" labels rotated to lie along the two outer confounder arrows
ax.text(2.28, 3.78, "backdoor", fontsize=8.5, color="#a21caf", fontweight="bold",
rotation=angle(3.0, 5.6, 1.8, 2.0), rotation_mode="anchor", ha="center", va="bottom")
ax.text(7.72, 3.78, "backdoor", fontsize=8.5, color="#a21caf", fontweight="bold",
rotation=angle(7.0, 5.6, 8.2, 2.0), rotation_mode="anchor", ha="center", va="bottom")
# legend, lifted clear of the prior_spend box
ax.plot([0.4, 1.2], [7.05, 7.05], color="#15803d", lw=3.5); ax.text(1.35, 7.05, "causal arrow (enrolled -> spend)", va="center", fontsize=8, color="#15803d")
ax.plot([0.4, 1.2], [6.66, 6.66], color="#a21caf", lw=2); ax.text(1.35, 6.66, "confounding arrows (open the backdoor)", va="center", fontsize=8, color="#86198f")
ax.set_title("The causal DAG: a backdoor path runs through the confounders", fontsize=11.5, fontweight="bold")
ax.text(5.0, 0.25, "Backdoor path: enrolled <- prior_spend -> spend (and via tenure_months). Adjusting for both\n"
"confounders blocks it, so the estimate drops from the naive +196 to +117, close to the true 120 dollars.",
ha="center", fontsize=8.4, color="#334155")
plt.tight_layout(); plt.show()
DEMO 2 · Regression adjustment¶
Put the confounders in the model beside the treatment. The enrolled coefficient is now the effect holding prior spend and tenure fixed, the honest apples-to-apples comparison.
reg = smf.ols("spend ~ enrolled + prior_spend + tenure_months", data=df).fit()
adj = reg.params["enrolled"]; ci = reg.conf_int().loc["enrolled"]
print(f"adjusted effect of enrolling = ${adj:.1f} 95% CI [${ci[0]:.0f}, ${ci[1]:.0f}]")
print(f"(naive was ${naive:.0f}; true is ${TRUE})")
print(reg.summary().tables[1])
adjusted effect of enrolling = $117.1 95% CI [$111, $123]
(naive was $196; true is $120)
=================================================================================
coef std err t P>|t| [0.025 0.975]
---------------------------------------------------------------------------------
Intercept 149.8621 4.714 31.790 0.000 140.617 159.107
enrolled 117.0829 3.108 37.671 0.000 110.988 123.178
prior_spend 0.9105 0.020 45.234 0.000 0.871 0.950
tenure_months 2.5510 0.076 33.674 0.000 2.402 2.700
=================================================================================
DEMO 3 · Propensity scores and IPW¶
A second route to the answer: model the probability of enrolling from the confounders (the propensity score), then reweight so treated and control groups look alike. The overlap of the two score distributions is what makes adjustment possible at all.
X = df[["prior_spend", "tenure_months"]].values
ps = LogisticRegression(max_iter=1000).fit(X, df.enrolled).predict_proba(X)[:, 1]
df["pscore"] = ps
fig, ax = plt.subplots(figsize=(8, 3.8))
ax.hist(df.loc[df.enrolled==0, "pscore"], bins=30, alpha=0.5, color=BL, label="not enrolled")
ax.hist(df.loc[df.enrolled==1, "pscore"], bins=30, alpha=0.5, color=FU, label="enrolled")
ax.set_xlabel("propensity score P(enroll | confounders)"); ax.set_title("Propensity overlap between groups"); ax.legend()
plt.tight_layout(); plt.show()
w = np.where(df.enrolled==1, 1/ps, 1/(1-ps)) # inverse-probability weights
ipw = (np.sum(w*df.enrolled*df.spend)/np.sum(w*df.enrolled)
- np.sum(w*(1-df.enrolled)*df.spend)/np.sum(w*(1-df.enrolled)))
print(f"IPW estimate of the effect = ${ipw:.1f} (regression ${adj:.0f}, true ${TRUE})")
# balance check: weighted confounder means should move closer together
for col in ["prior_spend", "tenure_months"]:
t = np.average(df.loc[df.enrolled==1, col], weights=w[df.enrolled==1])
ctl = np.average(df.loc[df.enrolled==0, col], weights=w[df.enrolled==0])
print(f" weighted {col:14s}: enrolled {t:7.1f} vs control {ctl:7.1f}")
# the same check in standard deviations, and how much a few large weights matter
t = df.enrolled.values == 1
for col in ["prior_spend", "tenure_months"]:
x = df[col].values; sd = np.sqrt((x[t].var() + x[~t].var()) / 2)
before = (x[t].mean() - x[~t].mean()) / sd
after = (np.average(x[t], weights=w[t]) - np.average(x[~t], weights=w[~t])) / sd
print(f" standardized difference in {col:14s}: {before:+.2f} before, {after:+.2f} after weighting")
print(f"largest weight {w.max():.0f}; enrolled customers with a score under 0.05: {(ps[t] < 0.05).sum()}")
rng = np.random.default_rng(0); boot = []
for _ in range(300): # refit the propensity model on each resample
d = df.sample(len(df), replace=True, random_state=int(rng.integers(1e9)))
Xd = d[["prior_spend", "tenure_months"]].values; e = d.enrolled.values == 1
p = LogisticRegression(max_iter=1000).fit(Xd, d.enrolled).predict_proba(Xd)[:, 1]
wd = np.where(e, 1/p, 1/(1-p))
boot.append(np.average(d.spend.values[e], weights=wd[e]) - np.average(d.spend.values[~e], weights=wd[~e]))
lo, hi = np.percentile(boot, [2.5, 97.5])
print(f"IPW bootstrap 95% interval: ${lo:.0f} to ${hi:.0f} (regression's CI: ${ci[0]:.0f} to ${ci[1]:.0f})")
IPW estimate of the effect = $100.9 (regression $117, true $120) weighted prior_spend : enrolled 186.0 vs control 201.8 weighted tenure_months : enrolled 31.0 vs control 31.2 standardized difference in prior_spend : +0.99 before, -0.25 after weighting standardized difference in tenure_months : +0.51 before, -0.01 after weighting largest weight 90; enrolled customers with a score under 0.05: 7
IPW bootstrap 95% interval: $81 to $122 (regression's CI: $111 to $123)
DEMO 4 · Four numbers, one picture¶
The naive estimate stands alone and wrong. Regression lands near the true 120 dollars. IPW lands lower, at about 101, and its bootstrap interval is wide enough to reach 120, because a handful of enrolled customers with tiny propensity scores carry very large weights and prior spend is still not balanced after weighting. The gap between the red bar and the rest is the confounding bias, and the gap between the two adjusted bars is a reminder to check balance and weights before trusting IPW.
vals = {"Naive\n(biased)": naive, "Regression\nadjusted": adj, "Propensity\nIPW": ipw}
fig, ax = plt.subplots(figsize=(7.5, 4))
bars = ax.bar(list(vals), list(vals.values()), color=[RD, GR, FU])
ax.axhline(TRUE, color="#111827", ls="--", lw=1.6); ax.text(2.35, TRUE+4, f"true ${TRUE}", ha="right", fontweight="bold")
for b, v in zip(bars, vals.values()): ax.text(b.get_x()+b.get_width()/2, v+4, f"${v:.0f}", ha="center", fontweight="bold")
ax.set_ylabel("estimated effect of enrolling ($)"); ax.set_title("Adjusting for the confounders removes most of the bias")
plt.tight_layout(); plt.show()
Wrap-up¶
The naive +196 dollars was not noisy, it was systematically wrong, because enrollees were bigger spenders before the program. Blocking the backdoor brought regression to +117 dollars, close to the true +120 dollars. IPW (+101 dollars) removed most of the bias but leaned on a few extreme weights, so its estimate is far less precise. Same discipline every time: draw the confounders, adjust for them, and state the no-unmeasured-confounders assumption you are leaning on. Next in Part XXIV: Survival Analysis.