{
 "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 43: The Propensity Score Theorem\n",
    "\n",
    "*Beginner to Ninja: Math for Causal Inference*\n",
    "\n",
    "One number, the probability of treatment given covariates, is enough to remove all confounding by those covariates. Why it works and how to use it.\n",
    "\n",
    "**Learning outcomes**\n",
    "\n",
    "- Prove the balancing property of the propensity score and the Rosenbaum–Rubin theorem.\n",
    "- Estimate propensity scores by logit and check covariate balance within propensity strata.\n",
    "- Use subclassification on the propensity score to estimate ATE and ATT.\n",
    "- Explain what the propensity score can and cannot do (it does not fix unobserved confounding)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Theory\n",
    "\n",
    "### The problem of many covariates\n",
    "\n",
    "With many covariates, exact matching on $X$ is impossible: almost every unit has a unique covariate combination. Rosenbaum & Rubin (1983) showed that a single summary suffices.\n",
    "\n",
    "### Two results\n",
    "\n",
    "The **propensity score** is $e(X) = P(T = 1 \\mid X)$.\n",
    "\n",
    "1. **Balancing property:** $X \\perp T \\mid e(X)$. Among units with the same propensity score, the distribution of covariates is the same for treated and controls.\n",
    "2. **Propensity score theorem:** if $(Y(0), Y(1)) \\perp T \\mid X$, then $(Y(0), Y(1)) \\perp T \\mid e(X)$.\n",
    "\n",
    "So instead of adjusting for a vector $X$, it is enough to adjust for the scalar $e(X)$: match, stratify or weight on it.\n",
    "\n",
    "### Using it in practice\n",
    "\n",
    "1. **Estimate** $e(X)$, typically by logit (Day 11) with flexible terms (interactions, squares) or machine learning.\n",
    "2. **Check overlap** (Day 42) and trim if needed.\n",
    "3. **Check balance:** within propensity strata (or after weighting), normalised differences of each covariate should be small. Balance is the diagnostic for whether the score model is good enough, not its fit or its $R^2$.\n",
    "4. **Estimate** effects: subclassification (5 strata typically remove about 90% of the bias from a single covariate, Cochran 1968), matching (Day 81), or weighting (Day 44).\n",
    "\n",
    "### What it cannot do\n",
    "\n",
    "The propensity score only balances the covariates used to build it. It is a tool for implementing unconfoundedness, not a substitute for it.\n",
    "\n",
    "## Part 1 · Warm-up"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 1 · A propensity score from a logit\n",
    "\n",
    "A logit for enrolling in a training programme has $\\log\\frac{e}{1 - e} = 1.0 - 0.1 \\times \\text{age} + 0.5 \\times \\text{unemployed}$. Compute $e$ for a 30-year-old unemployed worker.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Index: $1.0 - 3.0 + 0.5 = -1.5$.\n",
    "\n",
    "**Step 2.** $e = 1/(1 + e^{1.5}) = 1/(1 + 4.48) \\approx 0.18$.\n",
    "\n",
    "**Answer:** an 18% chance of enrolling.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 2 · Same score, different covariates\n",
    "\n",
    "A 25-year-old employed worker has index $1.0 - 2.5 + 0 = -1.5$. Compare with Exercise 1. What does the balancing property say about such pairs?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Both workers have the same index, hence $e \\approx 0.18$.\n",
    "\n",
    "**Step 2.** They differ in age and employment, but are equally likely to enrol.\n",
    "\n",
    "**Answer:** among units with $e = 0.18$, the mix of (age, employment) profiles is the same for enrollees and non-enrollees, so comparing them does not confound the effect with these covariates.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 3 · Subclassification estimate\n",
    "\n",
    "Five propensity strata of equal size have within-stratum treated-minus-control differences in earnings (\\$1,000) of 0.5, 1.0, 1.5, 1.2, 2.0. Estimate the ATE.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Equal stratum sizes give equal weights of 1/5.\n",
    "\n",
    "**Step 2.** ATE $\\approx (0.5 + 1.0 + 1.5 + 1.2 + 2.0)/5 = 6.2/5 = 1.24$.\n",
    "\n",
    "**Answer:** about \\$1,240.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 4 · ATT by subclassification\n",
    "\n",
    "In Exercise 3, the strata contain 5%, 10%, 20%, 25% and 40% of all treated units. Estimate the ATT.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Weight strata by their share of treated units.\n",
    "\n",
    "**Step 2.** ATT $\\approx 0.05(0.5) + 0.10(1.0) + 0.20(1.5) + 0.25(1.2) + 0.40(2.0)$.\n",
    "\n",
    "**Step 3.** $= 0.025 + 0.10 + 0.30 + 0.30 + 0.80 = 1.525$.\n",
    "\n",
    "**Answer:** about \\$1,525: the treated sit mostly in high-propensity strata, where effects are larger.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 5 · Balance within a stratum\n",
    "\n",
    "In one propensity stratum, treated workers average 11.2 years of schooling (sd 2.5) and controls 10.9 (sd 2.4). Compute the normalised difference. Is balance acceptable?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Pooled sd: $\\sqrt{(6.25 + 5.76)/2} \\approx 2.45$.\n",
    "\n",
    "**Step 2.** $\\Delta = 0.3/2.45 \\approx 0.12$.\n",
    "\n",
    "**Answer:** 0.12 < 0.25: acceptable. Check every covariate in every stratum.\n",
    "\n",
    "</details>\n",
    "\n",
    "## Part 2 · Proofs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 6 · The balancing property\n",
    "\n",
    "Prove $P(T = 1 \\mid X, e(X)) = P(T = 1 \\mid e(X))$, i.e. $X \\perp T \\mid e(X)$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** $e(X)$ is a function of $X$, so conditioning on $(X, e(X))$ is conditioning on $X$: $P(T = 1 \\mid X, e(X)) = e(X)$.\n",
    "\n",
    "**Step 2.** LIE: $P(T = 1 \\mid e(X)) = E[T \\mid e(X)] = E\\big[E[T \\mid X] \\mid e(X)\\big] = E[e(X) \\mid e(X)] = e(X)$.\n",
    "\n",
    "**Step 3.** Both equal $e(X)$, so knowing $X$ beyond $e(X)$ does not change the treatment probability: $T \\perp X \\mid e(X)$. $\\blacksquare$\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 7 · The propensity score theorem\n",
    "\n",
    "Assume $(Y(0), Y(1)) \\perp T \\mid X$. Prove $P(T = 1 \\mid Y(0), Y(1), e(X)) = e(X)$, i.e. unconfoundedness given $e(X)$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "Write $\\mathbf Y = (Y(0), Y(1))$.\n",
    "\n",
    "**Step 1.** LIE: $P(T = 1 \\mid \\mathbf Y, e(X)) = E\\big[E[T \\mid \\mathbf Y, X] \\mid \\mathbf Y, e(X)\\big]$ (conditioning on $X$ refines conditioning on $e(X)$).\n",
    "\n",
    "**Step 2.** Unconfoundedness given $X$: $E[T \\mid \\mathbf Y, X] = E[T \\mid X] = e(X)$.\n",
    "\n",
    "**Step 3.** So $P(T = 1 \\mid \\mathbf Y, e(X)) = E[e(X) \\mid \\mathbf Y, e(X)] = e(X)$.\n",
    "\n",
    "**Step 4.** This does not depend on $\\mathbf Y$, so $T \\perp \\mathbf Y \\mid e(X)$. $\\blacksquare$\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 8 · Adjusting for the score identifies the ATE\n",
    "\n",
    "Using Exercise 7, show $\\text{ATE} = E\\big[E[Y \\mid T = 1, e(X)] - E[Y \\mid T = 0, e(X)]\\big]$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Apply Day 42, Exercise 6 with $e(X)$ playing the role of $X$: unconfoundedness holds given $e(X)$ (Exercise 7), and overlap $0 < e < 1$ holds by assumption.\n",
    "\n",
    "**Step 2.** So $E[Y(t)] = E\\big[E[Y \\mid T = t, e(X)]\\big]$ for $t = 0, 1$; subtract. $\\blacksquare$\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 9 · Any balancing score works\n",
    "\n",
    "A **balancing score** is any $b(X)$ with $X \\perp T \\mid b(X)$. Show that $X$ itself is a balancing score, and that $e(X)$ is the *coarsest* one (it is a function of every balancing score).\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Given $X$, $X$ is fixed, so trivially independent of $T$: $X$ is a balancing score (the finest).\n",
    "\n",
    "**Step 2.** If $b(X)$ is balancing, then $e(X) = P(T = 1 \\mid X) = P(T = 1 \\mid X, b(X)) = P(T = 1 \\mid b(X))$, using the balancing property in the last step.\n",
    "\n",
    "**Step 3.** So $e(X)$ can be computed from $b(X)$: it is a function of every balancing score. $\\blacksquare$\n",
    "\n",
    "The propensity score achieves the maximum dimension reduction.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 10 · The propensity score does not fix hidden bias\n",
    "\n",
    "Suppose $T$ depends on $X$ and an unobserved $U$ that also affects $Y$. Explain why adjusting for $e(X) = P(T = 1 \\mid X)$ does not remove confounding by $U$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Exercise 7 needed unconfoundedness given $X$; here it fails because of $U$.\n",
    "\n",
    "**Step 2.** Within a stratum of $e(X)$, treated units still have different $U$ (hence different $Y(0)$) than controls.\n",
    "\n",
    "**Answer:** the propensity score balances only what goes into it. Good balance on $X$ says nothing about $U$.\n",
    "\n",
    "</details>\n",
    "\n",
    "## Part 3 · Applied"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 11 · Subclassification in a training study\n",
    "\n",
    "A training study estimates propensity scores and forms five strata. Within strata the effects and sizes are:\n",
    "\n",
    "| Stratum | Treated | Controls | Effect (\\$1,000) |\n",
    "|---|---|---|---|\n",
    "| 1 | 5 | 395 | 0.4 |\n",
    "| 2 | 10 | 190 | 0.8 |\n",
    "| 3 | 20 | 80 | 1.2 |\n",
    "| 4 | 40 | 40 | 1.6 |\n",
    "| 5 | 60 | 15 | 2.0 |\n",
    "\n",
    "Estimate the ATE and the ATT.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Stratum sizes: 400, 200, 100, 80, 75; total 855. Treated: 135.\n",
    "\n",
    "**Step 2.** ATE weights $N_s/855$: $\\frac{400(0.4) + 200(0.8) + 100(1.2) + 80(1.6) + 75(2.0)}{855} = \\frac{160 + 160 + 120 + 128 + 150}{855} = \\frac{718}{855} \\approx 0.84$.\n",
    "\n",
    "**Step 3.** ATT weights $N_{1s}/135$: $\\frac{5(0.4) + 10(0.8) + 20(1.2) + 40(1.6) + 60(2.0)}{135} = \\frac{2 + 8 + 24 + 64 + 120}{135} = \\frac{218}{135} \\approx 1.61$.\n",
    "\n",
    "**Answer:** ATE ≈ \\$840, ATT ≈ \\$1,610: the treated are concentrated where effects are large. Stratum 1's 5 treated and stratum 5's 15 controls signal thin overlap at the extremes.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 12 · Checking the score model\n",
    "\n",
    "After subclassification, baseline earnings are balanced in four strata but in stratum 5 treated workers earn \\$3,000 more (normalised difference 0.6). What should the analyst do?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Imbalance means the score model does not capture how earnings relate to treatment at high propensities.\n",
    "\n",
    "**Step 2.** Re-specify: add earnings squared, interactions (earnings × unemployed), or split stratum 5.\n",
    "\n",
    "**Step 3.** Re-check balance; iterate. Do not look at outcomes while doing this, so the model is chosen for balance, not for the answer.\n",
    "\n",
    "**Answer:** refine the propensity model until balance is achieved, blind to outcomes.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 13 · Smoking and birth weight\n",
    "\n",
    "A study of maternal smoking and birth weight uses a propensity score based on age, education, parity and prenatal visits. Name a plausible unobserved confounder and the likely direction of bias.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Candidates: alcohol use, diet, stress, or partner's smoking: correlated with smoking and harmful to the baby.\n",
    "\n",
    "**Step 2.** These make smokers' babies lighter even apart from smoking, so the estimated harm of smoking is overstated (biased away from zero).\n",
    "\n",
    "**Answer:** e.g. alcohol use; the propensity score cannot fix it, but a sensitivity analysis can show how strong it would need to be.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 14 · Many covariates, few treated\n",
    "\n",
    "A hospital study has 80 treated patients and 60 candidate covariates. Why can a saturated logit for the propensity score be a problem, and what are the alternatives?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** With 60 covariates and 80 treated, the logit overfits: estimated scores pile up near 0 and 1, sometimes with perfect separation (Day 25).\n",
    "\n",
    "**Step 2.** Overfitted scores exaggerate apparent non-overlap and produce extreme weights.\n",
    "\n",
    "**Step 3.** Alternatives: choose covariates by subject-matter knowledge (those affecting both treatment and outcome), penalised logit (ridge or lasso), or methods that target balance directly (entropy balancing, Day 84).\n",
    "\n",
    "**Answer:** a propensity model is a means to balance, not a prediction contest.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 15 · Is a high-AUC propensity model good news?\n",
    "\n",
    "A propensity model predicts enrolment in a job-search programme almost perfectly (area under the ROC curve 0.98). Is that good or bad for estimating the programme's effect?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Near-perfect prediction means scores close to 0 or 1 for most units.\n",
    "\n",
    "**Step 2.** That signals poor overlap: few comparable treated and control units.\n",
    "\n",
    "**Answer:** bad. Unlike in prediction tasks, a highly accurate propensity model warns that causal comparisons will rely on few units or on extrapolation.\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",
    "rng = np.random.default_rng(43)\n",
    "n = 20_000\n",
    "age = rng.normal(35, 10, n); unemp = rng.random(n) < 0.4\n",
    "X = np.column_stack([np.ones(n), age, unemp])\n",
    "T = rng.random(n) < 1 / (1 + np.exp(-(1.0 - 0.06 * age + 1.0 * unemp)))\n",
    "y0 = 20 + 0.3 * age - 5 * unemp + rng.normal(0, 4, n)\n",
    "y1 = y0 + 1 + 2 * unemp                                       # effect larger for the unemployed\n",
    "y = np.where(T, y1, y0)\n",
    "\n",
    "nll = lambda b: -np.sum(T * (X @ b) - np.log1p(np.exp(X @ b)))\n",
    "b = optimize.minimize(nll, np.zeros(3)).x\n",
    "e = 1 / (1 + np.exp(-X @ b))\n",
    "strata = np.digitize(e, np.quantile(e, [0.2, 0.4, 0.6, 0.8]))\n",
    "ate = sum((strata == s).mean() * (y[T & (strata == s)].mean() - y[~T & (strata == s)].mean()) for s in range(5))\n",
    "att = sum((T & (strata == s)).sum() / T.sum() * (y[T & (strata == s)].mean() - y[~T & (strata == s)].mean()) for s in range(5))\n",
    "# five strata remove most, but not all, of the confounding bias\n",
    "print(f\"naive {y[T].mean() - y[~T].mean():.2f}; subclassification ATE {ate:.2f} (truth {(y1 - y0).mean():.2f}), \"\n",
    "      f\"ATT {att:.2f} (truth {(y1 - y0)[T].mean():.2f})\")\n",
    "for name, v in ((\"age\", age), (\"unemployed\", unemp.astype(float))):\n",
    "    nd = [(v[T & (strata == s)].mean() - v[~T & (strata == s)].mean()) / v.std() for s in range(5)]\n",
    "    print(f\"normalised differences within strata, {name}:\", np.round(nd, 2))\n",
    "\n",
    "print(\"Exercise 11: ATE\", round(718 / 855, 3), \" ATT\", round(218 / 135, 3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Key takeaways\n",
    "\n",
    "- Conditional on the propensity score, covariates are balanced between treated and controls.\n",
    "- If treatment is unconfounded given X, it is unconfounded given e(X): one number replaces many covariates.\n",
    "- Judge a propensity model by the balance it achieves, never by its predictive accuracy; it cannot fix unobserved confounding."
   ]
  }
 ]
}