nb4.2 — Staggered DiD: TWFE vs. Callaway–Sant’Anna

Week 4, Chapter 4.2. Companion notebook.

In Chapter 4.1 you built the two-way fixed-effects (TWFE) regression

Y_{it} = \alpha_i + \lambda_t + \beta D_{it} + \varepsilon_{it},

and trusted the single coefficient \hat\beta as the effect of a policy. Chapter 4.2 is where that trust breaks. When units adopt treatment at different times (staggered timing) and the treatment effect differs across units or grows over time (heterogeneous, dynamic effects), \hat\beta becomes a contaminated weighted average of many little difference-in-differences — some of them forbidden, because they quietly use already-treated units as the control group.

We rebuild Priya’s world. A wave of states adopts a climate-risk-disclosure rule for property insurers; her outcome Y_{it} is the non-renewal rate in high-risk ZIP codes. Disclosure is meant to lower non-renewals, and its bite deepens the longer the rule is in force. Staggered timing plus a building effect — the exact recipe for the crisis.

This notebook does five things, matching the chapter: 1. Build the §3 sign-flip world (two states, no never-treated unit) and watch TWFE come out the wrong sign. 2. Trace the Goodman-Bacon intuition by hand — the forbidden “late vs. already-treated” comparison. 3. Add a never-treated group and estimate group-time ATTs with Callaway–Sant’Anna (differences.ATTgt), recovering the truth. 4. Scale up to a realistic Priya panel (2017 / 2019 / 2021 waves + never-adopters, with noise). 5. Contrast the TWFE event study against the CS event study, visually.

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 pyfixest as pf
from differences import ATTgt   # Callaway-Sant'Anna group-time ATT (differences 0.3.0)

# One generator, pinned, drives the whole notebook (matches the chapter seed).
rng = np.random.default_rng(20260528)

plt.rcParams["figure.figsize"] = (9, 5)
plt.rcParams["axes.grid"] = True
plt.rcParams["grid.alpha"] = 0.3

print("numpy", np.__version__)
print("pandas", pd.__version__)
print("pyfixest", pf.__version__)

1. A staggered DGP with heterogeneous, building treatment effects

The trick we are about to pull. We construct the untreated potential outcome so parallel trends holds perfectly, then add a treatment effect that is negative and grows in magnitude with event time (years since adoption). Disclosure helps, and helps more the longer it has been in force. Because we built the data, we know the truth and can check what each estimator returns against it.

Each state i adopts the rule in a cohort year G_i (its group), or never (G_i=\infty). The treatment switch is D_{it}=\mathbf 1\{t\ge G_i\} — once on, it stays on (an absorbing treatment). Write event time e = t - G_i (negative before adoption, 0 in the adoption year). The per-period effect at event time e\ge 0 is -(e+1): -1 in the first treated year, -2 in the second, and so on.

def make_world(adopt, Tmax, baseline=10.0):
    '''Staggered panel. Untreated potential outcome = baseline everywhere (parallel trends holds
    EXACTLY); per-period effect at event time e>=0 is -(e+1), i.e. -1,-2,-3,... (negative & BUILDING).

    adopt : dict {state_label: adoption_year or np.inf for never-treated}
    Returns a tidy DataFrame; never-treated cohort stored as np.nan (the differences convention).'''
    rows = []
    for state, g in adopt.items():
        for t in range(1, Tmax + 1):
            treated = (g != np.inf) and (t >= g)
            e = t - g
            effect = -(e + 1) if treated else 0.0
            rows.append({
                "state": state,
                "t": t,
                "G": (np.nan if g == np.inf else float(g)),  # never-treated -> NaN
                "D": int(treated),
                "Y": baseline + effect,
            })
    return pd.DataFrame(rows)

def true_mean_att(df):
    '''Average of the TRUE effects over all treated unit-time cells (the honest target).'''
    tr = df.loc[df["D"] == 1].copy()
    return (-((tr["t"] - tr["G"]) + 1)).mean()

2. World A — the sign flip (two states, no never-treated unit)

This is the smallest world where TWFE gets the sign wrong. State E (early) adopts in year 2; state L (late) adopts in year 5; eight years total. By year 5 every state has the rule, so there is no never-treated group at all — and that, as we will see, is exactly the condition that lets the disaster run to completion.

