{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "cells": [
  {
   "id": "m09-l2-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 9.2 · Two-Stage Least Squares and Weak Instruments\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- Estimate 2SLS by hand and with linearmodels, and see why manual second stages get the standard errors wrong\n",
    "- Estimate the return to schooling with college proximity as an instrument (Card, 1995)\n",
    "- Read first-stage diagnostics and apply the F > 10 rule of thumb (and its modern caveats)\n",
    "- Show by simulation how weak instruments bias 2SLS toward OLS\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": "m09-l2-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2SLS in two regressions"
   ]
  },
  {
   "id": "m09-l2-2",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q numpy pandas statsmodels linearmodels matplotlib causaldata\n",
    "import numpy as np, pandas as pd\n",
    "import statsmodels.api as sm\n",
    "from linearmodels.iv import IV2SLS\n",
    "import matplotlib.pyplot as plt\n",
    "rng = np.random.default_rng(101)"
   ]
  },
  {
   "id": "m09-l2-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The ability-bias simulation"
   ]
  },
  {
   "id": "m09-l2-4",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** What is the earnings return to a year of education when ability is unobserved?"
   ]
  },
  {
   "id": "m09-l2-5",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "N = 5_000\n",
    "ability = rng.normal(0, 1, N)\n",
    "near_college = rng.binomial(1, 0.5, N)\n",
    "education = 12 + 1.5 * near_college + 2.0 * ability + rng.normal(0, 1, N)\n",
    "earnings = 25_000 + 3_000 * education + 5_000 * ability + rng.normal(0, 4_000, N)\n",
    "df = pd.DataFrame({\"earnings\": earnings, \"education\": education, \"near_college\": near_college, \"const\": 1.0})\n",
    "\n",
    "ols = sm.OLS(df.earnings, df[[\"const\", \"education\"]]).fit()\n",
    "\n",
    "# 2SLS by hand (point estimate only!)\n",
    "d_hat = sm.OLS(df.education, df[[\"const\", \"near_college\"]]).fit().fittedvalues\n",
    "manual = sm.OLS(df.earnings, sm.add_constant(d_hat.rename(\"education_hat\"))).fit()\n",
    "\n",
    "iv = IV2SLS(dependent=df.earnings, exog=df.const, endog=df.education, instruments=df.near_college).fit(cov_type=\"robust\")\n",
    "print(f\"OLS:          {ols.params['education']:,.0f}\")\n",
    "print(f\"Manual 2SLS:  {manual.params['education_hat']:,.0f}  (SE {manual.bse['education_hat']:,.0f}: WRONG)\")\n",
    "print(f\"IV2SLS:       {iv.params['education']:,.0f}  (SE {iv.std_errors['education']:,.0f}: correct)\")"
   ]
  },
  {
   "id": "m09-l2-6",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Real data: distance to college"
   ]
  },
  {
   "id": "m09-l2-7",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** Does growing up near a college reveal a higher return to schooling than OLS suggests?"
   ]
  },
  {
   "id": "m09-l2-8",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "from causaldata import close_college\n",
    "card = close_college.load_pandas().data.dropna()\n",
    "controls = \"1 + exper + I(exper**2) + black + south + smsa + married\"\n",
    "ols_card = IV2SLS.from_formula(f\"lwage ~ {controls} + educ\", card).fit(cov_type=\"robust\")\n",
    "iv_card = IV2SLS.from_formula(f\"lwage ~ {controls} + [educ ~ nearc4]\", card).fit(cov_type=\"robust\")\n",
    "print(f\"OLS return:  {100 * ols_card.params['educ']:.1f}% per year (SE {100 * ols_card.std_errors['educ']:.1f})\")\n",
    "print(f\"2SLS return: {100 * iv_card.params['educ']:.1f}% per year (SE {100 * iv_card.std_errors['educ']:.1f})\")\n",
    "print(iv_card.first_stage.diagnostics.round(4))"
   ]
  },
  {
   "id": "m09-l2-9",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Weak instruments"
   ]
  },
  {
   "id": "m09-l2-10",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** How does the 2SLS estimate behave as the instrument gets weaker?"
   ]
  },
  {
   "id": "m09-l2-11",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def simulate_iv(strength, n=1_000, reps=300):\n",
    "    est, fstats = [], []\n",
    "    for _ in range(reps):\n",
    "        a = rng.normal(size=n); z = rng.binomial(1, 0.5, n)\n",
    "        educ = 12 + strength * z + 2 * a + rng.normal(size=n)\n",
    "        y = 25_000 + 3_000 * educ + 5_000 * a + rng.normal(0, 4_000, n)\n",
    "        d = pd.DataFrame({\"y\": y, \"educ\": educ, \"z\": z, \"const\": 1.0})\n",
    "        r = IV2SLS(d.y, d.const, d.educ, d.z).fit(cov_type=\"robust\")\n",
    "        est.append(r.params[\"educ\"]); fstats.append(r.first_stage.diagnostics[\"f.stat\"].iloc[0])\n",
    "    return np.median(est), np.median(fstats)\n",
    "\n",
    "rows = [(s, *simulate_iv(s)) for s in [1.5, 0.5, 0.25, 0.1, 0.05]]\n",
    "print(pd.DataFrame(rows, columns=[\"first-stage coefficient\", \"median 2SLS\", \"median F\"]).round(2))"
   ]
  },
  {
   "id": "m09-l2-12",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ✍️ Exercise\n",
    "\n",
    "For the Card data, compute the **reduced form** (lwage on nearc4 and controls) and the **first stage**\n",
    "(educ on nearc4 and controls) with statsmodels. Check that their ratio equals the 2SLS coefficient."
   ]
  },
  {
   "id": "m09-l2-13",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your answer here"
   ]
  },
  {
   "id": "m09-l2-14",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist and recap"
   ]
  }
 ]
}