{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "cells": [
  {
   "id": "m11-l3-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 11.3 · Staggered Adoption: When TWFE Fails, and What to Use Instead\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- Show by simulation that TWFE can be badly biased when treatment is staggered and effects grow over time\n",
    "- Explain \"forbidden comparisons\" (already-treated units used as controls)\n",
    "- Estimate cohort-specific 2×2 effects against never-treated units (the Callaway–Sant'Anna idea)\n",
    "- Apply Gardner's two-stage DiD (did2s) to the castle-doctrine laws and homicide\n",
    "\n",
    "> Read the full explanations on the course website, then run each cell here in order (**Runtime → Run all** works on a fresh runtime)."
   ]
  },
  {
   "id": "m11-l3-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Staggered roll-outs"
   ]
  },
  {
   "id": "m11-l3-2",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q pyfixest numpy pandas matplotlib causaldata\n",
    "import numpy as np, pandas as pd\n",
    "import pyfixest as pf\n",
    "import matplotlib.pyplot as plt\n",
    "rng = np.random.default_rng(3)"
   ]
  },
  {
   "id": "m11-l3-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A simulation where TWFE gets the sign wrong"
   ]
  },
  {
   "id": "m11-l3-4",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** Can TWFE report a negative effect when every unit's true effect is positive?"
   ]
  },
  {
   "id": "m11-l3-5",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def staggered(never_treated=0, units_per_cohort=30, T=12, cohorts=(2, 5, 8)):\n",
    "    rows = []\n",
    "    gs = [g for g in cohorts for _ in range(units_per_cohort)] + [0] * never_treated\n",
    "    for i, g in enumerate(gs):\n",
    "        a = rng.normal(0, 1)\n",
    "        for t in range(T):\n",
    "            treated = int(g > 0 and t >= g)\n",
    "            tau = (t - g + 1) * 1.0 if treated else 0.0           # effect grows with time since adoption\n",
    "            rows.append((i, t, g, treated, tau, a + 0.3 * t + tau + rng.normal(0, 0.5)))\n",
    "    return pd.DataFrame(rows, columns=[\"unit\", \"t\", \"g\", \"D\", \"tau\", \"y\"])\n",
    "\n",
    "sim = staggered()\n",
    "twfe = pf.feols(\"y ~ D | unit + t\", data=sim, vcov={\"CRV1\": \"unit\"})\n",
    "print(f\"True average effect on the treated: {sim.tau[sim.D == 1].mean():.2f}\")\n",
    "print(f\"TWFE estimate:                      {twfe.coef()['D']:.2f}\")"
   ]
  },
  {
   "id": "m11-l3-6",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Clean comparisons: cohort-by-cohort DiD"
   ]
  },
  {
   "id": "m11-l3-7",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** What is the effect for each adoption cohort at each point in time, using only clean controls?"
   ]
  },
  {
   "id": "m11-l3-8",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def att_gt(data):\n",
    "    out = []\n",
    "    for g in sorted(data.g[data.g > 0].unique()):\n",
    "        for t in range(int(g), data.t.max() + 1):\n",
    "            base = g - 1                                              # last pre-period of cohort g\n",
    "            clean = data[(data.g == 0) | (data.g > t)]                # never or not-yet treated at t\n",
    "            if clean.unit.nunique() == 0:\n",
    "                continue\n",
    "            ch_g = data[(data.g == g) & (data.t == t)].y.mean() - data[(data.g == g) & (data.t == base)].y.mean()\n",
    "            ch_c = clean[clean.t == t].y.mean() - clean[clean.t == base].y.mean()\n",
    "            out.append((g, t, t - g, ch_g - ch_c, (data.g == g).sum() / data.t.nunique()))\n",
    "    return pd.DataFrame(out, columns=[\"g\", \"t\", \"e\", \"att\", \"n\"])\n",
    "\n",
    "gt = att_gt(sim)\n",
    "overall = np.average(gt.att, weights=gt.n)\n",
    "print(gt.round(2).head(8))\n",
    "print(f\"Aggregated ATT (clean comparisons): {overall:.2f} | truth over the same cells: \"\n",
    "      f\"{np.mean([(t - g + 1) for g, t in zip(gt.g, gt.t)]):.2f}\")"
   ]
  },
  {
   "id": "m11-l3-9",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "sim_nt = staggered(never_treated=30)\n",
    "two_stage = pf.did2s(sim_nt, yname=\"y\", first_stage=\"~ 0 | unit + t\", second_stage=\"~ D\",\n",
    "                     treatment=\"D\", cluster=\"unit\")\n",
    "twfe_nt = pf.feols(\"y ~ D | unit + t\", data=sim_nt, vcov={\"CRV1\": \"unit\"})\n",
    "print(f\"Truth {sim_nt.tau[sim_nt.D == 1].mean():.2f} | TWFE {twfe_nt.coef()['D']:.2f} | did2s {two_stage.coef()['D']:.2f}\")"
   ]
  },
  {
   "id": "m11-l3-10",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Real data: castle-doctrine laws and homicides"
   ]
  },
  {
   "id": "m11-l3-11",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** Did castle-doctrine laws change homicide rates?"
   ]
  },
  {
   "id": "m11-l3-12",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "from causaldata import castle\n",
    "c = castle.load_pandas().data[[\"sid\", \"year\", \"post\", \"l_homicide\"]].copy()\n",
    "c[\"year\"] = c.year.astype(\"int64\")\n",
    "c[\"D\"] = (c.post > 0).astype(int)\n",
    "c[\"g\"] = c.sid.map(c[c.D == 1].groupby(\"sid\").year.min()).fillna(0).astype(int)\n",
    "print(c.drop_duplicates(\"sid\").g.value_counts().sort_index().rename(\"states by adoption year\"))\n",
    "\n",
    "twfe_c = pf.feols(\"l_homicide ~ D | sid + year\", data=c, vcov={\"CRV1\": \"sid\"})\n",
    "did2s_c = pf.did2s(c, yname=\"l_homicide\", first_stage=\"~ 0 | sid + year\", second_stage=\"~ D\",\n",
    "                   treatment=\"D\", cluster=\"sid\")\n",
    "for name, m in [(\"TWFE\", twfe_c), (\"did2s\", did2s_c)]:\n",
    "    print(f\"{name:6}: {m.coef()['D']:.3f} (SE {m.se()['D']:.3f})  ->  about {100 * (np.exp(m.coef()['D']) - 1):.0f}% more homicides\")"
   ]
  },
  {
   "id": "m11-l3-13",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ✍️ Exercise\n",
    "\n",
    "Build a robust event study for the castle data: compute ATT(g, t) by hand with `att_gt` (using\n",
    "`l_homicide` as `y`, `year` as `t`), then average by event time `e`. Plot it. Are the effects stable after\n",
    "adoption?"
   ]
  },
  {
   "id": "m11-l3-14",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your answer here"
   ]
  },
  {
   "id": "m11-l3-15",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist and recap"
   ]
  }
 ]
}