{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "cells": [
  {
   "id": "m07-l2-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 7.2 · From Partialling Out to Double Machine Learning\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- Rebuild the partially linear model as \"FWL with machine learning\"\n",
    "- Show that using the same data to fit and to residualise creates bias, and fix it with cross-fitting\n",
    "- Implement a cross-fitted DML estimator by hand with random forests\n",
    "- Explain why DML gives valid confidence intervals even though ML models are biased\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": "m07-l2-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The partially linear model"
   ]
  },
  {
   "id": "m07-l2-2",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q numpy pandas scikit-learn statsmodels matplotlib\n",
    "import numpy as np, pandas as pd\n",
    "from sklearn.ensemble import RandomForestRegressor\n",
    "from sklearn.model_selection import KFold\n",
    "import statsmodels.api as sm\n",
    "rng = np.random.default_rng(2)"
   ]
  },
  {
   "id": "m07-l2-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## High-dimensional, non-linear confounding"
   ]
  },
  {
   "id": "m07-l2-4",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** What is the wage effect of digital-skills training when 20 covariates confound it in unknown ways?"
   ]
  },
  {
   "id": "m07-l2-5",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "N, P = 3_000, 20\n",
    "X = rng.normal(0, 1, size=(N, P))\n",
    "m_x = 1 / (1 + np.exp(-(0.8 * X[:, 0] - 0.5 * X[:, 1] ** 2 + 0.4 * X[:, 2] * X[:, 3])))   # P(D=1 | X)\n",
    "D = rng.binomial(1, m_x)\n",
    "g_x = 2_000 * np.sin(2 * X[:, 0]) + 3_000 * X[:, 3] + 1_500 * X[:, 1] ** 2 + 1_000 * X[:, 2] * X[:, 3]\n",
    "Y = 5_000 * D + g_x + rng.normal(0, 4_000, N)\n",
    "TRUE = 5_000\n",
    "\n",
    "lin = sm.OLS(Y, sm.add_constant(np.column_stack([D, X]))).fit()\n",
    "print(f\"OLS with all 20 controls (linear): {lin.params[1]:,.0f}  (truth {TRUE:,})\")"
   ]
  },
  {
   "id": "m07-l2-6",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Naive ML plug-in vs residual-on-residual"
   ]
  },
  {
   "id": "m07-l2-7",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def rf():\n",
    "    return RandomForestRegressor(n_estimators=300, min_samples_leaf=5, max_features=0.5, n_jobs=-1, random_state=0)\n",
    "\n",
    "# Plug-in: one model for Y on (D, X)\n",
    "full = rf().fit(np.column_stack([D, X]), Y)\n",
    "plug_in = (full.predict(np.column_stack([np.ones(N), X])) - full.predict(np.column_stack([np.zeros(N), X]))).mean()\n",
    "\n",
    "# Residual-on-residual, in-sample (no cross-fitting)\n",
    "res_y = Y - rf().fit(X, Y).predict(X)\n",
    "res_d = D - rf().fit(X, D).predict(X)\n",
    "in_sample = (res_d @ res_y) / (res_d @ res_d)\n",
    "print(f\"Plug-in forest:                   {plug_in:,.0f}\")\n",
    "print(f\"Residual-on-residual, in-sample:  {in_sample:,.0f}\")"
   ]
  },
  {
   "id": "m07-l2-8",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cross-fitting by hand"
   ]
  },
  {
   "id": "m07-l2-9",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** How do I let machine learning control for confounders without overfitting my causal estimate?"
   ]
  },
  {
   "id": "m07-l2-10",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def dml_plr(Y, D, X, learner, k=5, seed=0):\n",
    "    res_y, res_d = np.zeros_like(Y, dtype=float), np.zeros_like(D, dtype=float)\n",
    "    for train, test in KFold(k, shuffle=True, random_state=seed).split(X):\n",
    "        res_y[test] = Y[test] - learner().fit(X[train], Y[train]).predict(X[test])\n",
    "        res_d[test] = D[test] - learner().fit(X[train], D[train]).predict(X[test])\n",
    "    theta = (res_d @ res_y) / (res_d @ res_d)\n",
    "    eps = res_y - theta * res_d\n",
    "    se = np.sqrt(np.mean(res_d ** 2 * eps ** 2) / np.mean(res_d ** 2) ** 2 / len(Y))\n",
    "    return theta, se\n",
    "\n",
    "theta, se = dml_plr(Y, D, X, rf)\n",
    "print(f\"Cross-fitted DML: {theta:,.0f}  (SE {se:,.0f}; 95% CI [{theta - 1.96*se:,.0f}, {theta + 1.96*se:,.0f}])\")"
   ]
  },
  {
   "id": "m07-l2-11",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ✍️ Exercise\n",
    "\n",
    "Run `dml_plr` with `k=2` and `k=10`, and with a weaker learner (a forest with `max_depth=2`). Which choice\n",
    "moves the estimate most? What does that say about tuning the nuisance models?"
   ]
  },
  {
   "id": "m07-l2-12",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your answer here"
   ]
  },
  {
   "id": "m07-l2-13",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist and recap"
   ]
  }
 ]
}