Look at the table below: every treatment effect is negative. A policy that helped in every single treated state-year. The true mean effect is unambiguously and substantially negative.

A = make_world({"E": 2, "L": 5}, Tmax=8)

# Show the treated-outcome table exactly as in the chapter (§3).
tableA = A.pivot(index="t", columns="state", values="Y")
print("World A — treated outcomes (untreated potential outcome = 10 everywhere):")
print(tableA.to_string())
print()
print(f"True mean ATT over treated cells: {true_mean_att(A):+.4f}   (chapter: -3.455)")

2.1 Run plain TWFE on World A

We estimate Y_{it} = \alpha_i + \lambda_t + \beta D_{it} + \varepsilon_{it} with pyfixest.feols, absorbing state and time fixed effects. Watch the sign.

mA = pf.feols("Y ~ D | state + t", data=A)
beta_A = float(mA.coef()["D"])

print(f"World A  TWFE  beta_hat = {beta_A:+.4f}")
print(f"World A  true mean ATT  = {true_mean_att(A):+.4f}")
print()
if beta_A > 0 and true_mean_att(A) < 0:
    print(">>> SIGN FLIP: every true effect is negative, yet TWFE reports a POSITIVE coefficient.")
    print(">>> The policy lowered non-renewals in every treated state-year; TWFE says it raised them.")

3. Why? The Goodman-Bacon intuition — the forbidden comparison

How can demeaning a table of all-negative effects produce a positive number? Goodman-Bacon (2021) proved that the staggered TWFE coefficient is identically a weighted average of every simple 2×2 DiD you could form. With staggered timing those 2×2s come in distinct flavors:

• Type 1 — treated vs. never-treated. Clean. (World A has none: no never-treated state exists.)
• Type 2 — earlier-treated vs. later-treated, before the later group switches on. Clean: the later group is still untreated, a legitimate control. (Here: L is the clean control for E during years 2–4.)
• Type 3 — later-treated vs. earlier-treated, after the earlier group switched on. Forbidden: the “control” (E) is already treated and its own effect is still moving. (Here: E is the contaminated control for L during years 5–8.)

Let us compute the Goodman-Bacon decomposition by hand for World A — just two 2×2 blocks — and watch the forbidden one come out the wrong sign and then dominate the weighted average.

# Goodman-Bacon (2021): with two timing groups and NO never-treated unit, the pooled beta_hat is
# EXACTLY a weighted average of just two 2x2 DiDs. Split the timeline into three windows:
#   PRE  = {year 1}      (before E adopts)
#   MID  = {2, 3, 4}     (E treated, L not yet  -> L is a CLEAN control here)
#   POST = {5, 6, 7, 8}  (both treated          -> E is an ALREADY-TREATED, forbidden control)
g = A.set_index(["state", "t"])["Y"]
def wmean(state, periods):
    return np.mean([g[(state, t)] for t in periods])

PRE, MID, POST = [1], [2, 3, 4], [5, 6, 7, 8]

# Block 1 (Type-2, CLEAN): treated = E, control = L, over PRE -> MID (L still untreated).
did_clean = (wmean("E", MID) - wmean("E", PRE)) - (wmean("L", MID) - wmean("L", PRE))
# Block 2 (Type-3, FORBIDDEN): treated = L, control = already-treated E, over MID -> POST.
did_forbidden = (wmean("L", POST) - wmean("L", MID)) - (wmean("E", POST) - wmean("E", MID))

print(f"Clean (Type-2)  2x2  [E treated, L clean control] : {did_clean:+.2f}   (right sign)")
print(f"Forbidden (Type-3) 2x2 [L treated, already-treated E as control] : {did_forbidden:+.2f}   (WRONG sign!)")
print()
print("Why the forbidden block is +1.0:  across POST, E's own effect keeps DEEPENING (-4..-7), so E falls")
print(f"  by {wmean('E', POST) - wmean('E', MID):+.1f} while newly-treated L falls only {wmean('L', POST) - wmean('L', MID):+.1f}.")
print("  L falling SLOWER than its 'control' makes L look like it ROSE +1.0 — pure contamination.")

