{
 "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-l2-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 11.2 · Two-Way Fixed Effects and Event Studies with pyfixest\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- Estimate a two-way fixed effects (TWFE) DiD with pyfixest and clustered standard errors\n",
    "- Build an event study with leads and lags relative to treatment\n",
    "- Read pre-treatment coefficients as a check on parallel trends\n",
    "- Plot dynamic effects with confidence intervals\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-l2-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Unit and time fixed effects"
   ]
  },
  {
   "id": "m11-l2-2",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q pyfixest numpy pandas matplotlib\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(42)"
   ]
  },
  {
   "id": "m11-l2-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TWFE with pyfixest"
   ]
  },
  {
   "id": "m11-l2-4",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** What is the average effect of the minimum wage increase on youth employment, controlling for municipality and month?"
   ]
  },
  {
   "id": "m11-l2-5",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "U, T, T0 = 60, 24, 12\n",
    "muni = np.repeat(np.arange(U), T); month = np.tile(np.arange(T), U)\n",
    "treated = (muni >= 30).astype(int)\n",
    "rel = np.where(treated == 1, month - T0, -99)                     # months since the policy (treated only)\n",
    "effect = np.where((treated == 1) & (rel >= 0), -0.02 - 0.03 * np.minimum(rel, 6) / 6, 0.0)\n",
    "y = (rng.normal(0.6, 0.05, U)[muni] + 0.03 * np.sin(month / 3) + effect + rng.normal(0, 0.015, len(muni)))\n",
    "df = pd.DataFrame({\"muni\": muni, \"month\": month, \"treated\": treated, \"rel\": rel,\n",
    "                   \"D\": ((treated == 1) & (month >= T0)).astype(int), \"employment\": y})\n",
    "\n",
    "twfe = pf.feols(\"employment ~ D | muni + month\", data=df, vcov={\"CRV1\": \"muni\"})\n",
    "twfe.summary()\n",
    "print(f\"True average effect in the post period: {effect[df.D == 1].mean():.4f}\")"
   ]
  },
  {
   "id": "m11-l2-6",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Event studies: effects by time since treatment"
   ]
  },
  {
   "id": "m11-l2-7",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** How did employment evolve month by month before and after the minimum wage increase?"
   ]
  },
  {
   "id": "m11-l2-8",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "df[\"rel_binned\"] = df.rel.clip(-8, 10)                            # bin distant periods\n",
    "df.loc[df.treated == 0, \"rel_binned\"] = -1                         # never-treated sit in the reference period\n",
    "es = pf.feols(\"employment ~ i(rel_binned, ref=-1) | muni + month\", data=df, vcov={\"CRV1\": \"muni\"})\n",
    "coefs = es.tidy()\n",
    "print(coefs[[\"Estimate\", \"Std. Error\"]].round(4).head(10))"
   ]
  },
  {
   "id": "m11-l2-9",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reading and plotting the event study"
   ]
  },
  {
   "id": "m11-l2-10",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "tab = es.tidy().reset_index()\n",
    "tab[\"k\"] = tab[\"Coefficient\"].str.extract(r\"\\[(?:T\\.)?(-?\\d+(?:\\.\\d+)?)\\]\")[0].astype(float)\n",
    "tab = pd.concat([tab, pd.DataFrame({\"k\": [-1.0], \"Estimate\": [0.0], \"Std. Error\": [0.0]})]).sort_values(\"k\")\n",
    "\n",
    "true_path = {k: (-0.02 - 0.03 * min(k, 6) / 6 if k >= 0 else 0.0) for k in tab.k}\n",
    "plt.figure(figsize=(10, 5))\n",
    "plt.errorbar(tab.k, tab.Estimate, yerr=1.96 * tab[\"Std. Error\"], fmt=\"o\", capsize=3, label=\"Estimate ± 95% CI\")\n",
    "plt.plot(list(true_path), list(true_path.values()), \"r--\", label=\"True effect\")\n",
    "plt.axhline(0, color=\"grey\"); plt.axvline(-0.5, color=\"black\", ls=\":\")\n",
    "plt.xlabel(\"Months relative to the policy\"); plt.ylabel(\"Effect on employment rate\")\n",
    "plt.title(\"Event study: leads are flat, lags show the effect building up\"); plt.legend(); plt.show()\n",
    "\n",
    "leads = [c for c in es.coef().index if \"-\" in c and \"-1]\" not in c]\n",
    "print(\"Joint test that all leads are zero:\")\n",
    "print(es.wald_test(R=np.eye(len(es.coef()))[[list(es.coef().index).index(c) for c in leads]]))"
   ]
  },
  {
   "id": "m11-l2-11",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ✍️ Exercise\n",
    "\n",
    "Add an anticipation effect: employment in the treated province starts falling two months **before** the\n",
    "policy (firms adjust in advance). Re-run the event study. What do the leads show, and how should you\n",
    "redefine the reference period?"
   ]
  },
  {
   "id": "m11-l2-12",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your answer here"
   ]
  },
  {
   "id": "m11-l2-13",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist and recap"
   ]
  }
 ]
}