{
 "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-l2-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 6.2 · Matching: Propensity Score and Mahalanobis Distance\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- Implement nearest-neighbour propensity score matching for the ATT\n",
    "- Implement Mahalanobis (covariate) matching and compare it with score matching\n",
    "- Check balance after matching with standardised differences\n",
    "- Bootstrap a standard error and compare estimates with the experimental benchmark\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-l2-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Rebuild the LaLonde sample"
   ]
  },
  {
   "id": "m06-l2-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",
    "from sklearn.neighbors import NearestNeighbors\n",
    "import matplotlib.pyplot as plt, seaborn as sns\n",
    "from causaldata import nsw_mixtape, cps_mixtape\n",
    "rng = np.random.default_rng(6)\n",
    "sns.set_theme(style=\"whitegrid\")\n",
    "\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",
    "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",
    "obs[\"ps\"] = smf.logit(PS_FORMULA, data=obs).fit(disp=0).predict(obs)\n",
    "obs[\"lps\"] = np.log(obs.ps / (1 - obs.ps))           # match on the log-odds: better spread\n",
    "COVS = [\"age\", \"educ\", \"black\", \"hisp\", \"marr\", \"nodegree\", \"re74\", \"re75\", \"u74\", \"u75\"]\n",
    "EXPERIMENTAL = 1_794"
   ]
  },
  {
   "id": "m06-l2-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Nearest-neighbour propensity score matching"
   ]
  },
  {
   "id": "m06-l2-4",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** What did the NSW programme do for its participants, using matched survey respondents as the counterfactual?"
   ]
  },
  {
   "id": "m06-l2-5",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def psm_att(data, k=1, score=\"lps\"):\n",
    "    t, c = data[data.treat == 1], data[data.treat == 0]\n",
    "    nn = NearestNeighbors(n_neighbors=k).fit(c[[score]])\n",
    "    _, idx = nn.kneighbors(t[[score]])\n",
    "    matched_y0 = c.re78.to_numpy()[idx].mean(axis=1)\n",
    "    return (t.re78.to_numpy() - matched_y0).mean(), c.iloc[idx.ravel()]\n",
    "\n",
    "for k in (1, 5):\n",
    "    att, _ = psm_att(obs, k=k)\n",
    "    print(f\"PSM, {k} neighbour(s): ATT = ${att:,.0f}   (experimental ${EXPERIMENTAL:,})\")"
   ]
  },
  {
   "id": "m06-l2-6",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Did matching balance the covariates?"
   ]
  },
  {
   "id": "m06-l2-7",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** After matching, do the comparison respondents look like the trainees?"
   ]
  },
  {
   "id": "m06-l2-8",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def std_diff(t, c, cols):\n",
    "    return pd.Series({v: (t[v].mean() - c[v].mean()) / np.sqrt((t[v].var() + c[v].var()) / 2) for v in cols})\n",
    "\n",
    "treated = obs[obs.treat == 1]\n",
    "_, matched_controls = psm_att(obs, k=1)\n",
    "bal = pd.DataFrame({\n",
    "    \"before\": std_diff(treated, obs[obs.treat == 0], COVS),\n",
    "    \"after\":  std_diff(treated, matched_controls, COVS),\n",
    "}).round(3)\n",
    "print(bal)\n",
    "\n",
    "bal.abs().plot.barh(figsize=(8, 4)); plt.axvline(0.1, ls=\"--\", color=\"red\")\n",
    "plt.title(\"|Standardised difference| before and after matching\"); plt.show()"
   ]
  },
  {
   "id": "m06-l2-9",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Mahalanobis matching"
   ]
  },
  {
   "id": "m06-l2-10",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def mahalanobis_att(data, cols):\n",
    "    t, c = data[data.treat == 1], data[data.treat == 0]\n",
    "    VI = np.linalg.inv(np.cov(data[cols].to_numpy(dtype=float), rowvar=False))\n",
    "    nn = NearestNeighbors(n_neighbors=1, metric=\"mahalanobis\", metric_params={\"VI\": VI}, algorithm=\"brute\")\n",
    "    nn.fit(c[cols].to_numpy(dtype=float))\n",
    "    _, idx = nn.kneighbors(t[cols].to_numpy(dtype=float))\n",
    "    return (t.re78.to_numpy() - c.re78.to_numpy()[idx[:, 0]]).mean()\n",
    "\n",
    "print(f\"Mahalanobis ATT: ${mahalanobis_att(obs, COVS):,.0f}\")"
   ]
  },
  {
   "id": "m06-l2-11",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Uncertainty with the bootstrap"
   ]
  },
  {
   "id": "m06-l2-12",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def one_bootstrap(seed):\n",
    "    b = obs.sample(frac=1, replace=True, random_state=seed)\n",
    "    b = b.assign(ps=smf.logit(PS_FORMULA, data=b).fit(disp=0).predict(b))\n",
    "    b = b.assign(lps=np.log(b.ps / (1 - b.ps)))\n",
    "    return psm_att(b, k=5)[0]\n",
    "\n",
    "boot = np.array([one_bootstrap(s) for s in range(100)])     # use 500+ in real work\n",
    "est = psm_att(obs, k=5)[0]\n",
    "print(f\"PSM (k=5) ATT ${est:,.0f}, bootstrap SE ${boot.std():,.0f}, \"\n",
    "      f\"95% CI [${np.percentile(boot, 2.5):,.0f}, ${np.percentile(boot, 97.5):,.0f}]\")"
   ]
  },
  {
   "id": "m06-l2-13",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ✍️ Exercise\n",
    "\n",
    "Match exactly on `black` and `nodegree` (only compare trainees with CPS respondents in the same cell) and\n",
    "by nearest propensity score within each cell. Does the estimate move?"
   ]
  },
  {
   "id": "m06-l2-14",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your answer here"
   ]
  },
  {
   "id": "m06-l2-15",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist and recap"
   ]
  }
 ]
}