{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "cells": [
  {
   "id": "m07-l1-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 7.1 · Doubly Robust Estimation: AIPW\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- Estimate an ATE by regression adjustment (the outcome-model approach)\n",
    "- Combine the outcome model and the propensity score into the augmented IPW (AIPW) estimator\n",
    "- Show by simulation that AIPW is consistent if either model is right (\"double robustness\")\n",
    "- Compute an AIPW standard error from its influence function\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": "m07-l1-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Two routes to the same counterfactual"
   ]
  },
  {
   "id": "m07-l1-2",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q numpy pandas scikit-learn statsmodels matplotlib seaborn\n",
    "import numpy as np, pandas as pd\n",
    "from sklearn.linear_model import LinearRegression, LogisticRegression\n",
    "import matplotlib.pyplot as plt, seaborn as sns\n",
    "rng = np.random.default_rng(3)\n",
    "sns.set_theme(style=\"whitegrid\")"
   ]
  },
  {
   "id": "m07-l1-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A Medicaid-style expansion with non-linear confounding"
   ]
  },
  {
   "id": "m07-l1-4",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** What is the effect of insurance coverage on a health index when enrolment depends non-linearly on income?"
   ]
  },
  {
   "id": "m07-l1-5",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def simulate(n=4_000):\n",
    "    income = rng.normal(0, 1, n)\n",
    "    chronic = rng.binomial(1, 0.3, n)\n",
    "    logit = -0.3 - 1.2 * income + 0.8 * chronic + 0.6 * income ** 2\n",
    "    d = rng.binomial(1, 1 / (1 + np.exp(-logit)))\n",
    "    y = 2.0 * d + 3 * income - 2 * chronic + 1.5 * income ** 2 + rng.normal(0, 2, n)\n",
    "    X = pd.DataFrame({\"income\": income, \"chronic\": chronic})\n",
    "    return X, d, y\n",
    "\n",
    "X, d, y = simulate()\n",
    "TRUE_ATE = 2.0\n",
    "print(f\"Naive difference: {y[d == 1].mean() - y[d == 0].mean():.2f} (truth {TRUE_ATE})\")"
   ]
  },
  {
   "id": "m07-l1-6",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Double robustness by simulation"
   ]
  },
  {
   "id": "m07-l1-7",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** Which estimator survives if I get one of the two models wrong?"
   ]
  },
  {
   "id": "m07-l1-8",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def features(X, correct):\n",
    "    F = X.copy()\n",
    "    if correct:\n",
    "        F[\"income2\"] = F.income ** 2\n",
    "    return F.to_numpy()\n",
    "\n",
    "def estimators(X, d, y, outcome_ok, ps_ok):\n",
    "    Fo, Fp = features(X, outcome_ok), features(X, ps_ok)\n",
    "    mu1 = LinearRegression().fit(Fo[d == 1], y[d == 1]).predict(Fo)\n",
    "    mu0 = LinearRegression().fit(Fo[d == 0], y[d == 0]).predict(Fo)\n",
    "    e = LogisticRegression(C=1e6, max_iter=1_000).fit(Fp, d).predict_proba(Fp)[:, 1]\n",
    "    e = np.clip(e, 0.01, 0.99)\n",
    "    ra = np.mean(mu1 - mu0)\n",
    "    ipw = np.mean(d * y / e) - np.mean((1 - d) * y / (1 - e))\n",
    "    psi = mu1 - mu0 + d * (y - mu1) / e - (1 - d) * (y - mu0) / (1 - e)\n",
    "    return ra, ipw, psi.mean(), psi\n",
    "\n",
    "results = []\n",
    "for outcome_ok in (True, False):\n",
    "    for ps_ok in (True, False):\n",
    "        reps = [estimators(*simulate(), outcome_ok, ps_ok)[:3] for _ in range(100)]\n",
    "        ra, ipw, aipw = np.mean(reps, axis=0)\n",
    "        results.append((outcome_ok, ps_ok, ra, ipw, aipw))\n",
    "table = pd.DataFrame(results, columns=[\"outcome model OK\", \"PS model OK\", \"Regression adj.\", \"IPW\", \"AIPW\"]).round(2)\n",
    "print(table, f\"\\n(truth = {TRUE_ATE})\")"
   ]
  },
  {
   "id": "m07-l1-9",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Standard errors from the influence function"
   ]
  },
  {
   "id": "m07-l1-10",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "ate, psi = estimators(X, d, y, True, True)[2:]\n",
    "se = psi.std(ddof=1) / np.sqrt(len(psi))\n",
    "print(f\"AIPW ATE = {ate:.3f}, SE = {se:.3f}, 95% CI [{ate - 1.96 * se:.3f}, {ate + 1.96 * se:.3f}]\")"
   ]
  },
  {
   "id": "m07-l1-11",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ✍️ Exercise\n",
    "\n",
    "Repeat the simulation table with n = 500 instead of 4,000 and look at the **spread** (standard deviation)\n",
    "of the 100 estimates for each estimator, not just the mean. Which is most variable, and why?"
   ]
  },
  {
   "id": "m07-l1-12",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your answer here"
   ]
  },
  {
   "id": "m07-l1-13",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist and recap"
   ]
  }
 ]
}