{
 "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-l1-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 6.1 · Selection on Observables and the Propensity Score\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- State the two identifying assumptions: conditional independence (unconfoundedness) and overlap\n",
    "- Rebuild LaLonde's test: the NSW trainees against a survey (CPS) comparison group\n",
    "- Estimate a propensity score with a logit and check overlap visually and numerically\n",
    "- Trim the sample to the region of common support\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-l1-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## When we cannot randomise"
   ]
  },
  {
   "id": "m06-l1-2",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q numpy pandas statsmodels scikit-learn 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",
    "from causaldata import nsw_mixtape, cps_mixtape\n",
    "sns.set_theme(style=\"whitegrid\")"
   ]
  },
  {
   "id": "m06-l1-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LaLonde's challenge: experimental vs observational"
   ]
  },
  {
   "id": "m06-l1-4",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** Can observational data recover what the NSW experiment found?"
   ]
  },
  {
   "id": "m06-l1-5",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "nsw = nsw_mixtape.load_pandas().data\n",
    "cps = cps_mixtape.load_pandas().data\n",
    "obs = pd.concat([nsw[nsw.treat == 1], cps], ignore_index=True)     # NSW trainees + CPS non-trainees\n",
    "obs[\"u74\"] = (obs.re74 == 0).astype(int)                            # unemployed in 1974\n",
    "obs[\"u75\"] = (obs.re75 == 0).astype(int)\n",
    "EXPERIMENTAL = 1_794\n",
    "\n",
    "naive = obs.re78[obs.treat == 1].mean() - obs.re78[obs.treat == 0].mean()\n",
    "ols = smf.ols(\"re78 ~ treat + age + educ + black + hisp + marr + nodegree + re74 + re75\",\n",
    "              data=obs).fit(cov_type=\"HC1\").params[\"treat\"]\n",
    "print(f\"Treated: {obs.treat.sum()} | CPS comparisons: {(obs.treat == 0).sum():,}\")\n",
    "print(f\"Experimental benchmark: ${EXPERIMENTAL:,}\")\n",
    "print(f\"Naive difference:       ${naive:,.0f}\")\n",
    "print(f\"OLS with controls:      ${ols:,.0f}\")\n",
    "print(obs.groupby(\"treat\")[[\"age\", \"educ\", \"black\", \"marr\", \"nodegree\", \"re74\", \"re75\"]].mean().round(2).T)"
   ]
  },
  {
   "id": "m06-l1-6",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Estimating the propensity score"
   ]
  },
  {
   "id": "m06-l1-7",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** How likely was each person to join the NSW programme, given their characteristics?"
   ]
  },
  {
   "id": "m06-l1-8",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "ps_formula = (\"treat ~ age + I(age**2) + I(age**3) + educ + I(educ**2) + marr + nodegree\"\n",
    "              \" + black + hisp + re74 + re75 + u74 + u75 + educ:re74\")\n",
    "ps_model = smf.logit(ps_formula, data=obs).fit(disp=0)\n",
    "obs[\"ps\"] = ps_model.predict(obs)\n",
    "print(obs.groupby(\"treat\")[\"ps\"].describe().round(4))"
   ]
  },
  {
   "id": "m06-l1-9",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overlap and common support"
   ]
  },
  {
   "id": "m06-l1-10",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(13, 4))\n",
    "sns.histplot(data=obs, x=\"ps\", hue=\"treat\", bins=50, stat=\"density\", common_norm=False, ax=axes[0])\n",
    "axes[0].set(title=\"Propensity score by group\", yscale=\"log\")\n",
    "obs[\"logit_ps\"] = np.log(obs.ps / (1 - obs.ps))\n",
    "sns.kdeplot(data=obs, x=\"logit_ps\", hue=\"treat\", common_norm=False, ax=axes[1])\n",
    "axes[1].set(title=\"Log-odds scale shows the overlap region\")\n",
    "plt.tight_layout(); plt.show()"
   ]
  },
  {
   "id": "m06-l1-11",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** Which comparison people are actually comparable to the trainees?"
   ]
  },
  {
   "id": "m06-l1-12",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "lo, hi = obs.loc[obs.treat == 1, \"ps\"].min(), obs.loc[obs.treat == 0, \"ps\"].max()\n",
    "support = obs[(obs.ps >= lo) & (obs.ps <= hi)]\n",
    "print(f\"Common support: [{lo:.5f}, {hi:.3f}]\")\n",
    "print(f\"Kept {len(support):,} of {len(obs):,} rows; treated kept: {support.treat.sum()} of {obs.treat.sum()}\")"
   ]
  },
  {
   "id": "m06-l1-13",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise"
   ]
  },
  {
   "id": "m06-l1-14",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ✍️ Exercise\n",
    "\n",
    "Re-estimate the propensity score with a **simple** logit (linear terms only: age, educ, black, hisp, marr,\n",
    "nodegree, re74, re75). Compare the overlap plot and the number of treated units outside common support.\n",
    "Does the specification of the propensity score matter?"
   ]
  },
  {
   "id": "m06-l1-15",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your answer here"
   ]
  },
  {
   "id": "m06-l1-16",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist and recap"
   ]
  }
 ]
}