# Goodman-Bacon's punchline: beta_hat = w_clean * did_clean + w_forbidden * did_forbidden, weights sum to 1
# and are set by TIMING-VARIANCE alone (not by which comparison is trustworthy). Recover them:
w_clean = (beta_A - did_forbidden) / (did_clean - did_forbidden)
w_forbidden = 1.0 - w_clean

print(f"Goodman-Bacon weights:  clean = {w_clean:.2f}   forbidden = {w_forbidden:.2f}")
print(f"Identity check: {w_clean:.2f}*({did_clean:+.1f}) + {w_forbidden:.2f}*({did_forbidden:+.1f}) "
      f"= {w_clean*did_clean + w_forbidden*did_forbidden:+.2f}  ==  TWFE beta = {beta_A:+.2f}")
print()
print(">>> The forbidden comparison carries 80% of the weight and is positive, so it DOMINATES and drags")
print(">>> beta_hat to +0.40. The weighting cares only about treatment timing, not about clean-vs-forbidden.")

Hold the mechanism. TWFE used an already-treated unit (E), whose own effect was still changing, as a control for a newly-treated unit (L), and the changing effect of the contaminated control got subtracted from the thing we wanted to measure. With no never-treated state to dilute it, that one forbidden comparison dominates the pooled average and drags \hat\beta positive. The fix in §5 is to refuse to ever form a Type-3 comparison.

4. World B — add a never-treated state, contain the catastrophe

Add a single genuine never-treated state N. Now clean (Type-1) comparisons exist. We use the six-year version (so the cohort dynamics stay short and the numbers match the chapter exactly). Two things should happen:

1. TWFE’s sign is rescued — but it is still badly attenuated (the clean comparisons fight the forbidden ones to a draw that lands near -1.23 instead of the true -2.571).
2. The clean group-time estimator (next section) recovers the truth exactly.

B = make_world({"E": 2, "L": 5, "N": np.inf}, Tmax=6)

mB = pf.feols("Y ~ D | state + t", data=B)
beta_B = float(mB.coef()["D"])

print(f"World B  TWFE beta_hat = {beta_B:+.4f}   (chapter: -1.23, attenuated; RIGHT sign now)")
print(f"World B  true mean ATT = {true_mean_att(B):+.4f}   (chapter: -2.571)")
print()
print("Sign is rescued, magnitude is not: TWFE recovers less than half the true effect.")

5. The cure: Callaway–Sant’Anna group-time ATTs (clean controls)

Callaway and Sant’Anna (2021) estimate a group-time effect \text{ATT}(g,t): the average effect at calendar time t on the cohort that first adopted in year g. Each one is an honest 2×2 DiD against a clean control — never-treated or not-yet-treated units, never an already-treated unit:

\widehat{\text{ATT}}(g,t) = \big(\bar Y_{g,t} - \bar Y_{g,\,g-1}\big) - \big(\bar Y_{C,t} - \bar Y_{C,\,g-1}\big).

We use the differences package. Its ATTgt wants a MultiIndex (entity, time) DataFrame and a cohort column in which never-treated units are coded np.nan (which is exactly how make_world stores cohort G).

# differences.ATTgt API:  ATTgt(data=<MultiIndex (entity, time)>, cohort_column=...).fit(formula="Y", ...)
B_idx = B.set_index(["state", "t"]).copy()          # MultiIndex (entity, time); never-treated G is NaN

att_B = ATTgt(data=B_idx, cohort_column="G")
gt_B = att_B.fit(
    formula="Y",                       # outcome only; parallel trends, no covariates
    control_group="never_treated",     # clean controls = the never-treated state N
    progress_bar=False,
)

# group-time ATT(g,t) grid: gt_B is an ATGgtResult; .to_pandas() returns the table.
gt_df = gt_B.to_pandas()
att_col = [c for c in gt_df.columns if c[-1] == "ATT"][0]   # ('ATTgtElements', '', 'ATT')
gt_tbl = gt_df[[att_col]].copy()
gt_tbl.columns = ["ATT"]
gt_tbl = gt_tbl.reset_index()                               # index: cohort, base_period, time
gt_tbl["event_time"] = gt_tbl["time"] - gt_tbl["cohort"]

