{
 "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-l2-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 8.2 · Causal Forests with EconML\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- Explain how a causal forest adapts random forests to estimate treatment effects honestly\n",
    "- Fit a CausalForestDML model with EconML, separating effect modifiers (X) from controls (W)\n",
    "- Report CATEs with confidence intervals and identify the main drivers of heterogeneity\n",
    "- Test whether heterogeneity is real with a calibration-style check\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-l2-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What a causal forest does differently"
   ]
  },
  {
   "id": "m08-l2-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 lightgbm import LGBMRegressor, LGBMClassifier\n",
    "import matplotlib.pyplot as plt, seaborn as sns\n",
    "rng = np.random.default_rng(11)\n",
    "sns.set_theme(style=\"whitegrid\")"
   ]
  },
  {
   "id": "m08-l2-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fitting CausalForestDML"
   ]
  },
  {
   "id": "m08-l2-4",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** What is the estimated effect of the stipend for each unemployed worker, with a confidence interval?"
   ]
  },
  {
   "id": "m08-l2-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",
    "        \"female\": rng.binomial(1, 0.5, n).astype(float),\n",
    "    })\n",
    "    W = pd.DataFrame({\"prior_earnings\": rng.lognormal(10.3, 0.4, n) / 1_000, \"region\": rng.integers(0, 5, 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) + 0.02 * (W.prior_earnings - 30))))\n",
    "    d = rng.binomial(1, p)\n",
    "    y = 15_000 + 400 * W.prior_earnings + 300 * (X.years_education - 12) - 120 * X.months_unemployed \\\n",
    "        + 500 * W.region + d * tau + rng.normal(0, 3_000, n)\n",
    "    return X, W, d, y.to_numpy(), tau.to_numpy()\n",
    "\n",
    "X, W, d, y, tau = make_data(8_000)\n",
    "X_test, W_test, _, _, tau_test = make_data(2_000)\n",
    "\n",
    "cf = CausalForestDML(\n",
    "    model_y=LGBMRegressor(n_estimators=300, learning_rate=0.05, verbose=-1),\n",
    "    model_t=LGBMClassifier(n_estimators=300, learning_rate=0.05, verbose=-1),\n",
    "    discrete_treatment=True, n_estimators=1_000, min_samples_leaf=20, cv=5, random_state=0,\n",
    ")\n",
    "cf.fit(y, d, X=X, W=W)\n",
    "cate = cf.effect(X_test)\n",
    "lo, hi = cf.effect_interval(X_test, alpha=0.05)\n",
    "print(f\"ATE estimate: {cf.ate(X):,.0f} (true {tau.mean():,.0f})\")\n",
    "print(f\"PEHE: {np.sqrt(np.mean((cate - tau_test) ** 2)):,.0f}\")\n",
    "print(f\"Coverage of 95% CATE intervals: {np.mean((tau_test >= lo) & (tau_test <= hi)):.2f}\")"
   ]
  },
  {
   "id": "m08-l2-6",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Intervals and drivers of heterogeneity"
   ]
  },
  {
   "id": "m08-l2-7",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** Which characteristics drive differences in the stipend's effect?"
   ]
  },
  {
   "id": "m08-l2-8",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "imp = pd.Series(cf.feature_importances_, index=X.columns).sort_values()\n",
    "imp.plot.barh(figsize=(7, 3), title=\"Causal forest: drivers of effect heterogeneity\"); plt.show()\n",
    "\n",
    "order = np.argsort(X_test.months_unemployed.to_numpy())\n",
    "plt.figure(figsize=(10, 4))\n",
    "plt.fill_between(X_test.months_unemployed.to_numpy()[order], lo[order], hi[order], alpha=.15, label=\"95% interval\")\n",
    "plt.scatter(X_test.months_unemployed, cate, s=4, alpha=.4, label=\"Estimated CATE\")\n",
    "plt.scatter(X_test.months_unemployed, tau_test, s=4, alpha=.4, label=\"True CATE\")\n",
    "plt.xlabel(\"Months unemployed\"); plt.ylabel(\"Effect on earnings ($)\"); plt.legend(); plt.show()"
   ]
  },
  {
   "id": "m08-l2-9",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Is the heterogeneity real?"
   ]
  },
  {
   "id": "m08-l2-10",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "from sklearn.model_selection import cross_val_predict\n",
    "X_t, W_t, d_t, y_t, tau_t = make_data(4_000)                       # fresh sample, never seen by the forest\n",
    "F = pd.concat([X_t, W_t], axis=1)\n",
    "e_hat = cross_val_predict(LGBMClassifier(n_estimators=200, verbose=-1), F, d_t, cv=5, method=\"predict_proba\")[:, 1].clip(0.02, 0.98)\n",
    "# Outcome models by treatment arm, cross-fitted, for doubly robust (AIPW) scores\n",
    "mu1, mu0 = np.zeros(len(y_t)), np.zeros(len(y_t))\n",
    "from sklearn.model_selection import KFold\n",
    "for tr, te in KFold(5, shuffle=True, random_state=0).split(F):\n",
    "    m1 = LGBMRegressor(n_estimators=300, verbose=-1).fit(F.iloc[tr][d_t[tr] == 1], y_t[tr][d_t[tr] == 1])\n",
    "    m0 = LGBMRegressor(n_estimators=300, verbose=-1).fit(F.iloc[tr][d_t[tr] == 0], y_t[tr][d_t[tr] == 0])\n",
    "    mu1[te], mu0[te] = m1.predict(F.iloc[te]), m0.predict(F.iloc[te])\n",
    "dr = mu1 - mu0 + d_t * (y_t - mu1) / e_hat - (1 - d_t) * (y_t - mu0) / (1 - e_hat)\n",
    "\n",
    "quartile = pd.qcut(cf.effect(X_t), 4, labels=[\"Q1 (lowest)\", \"Q2\", \"Q3\", \"Q4 (highest)\"])\n",
    "check = pd.DataFrame({\"quartile\": quartile, \"doubly robust effect\": dr, \"true\": tau_t})\n",
    "print(check.groupby(\"quartile\", observed=True).mean().round(0))"
   ]
  },
  {
   "id": "m08-l2-11",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist and recap"
   ]
  }
 ]
}