{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Day 44: IPW & Hájek Estimators of the ATE\n",
    "\n",
    "*Beginner to Ninja: Math for Causal Inference*\n",
    "\n",
    "Re-weighting treated and control units by the inverse of their propensity scores to recreate the population each group is missing.\n",
    "\n",
    "**Learning outcomes**\n",
    "\n",
    "- Prove the inverse-probability-weighting identities for E[Y(1)], E[Y(0)] and the ATT.\n",
    "- Compute Horvitz–Thompson and Hájek IPW estimates of the ATE and ATT by hand.\n",
    "- Diagnose extreme weights with the effective sample size and understand trimming and stabilisation.\n",
    "- Relate IPW to weighted regression and to survey weighting."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Theory\n",
    "\n",
    "### The idea\n",
    "\n",
    "Under unconfoundedness and overlap (Day 42), a treated unit with propensity $e(x)$ represents $1/e(x)$ units with the same covariates, some of whom were not treated. Weighting each treated unit by $1/e(X)$ rebuilds the whole population's $Y(1)$ distribution; weighting controls by $1/(1 - e(X))$ rebuilds $Y(0)$:\n",
    "\n",
    "$$\n",
    "E\\Big[\\frac{TY}{e(X)}\\Big] = E[Y(1)], \\qquad E\\Big[\\frac{(1 - T)Y}{1 - e(X)}\\Big] = E[Y(0)].\n",
    "$$\n",
    "\n",
    "(This is Day 14's missing-data result applied twice.)\n",
    "\n",
    "### Estimators\n",
    "\n",
    "**Horvitz–Thompson (HT):**\n",
    "\n",
    "$$\n",
    "\\hat\\tau_{HT} = \\frac1n\\sum_i\\frac{T_iY_i}{\\hat e(X_i)} - \\frac1n\\sum_i\\frac{(1 - T_i)Y_i}{1 - \\hat e(X_i)}.\n",
    "$$\n",
    "\n",
    "**Hájek (normalised):** divide each weighted sum by the sum of its weights instead of $n$:\n",
    "\n",
    "$$\n",
    "\\hat\\tau_{H} = \\frac{\\sum_iT_iY_i/\\hat e_i}{\\sum_iT_i/\\hat e_i} - \\frac{\\sum_i(1 - T_i)Y_i/(1 - \\hat e_i)}{\\sum_i(1 - T_i)/(1 - \\hat e_i)}.\n",
    "$$\n",
    "\n",
    "Hájek is invariant to shifting $Y$ by a constant, keeps each mean within the range of the data, and usually has lower variance. It equals a weighted regression of $Y$ on $T$ (Day 24).\n",
    "\n",
    "**ATT:** treated units get weight 1; controls get $\\hat e/(1 - \\hat e)$ (the odds of treatment), which reshapes the controls to look like the treated.\n",
    "\n",
    "### Weights and precision\n",
    "\n",
    "When $\\hat e$ is near 0 or 1, a few units get huge weights and dominate the estimate. Diagnostics and remedies:\n",
    "\n",
    "- **Effective sample size:** $\\text{ESS} = (\\sum w_i)^2/\\sum w_i^2$, the number of equally weighted units with the same precision.\n",
    "- **Trimming** units with extreme $\\hat e$ (changes the estimand, Day 42) or **truncating** weights at a cap (adds a little bias).\n",
    "- Using the *estimated* propensity score is, perhaps surprisingly, more efficient than using the true one (Hirano, Imbens & Ridder 2003).\n",
    "\n",
    "## Part 1 · Warm-up"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 1 · Individual weights\n",
    "\n",
    "A treated household has $e = 0.25$ and a control household also has $e = 0.25$. What ATE weights do they get, and what do the weights mean?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Treated: $1/0.25 = 4$. It stands in for itself and three similar households that were not treated.\n",
    "\n",
    "**Step 2.** Control: $1/(1 - 0.25) \\approx 1.33$. Controls with this profile are common, so each represents only a third more households.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 2 · An ATT weight\n",
    "\n",
    "What ATT weight does the control household in Exercise 1 get?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Odds of treatment: $e/(1 - e) = 0.25/0.75 = 1/3$.\n",
    "\n",
    "**Answer:** 1/3. Profiles that are rarely treated get small weight when we want controls to resemble the treated.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 3 · Effective sample size\n",
    "\n",
    "Four units have weights 1, 1, 1 and 10. Compute the ESS.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** $\\sum w = 13$ and $\\sum w^2 = 1 + 1 + 1 + 100 = 103$.\n",
    "\n",
    "**Step 2.** ESS $= 13^2/103 \\approx 1.64$.\n",
    "\n",
    "**Answer:** four units carry the information of fewer than two: one unit dominates.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 4 · Hájek is location-invariant\n",
    "\n",
    "Add 100 to every outcome. How do the HT and Hájek estimates of $E[Y(1)]$ change?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Hájek: weighted average of $Y + 100$ is the old weighted average plus 100 (weights sum to one after normalising).\n",
    "\n",
    "**Step 2.** HT: adds $100 \\times \\frac1n\\sum_iT_i/\\hat e_i$, which equals 100 only if the sample weights average exactly to 1.\n",
    "\n",
    "**Answer:** Hájek shifts by exactly 100; HT shifts by 100 times a random factor, so it is not location-invariant.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 5 · Truncation\n",
    "\n",
    "Weights above 20 are capped at 20. Which quantity does this change, and what is the trade-off?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** The weighted means no longer exactly reproduce the population, so a small bias appears.\n",
    "\n",
    "**Step 2.** Variance falls, often a lot.\n",
    "\n",
    "**Answer:** a bias–variance trade-off (Day 8); report results with and without truncation.\n",
    "\n",
    "</details>\n",
    "\n",
    "## Part 2 · Proofs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 6 · The IPW identity\n",
    "\n",
    "Under unconfoundedness and $e(X) > 0$, prove $E[TY/e(X)] = E[Y(1)]$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** $TY = TY(1)$ (the observed outcome of treated units).\n",
    "\n",
    "**Step 2.** LIE: $E\\Big[\\frac{TY(1)}{e(X)}\\Big] = E\\Big[\\frac{E[TY(1) \\mid X]}{e(X)}\\Big]$.\n",
    "\n",
    "**Step 3.** Unconfoundedness: $E[TY(1) \\mid X] = E[T \\mid X]\\,E[Y(1) \\mid X] = e(X)\\,E[Y(1) \\mid X]$.\n",
    "\n",
    "**Step 4.** The $e(X)$ cancels: $E\\big[E[Y(1) \\mid X]\\big] = E[Y(1)]$. $\\blacksquare$\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 7 · ATT weights\n",
    "\n",
    "Under $Y(0) \\perp T \\mid X$ and $e(X) < 1$, prove $E\\Big[\\frac{(1 - T)e(X)}{1 - e(X)}Y\\Big] = P(T = 1)\\,E[Y(0) \\mid T = 1]$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** $(1 - T)Y = (1 - T)Y(0)$.\n",
    "\n",
    "**Step 2.** LIE and unconfoundedness: $E[(1 - T)Y(0) \\mid X] = (1 - e(X))E[Y(0) \\mid X]$.\n",
    "\n",
    "**Step 3.** So the expectation equals $E\\big[e(X)E[Y(0) \\mid X]\\big]$.\n",
    "\n",
    "**Step 4.** Also $P(T = 1)E[Y(0) \\mid T = 1] = E[TY(0)] = E\\big[E[T \\mid X]E[Y(0) \\mid X]\\big] = E\\big[e(X)E[Y(0) \\mid X]\\big]$ (unconfoundedness again). The two agree. $\\blacksquare$\n",
    "\n",
    "Dividing by $P(T = 1)$ (or normalising the weights) gives $E[Y(0) \\mid T = 1]$.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 8 · Hájek is weighted least squares\n",
    "\n",
    "Show that regressing $Y$ on an intercept and $T$ with weights $w_i = T_i/\\hat e_i + (1 - T_i)/(1 - \\hat e_i)$ gives the Hájek estimator.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** WLS with columns $(1, T)$ fits a separate weighted mean in each arm (Day 24, Exercise 12).\n",
    "\n",
    "**Step 2.** The treated arm's weighted mean is $\\sum T_iY_i/\\hat e_i\\big/\\sum T_i/\\hat e_i$; the control arm's is the analogous expression with $1/(1 - \\hat e_i)$.\n",
    "\n",
    "**Step 3.** The slope on $T$ is their difference: $\\hat\\tau_H$. $\\blacksquare$\n",
    "\n",
    "In practice, use robust SEs that account for estimating $\\hat e$ (or bootstrap the whole procedure).\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 9 · Why HT is less stable\n",
    "\n",
    "Show that the HT estimate of $E[Y(1)]$ can be written as $\\hat\\mu_{H}\\cdot\\frac1n\\sum_iT_i/\\hat e_i$, where $\\hat\\mu_H$ is the Hájek mean, and explain the consequence.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** By definition, $\\hat\\mu_H = \\sum_iT_iY_i/\\hat e_i\\big/\\sum_iT_i/\\hat e_i$.\n",
    "\n",
    "**Step 2.** Multiply both sides by $\\frac1n\\sum_iT_i/\\hat e_i$: $\\frac1n\\sum_iT_iY_i/\\hat e_i$, the HT mean.\n",
    "\n",
    "**Step 3.** The factor has expectation 1 but varies from sample to sample, especially with extreme weights.\n",
    "\n",
    "**Answer:** HT multiplies the Hájek mean by a noisy factor, adding variance proportional to the *level* of $Y$.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 10 · Weights that explode\n",
    "\n",
    "Show that $\\operatorname{Var}(TY/e(X))$ involves $E[Y(1)^2/e(X)]$, and explain why limited overlap can make IPW very imprecise.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** $E[(TY/e)^2] = E[TY(1)^2/e^2] = E\\big[E[T \\mid X]E[Y(1)^2 \\mid X]/e(X)^2\\big] = E[Y(1)^2/e(X)]$.\n",
    "\n",
    "**Step 2.** If $e(X)$ is close to 0 for a non-negligible share of units, this expectation is huge (or infinite if $e$ can be arbitrarily close to 0).\n",
    "\n",
    "**Answer:** IPW variance is driven by the smallest propensities; poor overlap means imprecise IPW, which is why diagnostics come first.\n",
    "\n",
    "</details>\n",
    "\n",
    "## Part 3 · Applied\n",
    "\n",
    "Six households in a cash-transfer evaluation, with estimated propensities $\\hat e$ and monthly savings $Y$ (\\$10):\n",
    "\n",
    "| Household | $T$ | $\\hat e$ | $Y$ |\n",
    "|---|---|---|---|\n",
    "| 1 | 1 | 0.8 | 20 |\n",
    "| 2 | 1 | 0.5 | 16 |\n",
    "| 3 | 1 | 0.2 | 10 |\n",
    "| 4 | 0 | 0.8 | 14 |\n",
    "| 5 | 0 | 0.5 | 12 |\n",
    "| 6 | 0 | 0.2 | 8 |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 11 · Naive and Horvitz–Thompson\n",
    "\n",
    "Compute the naive difference in means and the HT estimate of the ATE.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Naive: treated mean $(20 + 16 + 10)/3 \\approx 15.33$; control mean $(14 + 12 + 8)/3 \\approx 11.33$; difference 4.0.\n",
    "\n",
    "**Step 2.** HT for $Y(1)$: $\\frac16(20/0.8 + 16/0.5 + 10/0.2) = \\frac16(25 + 32 + 50) = 107/6 \\approx 17.83$.\n",
    "\n",
    "**Step 3.** HT for $Y(0)$: $\\frac16(14/0.2 + 12/0.5 + 8/0.8) = \\frac16(70 + 24 + 10) = 104/6 \\approx 17.33$.\n",
    "\n",
    "**Step 4.** $\\hat\\tau_{HT} = 17.83 - 17.33 = 0.50$.\n",
    "\n",
    "**Answer:** naive 4.0 versus HT 0.5. The treated have higher propensities (likely better-off households), and weighting removes most of the gap.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 12 · Hájek\n",
    "\n",
    "Compute the Hájek estimate of the ATE for the same households.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Treated weights $1/\\hat e$: 1.25, 2, 5 (sum 8.25). Weighted mean: $(25 + 32 + 50)/8.25 = 107/8.25 \\approx 12.97$.\n",
    "\n",
    "**Step 2.** Control weights $1/(1 - \\hat e)$: 5, 2, 1.25 (sum 8.25). Weighted mean: $104/8.25 \\approx 12.61$.\n",
    "\n",
    "**Step 3.** $\\hat\\tau_H = 12.97 - 12.61 \\approx 0.36$.\n",
    "\n",
    "**Answer:** 0.36. With six units HT and Hájek differ noticeably; Hájek's means (≈ 13) stay within the range of the data, whereas HT's (≈ 17.8) do not.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 13 · The ATT\n",
    "\n",
    "Compute the IPW estimate of the ATT.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Treated mean (weight 1 each): 15.33.\n",
    "\n",
    "**Step 2.** Control ATT weights $\\hat e/(1 - \\hat e)$: household 4: $0.8/0.2 = 4$; household 5: 1; household 6: $0.2/0.8 = 0.25$ (sum 5.25).\n",
    "\n",
    "**Step 3.** Weighted control mean: $(4 \\times 14 + 1 \\times 12 + 0.25 \\times 8)/5.25 = (56 + 12 + 2)/5.25 = 70/5.25 \\approx 13.33$.\n",
    "\n",
    "**Step 4.** ATT $\\approx 15.33 - 13.33 = 2.0$.\n",
    "\n",
    "**Answer:** 2.0 (\\$20/month). Controls that look like the treated (high $\\hat e$) get most of the weight.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 14 · Extreme weights in a hospital study\n",
    "\n",
    "In a study of a surgical procedure, one untreated patient has $\\hat e = 0.99$ and gets an ATE weight of 100; the next largest weight is 8. The estimate changes from a 3-point to a 9-point mortality reduction when this patient is dropped. What should be reported?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** A single patient driving the estimate signals an overlap problem: almost everyone like this patient is treated.\n",
    "\n",
    "**Step 2.** Report the ESS, the weight distribution and estimates with trimming (e.g. $\\hat e \\in [0.05, 0.95]$) or truncation.\n",
    "\n",
    "**Step 3.** Consider a doubly robust estimator (Day 45), which leans on an outcome model where weights are extreme and is less sensitive to a single large weight.\n",
    "\n",
    "**Answer:** do not report a single IPW number; show its sensitivity to the extreme weight and redefine the estimand for the overlap population if needed.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 15 · Survey weights and IPW\n",
    "\n",
    "National surveys oversample poor regions and publish weights equal to the inverse of each household's sampling probability. Explain the parallel with IPW and why both rely on knowing (or estimating well) a probability.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Survey sampling: each sampled household represents $1/\\pi_i$ households in the population; weighting by $1/\\pi_i$ gives unbiased population means (Horvitz & Thompson, 1952).\n",
    "\n",
    "**Step 2.** IPW: each treated unit \"was sampled\" into treatment with probability $e(X_i)$, so weighting by $1/e(X_i)$ recovers the population mean of $Y(1)$.\n",
    "\n",
    "**Step 3.** Survey designers *know* $\\pi_i$; observational analysts must *estimate* $e(X)$ and assume no unobserved confounding.\n",
    "\n",
    "**Answer:** the same estimator; the difference is whether the selection probability is known by design or estimated under assumptions.\n",
    "\n",
    "</details>\n",
    "\n",
    "## Verify in Python"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy import optimize\n",
    "\n",
    "T = np.array([1, 1, 1, 0, 0, 0]); e = np.array([.8, .5, .2, .8, .5, .2]); Y = np.array([20, 16, 10, 14, 12, 8.])\n",
    "ht = (T * Y / e).mean() - ((1 - T) * Y / (1 - e)).mean()\n",
    "hj = (T * Y / e).sum() / (T / e).sum() - ((1 - T) * Y / (1 - e)).sum() / ((1 - T) / (1 - e)).sum()\n",
    "w = e / (1 - e); c = T == 0\n",
    "print(f\"naive {Y[T == 1].mean() - Y[T == 0].mean():.2f}, HT {ht:.2f}, Hajek {hj:.2f}, \"\n",
    "      f\"ATT {Y[T == 1].mean() - (w[c] * Y[c]).sum() / w[c].sum():.2f}\")\n",
    "\n",
    "# Monte Carlo: HT vs Hajek, true vs estimated propensity scores\n",
    "rng = np.random.default_rng(44)\n",
    "res = {\"HT\": [], \"Hajek\": [], \"Hajek (true e)\": []}\n",
    "for _ in range(2000):\n",
    "    n = 500; x = rng.normal(size=n); X = np.column_stack([np.ones(n), x])\n",
    "    e_true = 1 / (1 + np.exp(-1.2 * x)); T = rng.random(n) < e_true\n",
    "    y = 50 + 5 * x + 2 * T + rng.normal(0, 3, n)\n",
    "    b = optimize.minimize(lambda b: -np.sum(T * (X @ b) - np.log1p(np.exp(X @ b))), np.zeros(2)).x\n",
    "    eh = 1 / (1 + np.exp(-X @ b))\n",
    "    res[\"HT\"].append((T * y / eh).mean() - ((1 - T) * y / (1 - eh)).mean())\n",
    "    for key, ee in ((\"Hajek\", eh), (\"Hajek (true e)\", e_true)):\n",
    "        res[key].append((T * y / ee).sum() / (T / ee).sum() - ((1 - T) * y / (1 - ee)).sum() / ((1 - T) / (1 - ee)).sum())\n",
    "for k, v in res.items():\n",
    "    print(f\"{k:15s} mean {np.mean(v):.3f} (truth 2)  SD {np.std(v):.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Key takeaways\n",
    "\n",
    "- Weighting by 1/e and 1/(1 − e) recreates the population's Y(1) and Y(0) under unconfoundedness and overlap.\n",
    "- Prefer the normalised (Hájek) estimator; use odds weights for the ATT.\n",
    "- Extreme weights signal poor overlap: check the ESS and report sensitivity to trimming."
   ]
  }
 ]
}