print("Callaway-Sant'Anna group-time ATT(g, t)  [World B, clean never-treated control]:")
for _, r in gt_tbl.iterrows():
    tag = "" if r["event_time"] >= 0 else "  (pre-period placebo)"
    print(f"   ATT(g={int(r['cohort'])}, t={int(r['time'])}) = {r['ATT']:+.2f}   (event time e={int(r['event_time']):+d}){tag}")
print()
print("Post-adoption: cohort E (g=2) recovers its true -1,-2,-3,-4,-5; cohort L (g=5) recovers -1,-2. Exactly.")

5.1 Aggregate to a single overall ATT

The group-time grid is the honest output; now we average it on purpose. The simple / overall aggregation collapses the clean group-time effects (weighted by cohort size) into one number — the honest replacement for the broken pooled \hat\beta.

agg_simple_B = att_B.aggregate("simple", overall=True)
overall_B = float(agg_simple_B.iloc[0, 0])

print(f"World B  Callaway-Sant'Anna overall ATT = {overall_B:+.4f}   (chapter: -2.571)")
print(f"World B  true mean ATT                  = {true_mean_att(B):+.4f}")
print(f"World B  pooled TWFE beta_hat           = {beta_B:+.4f}   (attenuated)")
print()
print("CS recovers the truth to the decimal; TWFE is off by more than a full point.")
assert abs(overall_B - true_mean_att(B)) < 1e-6, "CS should recover the true overall ATT exactly here"

6. World C — Priya’s realistic panel (waves of 2017 / 2019 / 2021, plus noise)

The two-state worlds expose the mechanism, but Priya’s real data has many states, several adoption waves, idiosyncratic state levels, common time shocks, and sampling noise. We build it: 40 states, years 2013–2024, cohorts adopting in 2017, 2019, 2021, and a block of never-adopters.

The building effect here is -(e+1)\cdot 0.5 at event time e\ge 0 (negative, growing). We add state fixed effects, a mild common time trend, and Gaussian noise — none of which changes the lesson, but all of which makes the contrast realistic. With a genuine never-treated block present, TWFE keeps the right sign but stays attenuated; CS recovers the truth.

def make_priya_panel(rng, n_states=40, years=range(2013, 2025), slope=0.5):
    '''Realistic staggered panel: cohorts 2017/2019/2021 + never-treated; building effect -(e+1)*slope.'''
    labels = (["G2017"] * 8 + ["G2019"] * 10 + ["G2021"] * 10 + ["NEVER"] * 12)
    gmap = {"G2017": 2017, "G2019": 2019, "G2021": 2021, "NEVER": np.inf}
    rows = []
    for i in range(n_states):
        lab = labels[i % len(labels)]
        g = gmap[lab]
        state_fe = rng.normal(0.0, 1.0)             # permanent state differences
        for t in years:
            time_fe = 0.05 * (t - 2013)             # common (small) calendar trend
            treated = (g != np.inf) and (t >= g)
            e = t - g
            effect = -(e + 1) * slope if treated else 0.0
            y = 10.0 + state_fe + time_fe + effect + rng.normal(0.0, 0.3)
            rows.append({
                "state": f"s{i:02d}", "year": int(t),
                "G": (np.nan if g == np.inf else float(g)),
                "D": int(treated), "Y": y,
            })
    return pd.DataFrame(rows)

SLOPE = 0.5
C = make_priya_panel(rng, slope=SLOPE)
n_cohort = C.drop_duplicates("state").groupby("G", dropna=False).size()
print("States per cohort (NaN = never-treated):")
print(n_cohort.to_string())
print(f"\nPanel: {C['state'].nunique()} states x {C['year'].nunique()} years = {len(C)} rows")

6.1 The honest target, and what TWFE returns

The true mean ATT averages the true cell effects over all treated state-years. Then we run the static pooled TWFE.

def true_mean_att_C(df, slope):
    tr = df.loc[df["D"] == 1].copy()
    return (-((tr["year"] - tr["G"]) + 1) * slope).mean()

mC = pf.feols("Y ~ D | state + year", data=C)
beta_C = float(mC.coef()["D"])
true_C = true_mean_att_C(C, SLOPE)

