{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "cells": [
  {
   "id": "m03-l2-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 3.2 · Omitted Variable Bias and Partialling Out\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- Simulate an unobserved confounder and measure the omitted variable bias (OVB) it creates\n",
    "- Verify the OVB formula numerically\n",
    "- Use multiple regression to \"close the backdoor\" and recover the causal return to education\n",
    "- Understand partialling out (Frisch–Waugh–Lovell), the engine of Double Machine Learning\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": "m03-l2-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "id": "m03-l2-2",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q numpy polars statsmodels matplotlib seaborn causaldata\n",
    "import numpy as np, polars as pl, 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(2026)\n",
    "sns.set_theme(style=\"darkgrid\")"
   ]
  },
  {
   "id": "m03-l2-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A world with ability bias"
   ]
  },
  {
   "id": "m03-l2-4",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** Why does a simple regression overstate the return to education?"
   ]
  },
  {
   "id": "m03-l2-5",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "N = 1_000\n",
    "TRUE_RETURN = 2_500\n",
    "ability = rng.normal(0, 1, N)\n",
    "education = np.clip(np.round(12 + 1.5 * ability + rng.normal(0, 1.2, N)), 8, 20)\n",
    "wage = 30_000 + TRUE_RETURN * education + 8_000 * ability + rng.normal(0, 5_000, N)\n",
    "df = pd.DataFrame({\"education\": education, \"ability\": ability, \"wage\": wage})\n",
    "\n",
    "naive = smf.ols(\"wage ~ education\", data=df).fit()\n",
    "controlled = smf.ols(\"wage ~ education + ability\", data=df).fit()\n",
    "print(f\"True return:          ${TRUE_RETURN:,.0f}\")\n",
    "print(f\"Naive estimate:       ${naive.params['education']:,.0f}\")\n",
    "print(f\"Controlling ability:  ${controlled.params['education']:,.0f}\")"
   ]
  },
  {
   "id": "m03-l2-6",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The OVB formula, checked by code"
   ]
  },
  {
   "id": "m03-l2-7",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "gamma = controlled.params[\"ability\"]                                  # effect of the omitted variable\n",
    "delta = smf.ols(\"ability ~ education\", data=df).fit().params[\"education\"]  # its link to the treatment\n",
    "print(f\"Controlled + gamma*delta = {controlled.params['education'] + gamma * delta:,.1f}\")\n",
    "print(f\"Naive slope              = {naive.params['education']:,.1f}\")"
   ]
  },
  {
   "id": "m03-l2-8",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Partialling out: the Frisch–Waugh–Lovell theorem"
   ]
  },
  {
   "id": "m03-l2-9",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** What does \"controlling for ability\" actually do to the data?"
   ]
  },
  {
   "id": "m03-l2-10",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "df[\"wage_res\"] = smf.ols(\"wage ~ ability\", data=df).fit().resid\n",
    "df[\"educ_res\"] = smf.ols(\"education ~ ability\", data=df).fit().resid\n",
    "fwl = smf.ols(\"wage_res ~ educ_res - 1\", data=df).fit()\n",
    "print(f\"FWL slope:        {fwl.params['educ_res']:,.2f}\")\n",
    "print(f\"Multiple OLS:     {controlled.params['education']:,.2f}\")\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(12, 4))\n",
    "sns.regplot(x=\"education\", y=\"wage\", data=df, ax=axes[0], scatter_kws={\"alpha\": .3}).set(title=\"Raw: confounded\")\n",
    "sns.regplot(x=\"educ_res\", y=\"wage_res\", data=df, ax=axes[1], scatter_kws={\"alpha\": .3}).set(title=\"Residualised: ability removed\")\n",
    "plt.tight_layout(); plt.show()"
   ]
  },
  {
   "id": "m03-l2-11",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Real data: Card's returns to schooling"
   ]
  },
  {
   "id": "m03-l2-12",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "from causaldata import close_college\n",
    "card = close_college.load_pandas().data\n",
    "m1 = smf.ols(\"lwage ~ educ\", data=card).fit(cov_type=\"HC1\")\n",
    "m2 = smf.ols(\"lwage ~ educ + exper + I(exper**2) + black + south + smsa + married\", data=card).fit(cov_type=\"HC1\")\n",
    "print(f\"Return to schooling, no controls:   {100*m1.params['educ']:.1f}% per year\")\n",
    "print(f\"Return to schooling, with controls: {100*m2.params['educ']:.1f}% per year\")"
   ]
  },
  {
   "id": "m03-l2-13",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ✍️ Exercise\n",
    "\n",
    "In the simulated world, suppose ability is measured with error: you only observe\n",
    "`ability_noisy = ability + noise` with noise standard deviation 1. Control for the noisy measure instead.\n",
    "How much bias remains? What does this suggest about \"controlling for\" imperfect proxies such as test scores?"
   ]
  },
  {
   "id": "m03-l2-14",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your answer here"
   ]
  },
  {
   "id": "m03-l2-15",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist and recap"
   ]
  }
 ]
}