{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "cells": [
  {
   "id": "m05-l2-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 5.2 · The Backdoor Criterion, Adjustment Sets and the Front Door\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- State the backdoor criterion and find valid adjustment sets programmatically\n",
    "- Show by simulation that controlling for a mediator removes the effect you want to measure\n",
    "- Estimate an effect with the front-door formula when the confounder is unobserved\n",
    "- Classify candidate controls as good, neutral or bad\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": "m05-l2-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Backdoor paths"
   ]
  },
  {
   "id": "m05-l2-2",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q numpy pandas statsmodels networkx\n",
    "import itertools, numpy as np, pandas as pd, networkx as nx\n",
    "import statsmodels.formula.api as smf\n",
    "rng = np.random.default_rng(21)\n",
    "d_sep = getattr(nx, \"is_d_separator\", None) or nx.d_separated"
   ]
  },
  {
   "id": "m05-l2-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Finding adjustment sets in code"
   ]
  },
  {
   "id": "m05-l2-4",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** Which variables must I control for to estimate the effect of job training on earnings?"
   ]
  },
  {
   "id": "m05-l2-5",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def backdoor_sets(g, treatment, outcome, observed, max_size=4):\n",
    "    g_back = g.copy()\n",
    "    g_back.remove_edges_from(list(g.out_edges(treatment)))           # keep only backdoor paths\n",
    "    forbidden = nx.descendants(g, treatment)\n",
    "    candidates = [v for v in observed if v not in forbidden and v not in (treatment, outcome)]\n",
    "    valid = []\n",
    "    for k in range(len(candidates) + 1):\n",
    "        for z in itertools.combinations(candidates, k):\n",
    "            if k <= max_size and d_sep(g_back, {treatment}, {outcome}, set(z)):\n",
    "                valid.append(set(z))\n",
    "    return valid\n",
    "\n",
    "dag = nx.DiGraph([\n",
    "    (\"Motivation\", \"Training\"), (\"Motivation\", \"Earnings\"),\n",
    "    (\"PriorEducation\", \"Training\"), (\"PriorEducation\", \"Earnings\"),\n",
    "    (\"Training\", \"Skills\"), (\"Skills\", \"Earnings\"),\n",
    "])\n",
    "\n",
    "print(\"If Motivation is observed:\")\n",
    "for s in backdoor_sets(dag, \"Training\", \"Earnings\", [\"Motivation\", \"PriorEducation\", \"Skills\"]):\n",
    "    print(\"  \", s)\n",
    "print(\"If Motivation is NOT observed:\",\n",
    "      backdoor_sets(dag, \"Training\", \"Earnings\", [\"PriorEducation\", \"Skills\"]) or \"no valid set\")"
   ]
  },
  {
   "id": "m05-l2-6",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Good and bad controls by simulation"
   ]
  },
  {
   "id": "m05-l2-7",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** What happens to the estimated training effect if I also control for skills, or for prior education?"
   ]
  },
  {
   "id": "m05-l2-8",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "N = 20_000\n",
    "motivation = rng.normal(size=N)\n",
    "prior_educ = rng.normal(size=N)\n",
    "training = rng.binomial(1, 1 / (1 + np.exp(-(0.8 * motivation + 0.8 * prior_educ))))\n",
    "skills = 2.0 * training + rng.normal(size=N)                       # mediator\n",
    "earnings = 1_000 * skills + 1_500 * motivation + 1_200 * prior_educ + rng.normal(0, 1_000, N)\n",
    "df = pd.DataFrame(dict(training=training, skills=skills, motivation=motivation,\n",
    "                       prior_educ=prior_educ, earnings=earnings))\n",
    "TRUE_TOTAL = 2.0 * 1_000\n",
    "\n",
    "specs = {\n",
    "    \"no controls\": \"earnings ~ training\",\n",
    "    \"backdoor set {Motivation, PriorEducation}\": \"earnings ~ training + motivation + prior_educ\",\n",
    "    \"+ mediator Skills (bad control)\": \"earnings ~ training + motivation + prior_educ + skills\",\n",
    "}\n",
    "print(f\"True total effect: {TRUE_TOTAL:,.0f}\")\n",
    "for label, f in specs.items():\n",
    "    print(f\"{label:45s} {smf.ols(f, data=df).fit().params['training']:8,.0f}\")"
   ]
  },
  {
   "id": "m05-l2-9",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The front-door criterion"
   ]
  },
  {
   "id": "m05-l2-10",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "step1 = smf.ols(\"skills ~ training\", data=df).fit().params[\"training\"]\n",
    "step2 = smf.ols(\"earnings ~ skills + training\", data=df).fit().params[\"skills\"]\n",
    "print(f\"Front-door estimate: {step1 * step2:,.0f} (true {TRUE_TOTAL:,.0f}), \"\n",
    "      \"without ever observing motivation in the formula\")"
   ]
  },
  {
   "id": "m05-l2-11",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist and recap"
   ]
  }
 ]
}