{
 "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-l1-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 11.1 · Difference-in-Differences: the 2×2 Design and Parallel Trends\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- Compute a difference-in-differences estimate from four group means\n",
    "- Explain the parallel trends assumption and why before/after or treated/control comparisons alone fail\n",
    "- Estimate the same DiD with a regression interaction and clustered standard errors\n",
    "- Inspect pre-treatment trends graphically\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-l1-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Two bad comparisons make one good one"
   ]
  },
  {
   "id": "m11-l1-2",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q numpy pandas statsmodels matplotlib seaborn\n",
    "import numpy as np, pandas as pd\n",
    "import statsmodels.formula.api as smf\n",
    "import matplotlib.pyplot as plt, seaborn as sns\n",
    "rng = np.random.default_rng(42)\n",
    "sns.set_theme(style=\"whitegrid\")"
   ]
  },
  {
   "id": "m11-l1-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A provincial minimum wage increase"
   ]
  },
  {
   "id": "m11-l1-4",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** Did the minimum wage increase reduce youth employment in the treated province?"
   ]
  },
  {
   "id": "m11-l1-5",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "N_UNITS, N_TIME, T0 = 50, 10, 5\n",
    "muni = np.repeat(np.arange(N_UNITS), N_TIME)\n",
    "time = np.tile(np.arange(N_TIME), N_UNITS)\n",
    "treated = (muni >= 25).astype(int)\n",
    "post = (time >= T0).astype(int)\n",
    "\n",
    "muni_level = rng.normal(0.60, 0.05, N_UNITS)[muni] + 0.05 * treated     # treated province has higher employment\n",
    "common_shock = (-0.01 * time)                                          # a shared downturn\n",
    "employment = muni_level + common_shock - 0.05 * treated * post + rng.normal(0, 0.02, len(muni))\n",
    "df = pd.DataFrame({\"muni\": muni, \"time\": time, \"treated\": treated, \"post\": post, \"employment\": employment})\n",
    "\n",
    "means = df.groupby([\"treated\", \"post\"]).employment.mean().unstack()\n",
    "print(means.round(4))\n",
    "before_after = means.loc[1, 1] - means.loc[1, 0]\n",
    "treat_control = means.loc[1, 1] - means.loc[0, 1]\n",
    "did = (means.loc[1, 1] - means.loc[1, 0]) - (means.loc[0, 1] - means.loc[0, 0])\n",
    "print(f\"Before/after (treated only): {before_after:+.4f}\")\n",
    "print(f\"Treated vs control (post):   {treat_control:+.4f}\")\n",
    "print(f\"Difference-in-differences:   {did:+.4f}   (truth -0.05)\")"
   ]
  },
  {
   "id": "m11-l1-6",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DiD as a regression"
   ]
  },
  {
   "id": "m11-l1-7",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "fit = smf.ols(\"employment ~ treated * post\", data=df).fit(cov_type=\"cluster\", cov_kwds={\"groups\": df[\"muni\"]})\n",
    "print(fit.summary().tables[1])"
   ]
  },
  {
   "id": "m11-l1-8",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parallel trends: what we assume and what we can check"
   ]
  },
  {
   "id": "m11-l1-9",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "trend = df.groupby([\"time\", \"treated\"]).employment.mean().unstack()\n",
    "plt.figure(figsize=(10, 5))\n",
    "plt.plot(trend.index, trend[0], marker=\"o\", label=\"Control province\")\n",
    "plt.plot(trend.index, trend[1], marker=\"o\", label=\"Treated province\")\n",
    "counterfactual = trend[1].iloc[:T0].tolist() + (trend[1].iloc[T0 - 1] + (trend[0].iloc[T0:] - trend[0].iloc[T0 - 1])).tolist()\n",
    "plt.plot(trend.index, counterfactual, ls=\":\", color=\"grey\", label=\"Treated counterfactual (parallel trend)\")\n",
    "plt.axvline(T0 - 0.5, color=\"red\", ls=\"--\", label=\"Policy\")\n",
    "plt.xlabel(\"Month\"); plt.ylabel(\"Employment rate\"); plt.legend(); plt.title(\"Difference-in-differences\"); plt.show()"
   ]
  },
  {
   "id": "m11-l1-10",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise"
   ]
  },
  {
   "id": "m11-l1-11",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ✍️ Exercise\n",
    "\n",
    "Break parallel trends: give the treated province an extra upward trend of +0.005 per month from the start.\n",
    "Recompute the DiD. Then add a group-specific linear trend (`treated:time`) to the regression. Does that fix\n",
    "it here, and why is it risky in general?"
   ]
  },
  {
   "id": "m11-l1-12",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your answer here"
   ]
  },
  {
   "id": "m11-l1-13",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist and recap"
   ]
  }
 ]
}