{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "cells": [
  {
   "id": "m08-l1-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 8.1 · Conditional Average Treatment Effects and Meta-Learners\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- Define the CATE and explain why the ATE can hide who benefits\n",
    "- Implement S-, T- and X-learners with scikit-learn\n",
    "- Evaluate CATE estimates against the truth in a simulation (PEHE)\n",
    "- Summarise heterogeneity by group (GATEs)\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": "m08-l1-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## From one number to a function"
   ]
  },
  {
   "id": "m08-l1-2",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q numpy pandas scikit-learn lightgbm matplotlib seaborn\n",
    "import numpy as np, pandas as pd\n",
    "from lightgbm import LGBMRegressor, LGBMClassifier\n",
    "import matplotlib.pyplot as plt, seaborn as sns\n",
    "rng = np.random.default_rng(10)\n",
    "sns.set_theme(style=\"whitegrid\")"
   ]
  },
  {
   "id": "m08-l1-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A retraining programme with heterogeneous effects"
   ]
  },
  {
   "id": "m08-l1-4",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** Which unemployed workers gain most from a retraining stipend?"
   ]
  },
  {
   "id": "m08-l1-5",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def make_data(n):\n",
    "    X = pd.DataFrame({\n",
    "        \"months_unemployed\": rng.gamma(2, 4, n),\n",
    "        \"years_education\": rng.integers(8, 21, n).astype(float),\n",
    "        \"age\": rng.uniform(20, 60, n),\n",
    "        \"prior_earnings\": rng.lognormal(10.3, 0.4, n) / 1_000,     # thousands of dollars\n",
    "        \"female\": rng.binomial(1, 0.5, n).astype(float),\n",
    "    })\n",
    "    tau = 500 + 250 * np.minimum(X.months_unemployed, 24) - 150 * (X.years_education - 12)\n",
    "    p = 1 / (1 + np.exp(-(-0.5 + 0.05 * X.months_unemployed - 0.03 * (X.age - 40))))\n",
    "    d = rng.binomial(1, p)\n",
    "    y0 = 15_000 + 400 * X.prior_earnings + 300 * (X.years_education - 12) - 120 * X.months_unemployed \\\n",
    "         + rng.normal(0, 3_000, n)\n",
    "    y = y0 + d * tau\n",
    "    return X, d.astype(float), y.to_numpy(), tau.to_numpy()\n",
    "\n",
    "X, d, y, tau = make_data(6_000)\n",
    "X_test, _, _, tau_test = make_data(2_000)\n",
    "print(f\"True ATE: {tau.mean():,.0f} | effects range from {tau.min():,.0f} to {tau.max():,.0f}\")"
   ]
  },
  {
   "id": "m08-l1-6",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## S-, T- and X-learners"
   ]
  },
  {
   "id": "m08-l1-7",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** How do I turn an off-the-shelf boosting model into a treatment-effect estimator?"
   ]
  },
  {
   "id": "m08-l1-8",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def gbm():\n",
    "    return LGBMRegressor(n_estimators=300, learning_rate=0.05, num_leaves=15, verbose=-1)\n",
    "\n",
    "# S-learner\n",
    "s = gbm().fit(X.assign(d=d), y)\n",
    "cate_s = s.predict(X_test.assign(d=1.0)) - s.predict(X_test.assign(d=0.0))\n",
    "\n",
    "# T-learner\n",
    "mu1 = gbm().fit(X[d == 1], y[d == 1]); mu0 = gbm().fit(X[d == 0], y[d == 0])\n",
    "cate_t = mu1.predict(X_test) - mu0.predict(X_test)\n",
    "\n",
    "# X-learner\n",
    "d1 = y[d == 1] - mu0.predict(X[d == 1])          # imputed effects for the treated\n",
    "d0 = mu1.predict(X[d == 0]) - y[d == 0]          # imputed effects for the controls\n",
    "tau1 = gbm().fit(X[d == 1], d1); tau0 = gbm().fit(X[d == 0], d0)\n",
    "e = LGBMClassifier(n_estimators=200, learning_rate=0.05, verbose=-1).fit(X, d).predict_proba(X_test)[:, 1]\n",
    "cate_x = e * tau0.predict(X_test) + (1 - e) * tau1.predict(X_test)\n",
    "\n",
    "print(f\"Share treated: {d.mean():.2f}\")"
   ]
  },
  {
   "id": "m08-l1-9",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluating and summarising CATEs"
   ]
  },
  {
   "id": "m08-l1-10",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "pehe = {name: np.sqrt(np.mean((est - tau_test) ** 2))\n",
    "        for name, est in [(\"S-learner\", cate_s), (\"T-learner\", cate_t), (\"X-learner\", cate_x)]}\n",
    "print(pd.Series(pehe).round(0).rename(\"PEHE ($)\"))\n",
    "\n",
    "fig, axes = plt.subplots(1, 3, figsize=(15, 4), sharey=True)\n",
    "for ax, (name, est) in zip(axes, [(\"S\", cate_s), (\"T\", cate_t), (\"X\", cate_x)]):\n",
    "    ax.scatter(tau_test, est, s=5, alpha=.3); ax.plot([tau_test.min(), tau_test.max()], [tau_test.min(), tau_test.max()], \"r--\")\n",
    "    ax.set(title=f\"{name}-learner\", xlabel=\"True CATE\", ylabel=\"Estimated CATE\")\n",
    "plt.tight_layout(); plt.show()"
   ]
  },
  {
   "id": "m08-l1-11",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** How large is the stipend's effect for short-, medium- and long-term unemployed workers?"
   ]
  },
  {
   "id": "m08-l1-12",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "groups = pd.cut(X_test.months_unemployed, [0, 3, 12, 100], labels=[\"< 3 months\", \"3–12 months\", \"> 12 months\"])\n",
    "gate = pd.DataFrame({\"true\": tau_test, \"X-learner\": cate_x, \"group\": groups}).groupby(\"group\", observed=True).mean().round(0)\n",
    "print(gate)"
   ]
  },
  {
   "id": "m08-l1-13",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise"
   ]
  },
  {
   "id": "m08-l1-14",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ✍️ Exercise\n",
    "\n",
    "Make the treated group small: change the intercept of the participation model so only about 10% are\n",
    "treated. Re-run the three learners. Which one degrades most, and why?"
   ]
  },
  {
   "id": "m08-l1-15",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your answer here"
   ]
  },
  {
   "id": "m08-l1-16",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist and recap"
   ]
  }
 ]
}