print(f"World C  static TWFE beta_hat = {beta_C:+.4f}")
print(f"World C  true mean ATT        = {true_C:+.4f}")
print(f"          attenuation: TWFE recovers {100*beta_C/true_C:.0f}% of the truth (right sign, too small)")

6.2 Callaway–Sant’Anna on the realistic panel

Same recipe: MultiIndex (state, year), cohort column G (never-treated = NaN), never-treated controls. We pull the overall ATT and the event-study aggregation (effects by event time e=t-g, averaged across cohorts). The e<0 coefficients are genuine pre-trend / placebo checks; the e\ge 0 coefficients trace the dynamic build-up.

C_idx = C.set_index(["state", "year"]).copy()

att_C = ATTgt(data=C_idx, cohort_column="G")
att_C.fit(formula="Y", control_group="never_treated", progress_bar=False)

cs_overall_C = float(att_C.aggregate("simple", overall=True).iloc[0, 0])

ev_C = att_C.aggregate("event")          # index = relative_period (event time e)
# Column keys are MultiIndex; ATT sits under last level "ATT", its SE under last level "std_error".
cs_ev = ev_C[[c for c in ev_C.columns if c[-1] == "ATT"][0]]
cs_ev_se = ev_C[[c for c in ev_C.columns if c[-1] == "std_error"][0]]

print(f"World C  CS overall ATT = {cs_overall_C:+.4f}   (true mean ATT = {true_C:+.4f})")
print(f"World C  TWFE beta_hat  = {beta_C:+.4f}   (attenuated)")
print()
print("CS event-study coefficients (relative_period : ATT):")
for e, v in cs_ev.items():
    flag = "  <- pre-period (should be ~0)" if e < 0 else ""
    print(f"   e={int(e):+d}: {v:+.3f}{flag}")

6.3 The naive TWFE event study, for contrast

The Chapter 4.1 event study replaces D_{it} with relative-time dummies \mathbf 1\{t-G_i=e\} and reads off one coefficient per event time. We estimate it with pyfixest’s i(), dropping e=-1 as the reference and letting the never-treated states (relative time undefined) form the baseline control. Under staggered timing with dynamic effects, Sun–Abraham (2021) showed these leads/lags bleed into each other through the forbidden comparisons — so the post-period dynamics are distorted relative to the clean CS curve.

Cev = C.copy()
# relative time for adopters; never-treated -> -1 so they sit in the (excluded) reference bin = clean baseline
Cev["rel"] = np.where(Cev["G"].notna(), Cev["year"] - Cev["G"], np.nan)
Cev["rel_f"] = Cev["rel"].fillna(-1).astype(int)

m_es = pf.feols("Y ~ i(rel_f, ref=-1) | state + year", data=Cev)
es = m_es.coef()

# Parse "rel_f::E" -> integer event time E
twfe_ev = {}
for name, val in es.items():
    if name.startswith("rel_f::"):
        twfe_ev[int(name.split("::")[1])] = float(val)
twfe_ev = pd.Series(twfe_ev).sort_index()

print("Naive TWFE event-study coefficients (event_time : beta_e), e=-1 is the omitted reference:")
for e, v in twfe_ev.items():
    print(f"   e={int(e):+d}: {v:+.3f}")

7. The picture: TWFE event study vs. CS event study vs. the truth

We overlay three curves over event time: the true building effect -(e+1)\cdot 0.5, the Callaway–Sant’Anna event study (clean controls), and the naive TWFE event study. The CS curve hugs the truth and its pre-period sits flat at zero; the TWFE curve drifts off it in the post-period because of cohort-bleed. The right panel shows the headline numbers side by side: a single honest CS overall ATT against the attenuated pooled TWFE \hat\beta, both versus the true mean.

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 5))

# --- left: event-study curves ---
e_grid = np.arange(int(min(cs_ev.index.min(), twfe_ev.index.min())),
                   int(max(cs_ev.index.max(), twfe_ev.index.max())) + 1)
truth = np.array([-(e + 1) * SLOPE if e >= 0 else 0.0 for e in e_grid])

