{
 "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-l3-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 8.3 · Policy Learning: Who Should Be Treated?\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- Turn CATE estimates into a treatment rule under a budget constraint\n",
    "- Evaluate the value of a policy honestly with doubly robust scores on held-out data\n",
    "- Learn a transparent, shallow policy tree that a ministry can apply\n",
    "- Compare targeted, universal and random allocation in dollars\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-l3-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## From effects to decisions"
   ]
  },
  {
   "id": "m08-l3-2",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q econml lightgbm scikit-learn matplotlib seaborn\n",
    "import numpy as np, pandas as pd\n",
    "from econml.dml import CausalForestDML\n",
    "from econml.policy import DRPolicyTree\n",
    "from lightgbm import LGBMRegressor, LGBMClassifier\n",
    "from sklearn.model_selection import KFold\n",
    "import matplotlib.pyplot as plt\n",
    "rng = np.random.default_rng(12)\n",
    "pd.set_option(\"display.width\", 140)"
   ]
  },
  {
   "id": "m08-l3-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data, nuisance models and doubly robust scores"
   ]
  },
  {
   "id": "m08-l3-4",
   "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,\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",
    "    y = 15_000 + 400 * X.prior_earnings + 300 * (X.years_education - 12) - 120 * X.months_unemployed \\\n",
    "        + d * tau + rng.normal(0, 3_000, n)\n",
    "    return X, d, y.to_numpy(), tau.to_numpy()\n",
    "\n",
    "X_train, d_train, y_train, _ = make_data(8_000)       # used to learn the policy\n",
    "X_eval, d_eval, y_eval, tau_eval = make_data(4_000)   # used only to evaluate it\n",
    "COST = 2_000                                          # cost of one stipend"
   ]
  },
  {
   "id": "m08-l3-5",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** How do I score each person's treatment effect in a way that is unbiased on average?"
   ]
  },
  {
   "id": "m08-l3-6",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def dr_scores(X, d, y, k=5):\n",
    "    mu1, mu0, e = np.zeros(len(y)), np.zeros(len(y)), np.zeros(len(y))\n",
    "    for tr, te in KFold(k, shuffle=True, random_state=0).split(X):\n",
    "        Xtr, dtr, ytr = X.iloc[tr], d[tr], y[tr]\n",
    "        mu1[te] = LGBMRegressor(n_estimators=300, verbose=-1).fit(Xtr[dtr == 1], ytr[dtr == 1]).predict(X.iloc[te])\n",
    "        mu0[te] = LGBMRegressor(n_estimators=300, verbose=-1).fit(Xtr[dtr == 0], ytr[dtr == 0]).predict(X.iloc[te])\n",
    "        e[te] = LGBMClassifier(n_estimators=200, verbose=-1).fit(Xtr, dtr).predict_proba(X.iloc[te])[:, 1]\n",
    "    e = e.clip(0.02, 0.98)\n",
    "    return mu1 - mu0 + d * (y - mu1) / e - (1 - d) * (y - mu0) / (1 - e)\n",
    "\n",
    "gamma = dr_scores(X_eval, d_eval, y_eval)\n",
    "print(f\"DR estimate of ATE on the evaluation sample: {gamma.mean():,.0f} (true {tau_eval.mean():,.0f})\")"
   ]
  },
  {
   "id": "m08-l3-7",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Budget-constrained targeting"
   ]
  },
  {
   "id": "m08-l3-8",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** If we can fund stipends for only 30% of applicants, whom should we choose, and what do we gain?"
   ]
  },
  {
   "id": "m08-l3-9",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "cf = CausalForestDML(model_y=LGBMRegressor(n_estimators=300, verbose=-1),\n",
    "                     model_t=LGBMClassifier(n_estimators=300, verbose=-1),\n",
    "                     discrete_treatment=True, n_estimators=800, min_samples_leaf=20, random_state=0)\n",
    "cf.fit(y_train, d_train, X=X_train)\n",
    "cate_eval = cf.effect(X_eval)\n",
    "\n",
    "def value(policy):                       # net dollars per applicant vs treating nobody\n",
    "    return np.mean(policy * (gamma - COST))\n",
    "\n",
    "budget = 0.30\n",
    "rules = {\n",
    "    \"Treat nobody\":            np.zeros(len(X_eval)),\n",
    "    \"Treat everyone\":          np.ones(len(X_eval)),\n",
    "    \"Random 30%\":              rng.binomial(1, budget, len(X_eval)),\n",
    "    \"Longest unemployed 30%\":  (X_eval.months_unemployed >= X_eval.months_unemployed.quantile(1 - budget)).astype(int).to_numpy(),\n",
    "    \"Causal forest top 30%\":   (cate_eval >= np.quantile(cate_eval, 1 - budget)).astype(int),\n",
    "    \"Oracle top 30%\":          (tau_eval >= np.quantile(tau_eval, 1 - budget)).astype(int),\n",
    "}\n",
    "report = pd.DataFrame({\n",
    "    \"share treated\": {k: v.mean() for k, v in rules.items()},\n",
    "    \"DR value ($/applicant)\": {k: value(v) for k, v in rules.items()},\n",
    "    \"true value ($/applicant)\": {k: np.mean(v * (tau_eval - COST)) for k, v in rules.items()},\n",
    "}).round(2)\n",
    "print(report)"
   ]
  },
  {
   "id": "m08-l3-10",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## An interpretable policy tree"
   ]
  },
  {
   "id": "m08-l3-11",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "tree = DRPolicyTree(\n",
    "    max_depth=2, min_samples_leaf=200, random_state=0,\n",
    "    model_regression=LGBMRegressor(n_estimators=300, verbose=-1),\n",
    "    model_propensity=LGBMClassifier(n_estimators=200, verbose=-1),\n",
    ")\n",
    "# Encode the cost: treatment 'pays off' only if the effect exceeds COST\n",
    "tree.fit(y_train - COST * d_train, d_train, X=X_train)\n",
    "tree_policy = tree.predict(X_eval)\n",
    "print(f\"Policy tree treats {tree_policy.mean():.0%} | DR value {value(tree_policy):,.0f} $/applicant \"\n",
    "      f\"| true value {np.mean(tree_policy * (tau_eval - COST)):,.0f}\")\n",
    "tree.plot(feature_names=list(X_train.columns), treatment_names=[\"No stipend\", \"Stipend\"])\n",
    "plt.show()"
   ]
  },
  {
   "id": "m08-l3-12",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise"
   ]
  },
  {
   "id": "m08-l3-13",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ✍️ Exercise\n",
    "\n",
    "Plot the DR value of \"causal forest top q%\" for budgets q from 5% to 100%. Where does the value peak, and\n",
    "how does that relate to the share of people whose effect exceeds the cost?"
   ]
  },
  {
   "id": "m08-l3-14",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your answer here"
   ]
  },
  {
   "id": "m08-l3-15",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist and recap"
   ]
  }
 ]
}