{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "cells": [
  {
   "id": "m06-l3-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 6.3 · Inverse Probability Weighting, Trimming and Diagnostics\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- Explain how inverse probability weights create a pseudo-population in which treatment is random\n",
    "- Estimate the ATE and the ATT with IPW, using normalised (Hájek) weights\n",
    "- Diagnose extreme weights and stabilise them by trimming\n",
    "- Apply IPW to a targeted welfare programme and to the LaLonde data\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": "m06-l3-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The idea of reweighting"
   ]
  },
  {
   "id": "m06-l3-2",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q numpy pandas statsmodels matplotlib seaborn causaldata\n",
    "import numpy as np, 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(42)\n",
    "sns.set_theme(style=\"whitegrid\")"
   ]
  },
  {
   "id": "m06-l3-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A targeted tutoring programme"
   ]
  },
  {
   "id": "m06-l3-4",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** How many points does targeted tutoring add, once we account for whom it targets?"
   ]
  },
  {
   "id": "m06-l3-5",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "N = 2_000\n",
    "poverty = rng.normal(0, 1, N)\n",
    "p_true = 1 / (1 + np.exp(-(-0.5 + 1.2 * poverty)))\n",
    "treat = rng.binomial(1, p_true)\n",
    "score = 50 + 10 * treat - 15 * poverty + rng.normal(0, 5, N)\n",
    "df = pd.DataFrame({\"poverty\": poverty, \"treat\": treat, \"score\": score})\n",
    "\n",
    "df[\"ps\"] = smf.logit(\"treat ~ poverty\", data=df).fit(disp=0).predict(df)\n",
    "w_ate = np.where(df.treat == 1, 1 / df.ps, 1 / (1 - df.ps))\n",
    "\n",
    "def weighted_diff(y, d, w):\n",
    "    \"\"\"Hájek (normalised) estimator: weighted mean of treated minus weighted mean of controls.\"\"\"\n",
    "    return np.average(y[d == 1], weights=w[d == 1]) - np.average(y[d == 0], weights=w[d == 0])\n",
    "\n",
    "naive = df.score[df.treat == 1].mean() - df.score[df.treat == 0].mean()\n",
    "ipw = weighted_diff(df.score.to_numpy(), df.treat.to_numpy(), w_ate)\n",
    "wls = smf.wls(\"score ~ treat\", data=df, weights=w_ate).fit(cov_type=\"HC1\")\n",
    "print(f\"Naive: {naive:6.2f} | IPW (Hájek): {ipw:6.2f} | WLS: {wls.params['treat']:6.2f} (robust SE {wls.bse['treat']:.2f}) | truth 10\")"
   ]
  },
  {
   "id": "m06-l3-6",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "for g in (1, 0):\n",
    "    m = df.treat == g\n",
    "    print(f\"treat={g}: raw mean poverty {df.poverty[m].mean():+.3f} | weighted {np.average(df.poverty[m], weights=w_ate[m]):+.3f}\")"
   ]
  },
  {
   "id": "m06-l3-7",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extreme weights and trimming"
   ]
  },
  {
   "id": "m06-l3-8",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** Is my IPW estimate driven by a handful of observations with huge weights?"
   ]
  },
  {
   "id": "m06-l3-9",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def ess(w):\n",
    "    return w.sum() ** 2 / (w ** 2).sum()\n",
    "\n",
    "# A harsher DGP: strong targeting -> poor overlap\n",
    "p_hard = 1 / (1 + np.exp(-(-0.5 + 3.0 * poverty)))\n",
    "treat_h = rng.binomial(1, p_hard)\n",
    "score_h = 50 + 10 * treat_h - 15 * poverty + rng.normal(0, 5, N)\n",
    "h = pd.DataFrame({\"poverty\": poverty, \"treat\": treat_h, \"score\": score_h})\n",
    "h[\"ps\"] = smf.logit(\"treat ~ poverty\", data=h).fit(disp=0).predict(h)\n",
    "w_h = np.where(h.treat == 1, 1 / h.ps, 1 / (1 - h.ps))\n",
    "print(f\"Max weight {w_h.max():,.0f} | ESS {ess(w_h):,.0f} of {N:,}\")\n",
    "\n",
    "keep = (h.ps > 0.05) & (h.ps < 0.95)\n",
    "est_all = weighted_diff(h.score.to_numpy(), h.treat.to_numpy(), w_h)\n",
    "est_trim = weighted_diff(h.score[keep].to_numpy(), h.treat[keep].to_numpy(), w_h[keep])\n",
    "print(f\"IPW all: {est_all:.2f} | IPW trimmed to [0.05, 0.95]: {est_trim:.2f} (kept {keep.sum():,})\")\n",
    "\n",
    "sns.histplot(w_h, bins=80, log_scale=True); plt.title(\"IPW weights (log scale)\"); plt.show()"
   ]
  },
  {
   "id": "m06-l3-10",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ATT weights on the LaLonde data"
   ]
  },
  {
   "id": "m06-l3-11",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "from causaldata import nsw_mixtape, cps_mixtape\n",
    "nsw = nsw_mixtape.load_pandas().data\n",
    "obs = pd.concat([nsw[nsw.treat == 1], cps_mixtape.load_pandas().data], ignore_index=True)\n",
    "obs[\"u74\"] = (obs.re74 == 0).astype(int); obs[\"u75\"] = (obs.re75 == 0).astype(int)\n",
    "obs[\"ps\"] = smf.logit(\"treat ~ age + I(age**2) + I(age**3) + educ + I(educ**2) + marr + nodegree\"\n",
    "                      \" + black + hisp + re74 + re75 + u74 + u75 + educ:re74\", data=obs).fit(disp=0).predict(obs)\n",
    "\n",
    "w_att = np.where(obs.treat == 1, 1.0, obs.ps / (1 - obs.ps))\n",
    "att = weighted_diff(obs.re78.to_numpy(), obs.treat.to_numpy(), w_att)\n",
    "print(f\"IPW ATT: ${att:,.0f}  (experimental $1,794)\")\n",
    "print(f\"Control ESS: {ess(w_att[obs.treat == 0]):,.0f} of {(obs.treat == 0).sum():,} CPS respondents\")"
   ]
  },
  {
   "id": "m06-l3-12",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise"
   ]
  },
  {
   "id": "m06-l3-13",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ✍️ Exercise\n",
    "\n",
    "Write a function `ipw_bootstrap(data, reps=200)` that re-estimates the propensity score and the IPW ATT in\n",
    "each bootstrap sample of the LaLonde data and returns a 95% percentile interval. Does it contain $1,794?"
   ]
  },
  {
   "id": "m06-l3-14",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your answer here"
   ]
  },
  {
   "id": "m06-l3-15",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist and recap"
   ]
  }
 ]
}