ax1.axhline(0, color="0.6", lw=0.8)
ax1.axvline(-0.5, color="0.6", lw=0.8, ls=":")
ax1.plot(e_grid, truth, "k--", lw=2, label="true effect $-(e{+}1)\\cdot0.5$")
ax1.plot(cs_ev.index.astype(int), cs_ev.values, "o-", color="C0", lw=2,
         label="Callaway–Sant'Anna (clean)")
ax1.plot(twfe_ev.index.astype(int), twfe_ev.values, "s--", color="C3", lw=1.6,
         label="naive TWFE event study")
ax1.set_xlabel("event time  $e = t - G_i$  (years since adoption)")
ax1.set_ylabel("effect on non-renewal rate")
ax1.set_title("Event study: CS hugs the truth; naive TWFE drifts")
ax1.legend(loc="lower left", fontsize=9)

# --- right: headline numbers ---
names = ["true\nmean ATT", "Callaway–\nSant'Anna", "pooled\nTWFE"]
vals = [true_C, cs_overall_C, beta_C]
colors = ["k", "C0", "C3"]
bars = ax2.bar(names, vals, color=colors, alpha=0.8)
ax2.axhline(0, color="0.4", lw=0.8)
ax2.axhline(true_C, color="k", ls="--", lw=1, alpha=0.6)
for b, v in zip(bars, vals):
    ax2.text(b.get_x() + b.get_width()/2, v - 0.06, f"{v:+.3f}",
             ha="center", va="top", fontsize=10, color="white", fontweight="bold")
ax2.set_ylabel("overall ATT (non-renewal rate)")
ax2.set_title("Headline: CS recovers the truth; TWFE is attenuated")

fig.tight_layout()
fig.savefig("nb4.2_staggered_twfe_vs_cs.png", dpi=110, bbox_inches="tight")
print("Saved figure: nb4.2_staggered_twfe_vs_cs.png")
print(f"\nSide by side:  true = {true_C:+.4f} | CS = {cs_overall_C:+.4f} | TWFE = {beta_C:+.4f}")

What to carry forward

1. Pooled TWFE is not a safe default under staggered adoption. With staggered timing and heterogeneous/dynamic effects, \hat\beta is a contaminated weighted average. In World A (no clean control) it came out the wrong sign (+0.40 on all-negative effects); in Worlds B and C it kept the sign but was badly attenuated.
2. The root cause is the forbidden comparison — an already-treated unit (E) pressed into control duty for a newly-treated unit (L). When E’s own effect is still moving, that motion is mis-read as a counterfactual trend and subtracted from L’s estimate.
3. The cure is one idea: clean controls + transparent aggregation. Callaway–Sant’Anna estimate a grid of group-time \text{ATT}(g,t) against never- or not-yet-treated controls, then average on purpose. It recovered -2.571 exactly in World B and tracked the truth in noisy World C, where TWFE could not.
4. Report the diagnostics, not just the estimate. Show the event study whose pre-periods genuinely test parallel trends, and name your control group (never- vs. not-yet-treated) and aggregation weights explicitly.

Reconciliation note. Worlds A and B reproduce the chapter’s hand-built numbers exactly (TWFE +0.40 / -1.227; true mean -3.455 / -2.571; CS overall -2.571). World C is a new noisy panel built with rng=np.random.default_rng(20260528); its numbers differ from A/B by construction but the qualitative result is identical — TWFE attenuated, CS recovers the truth, sign always right once a never-treated group exists.

---

Your Turn

1. Kill the never-treated group in World C. Recode the NEVER block as a late 2023 cohort so no state is ever a clean control at the longest horizons. Re-run TWFE and CS. Does TWFE’s attenuation worsen? Does CS refuse to identify the longest event times (no clean control left)?
2. Switch the control group. Re-fit ATTgt with control_group="not_yet_treated". Compare the group-time grid and overall ATT to the never_treated version. Which standard errors are tighter, and why might the not-yet-treated comparison be more credible in real data?
3. Make the effect constant instead of building. Change make_priya_panel’s effect to a flat -1.5 for all treated cells. Re-run static TWFE. Confirm that with homogeneous effects TWFE is no longer attenuated — pin down which of the two crisis ingredients you just switched off.