{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "cells": [
  {
   "id": "m14-l3-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 14.3 · Omitted-Variable Sensitivity: Robustness Values and Oster's δ\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- Express sensitivity in terms of partial R² of an unobserved confounder (Cinelli and Hazlett)\n",
    "- Compute the robustness value of a regression estimate and benchmark it against observed covariates\n",
    "- Use coefficient stability (Oster's δ) to judge how much selection on unobservables would erase an effect\n",
    "- Write the sensitivity section of an executive memo\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": "m14-l3-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## From OVB formula to sensitivity analysis"
   ]
  },
  {
   "id": "m14-l3-2",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q numpy pandas statsmodels matplotlib causaldata\n",
    "import numpy as np, pandas as pd\n",
    "import statsmodels.formula.api as smf\n",
    "import matplotlib.pyplot as plt\n",
    "from causaldata import nhefs_complete"
   ]
  },
  {
   "id": "m14-l3-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Robustness values on the NHEFS smoking data"
   ]
  },
  {
   "id": "m14-l3-4",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** How strong would an unmeasured confounder have to be to explain away the weight gain from quitting smoking?"
   ]
  },
  {
   "id": "m14-l3-5",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "nh = nhefs_complete.load_pandas().data\n",
    "covs = \"sex + race + age + I(age**2) + C(education) + smokeintensity + I(smokeintensity**2) + smokeyrs + I(smokeyrs**2) + C(exercise) + C(active) + wt71 + I(wt71**2)\"\n",
    "fit = smf.ols(f\"wt82_71 ~ qsmk + {covs}\", data=nh).fit()\n",
    "est, se, t, dof = fit.params[\"qsmk\"], fit.bse[\"qsmk\"], fit.tvalues[\"qsmk\"], fit.df_resid\n",
    "\n",
    "def robustness_value(t, dof, q=1.0, alpha=None):\n",
    "    f = q * abs(t) / np.sqrt(dof)\n",
    "    if alpha is not None:                                  # RV for losing significance at level alpha\n",
    "        from scipy.stats import t as tdist\n",
    "        f = max(q * abs(t) / np.sqrt(dof) - tdist.ppf(1 - alpha / 2, dof - 1) / np.sqrt(dof - 1), 0)\n",
    "    return 0.5 * (np.sqrt(f ** 4 + 4 * f ** 2) - f ** 2)\n",
    "\n",
    "partial_r2_d = t ** 2 / (t ** 2 + dof)                     # share of outcome variance explained by D itself\n",
    "print(f\"Estimate: {est:.2f} kg (SE {se:.2f}, t = {t:.1f})\")\n",
    "print(f\"Partial R² of quitting with weight change: {partial_r2_d:.3f}\")\n",
    "print(f\"Robustness value (to bring the estimate to 0):         {robustness_value(t, dof):.3f}\")\n",
    "print(f\"Robustness value (to lose significance at 5%):          {robustness_value(t, dof, alpha=0.05):.3f}\")"
   ]
  },
  {
   "id": "m14-l3-6",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Benchmarking against observed covariates"
   ]
  },
  {
   "id": "m14-l3-7",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def partial_r2(full, var_terms):\n",
    "    \"\"\"Partial R² of a group of terms in a fitted model, via the change in R².\"\"\"\n",
    "    reduced_formula = full.model.formula\n",
    "    for term in var_terms:\n",
    "        reduced_formula = reduced_formula.replace(f\" + {term}\", \"\")\n",
    "    reduced = smf.ols(reduced_formula, data=nh).fit()\n",
    "    return (full.rsquared - reduced.rsquared) / (1 - reduced.rsquared)\n",
    "\n",
    "d_model = smf.ols(f\"qsmk ~ {covs}\", data=nh).fit()\n",
    "benchmarks = {\n",
    "    \"age\": [\"age\", \"I(age**2)\"],\n",
    "    \"smoking intensity\": [\"smokeintensity\", \"I(smokeintensity**2)\"],\n",
    "    \"baseline weight\": [\"wt71\", \"I(wt71**2)\"],\n",
    "}\n",
    "rows = [(name, partial_r2(d_model, terms), partial_r2(fit, terms)) for name, terms in benchmarks.items()]\n",
    "bench = pd.DataFrame(rows, columns=[\"covariate\", \"R² with quitting | X\", \"R² with weight change | D, X\"]).round(3)\n",
    "print(bench)\n",
    "print(f\"Robustness value: {robustness_value(t, dof):.3f}\")"
   ]
  },
  {
   "id": "m14-l3-8",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Coefficient stability: Oster's δ"
   ]
  },
  {
   "id": "m14-l3-9",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "short = smf.ols(\"wt82_71 ~ qsmk\", data=nh).fit()\n",
    "b0, r0 = short.params[\"qsmk\"], short.rsquared\n",
    "b1, r1 = fit.params[\"qsmk\"], fit.rsquared\n",
    "r_max = min(1.0, 1.3 * r1)\n",
    "delta = (b1 * (r1 - r0)) / ((b0 - b1) * (r_max - r1))\n",
    "print(f\"Uncontrolled: {b0:.2f} kg (R² {r0:.3f}) | Controlled: {b1:.2f} kg (R² {r1:.3f}) | R_max {r_max:.3f}\")\n",
    "print(f\"Oster's δ for a zero effect: {delta:.2f}\")"
   ]
  },
  {
   "id": "m14-l3-10",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ✍️ Exercise\n",
    "\n",
    "Run the robustness-value calculation for the LaLonde OLS regression (re78 on treat and covariates, NSW\n",
    "treated vs CPS controls, from Lesson 6.1). Compare its RV with the NHEFS one. Which finding is more robust,\n",
    "and is that consistent with what you learned about the LaLonde data?"
   ]
  },
  {
   "id": "m14-l3-11",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your answer here"
   ]
  },
  {
   "id": "m14-l3-12",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist and recap"
   ]
  }
 ]
}