{
 "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-l1-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 5.1 · DAGs and Structural Causal Models: Chains, Forks and Colliders\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- Draw a directed acyclic graph (DAG) for a policy question with networkx\n",
    "- Write the structural causal model (SCM) that a DAG represents and simulate it\n",
    "- Recognise the three building blocks (chains, forks, colliders) and what each implies for correlation\n",
    "- Use d-separation to read off which variables should be independent\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-l1-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## From equations to arrows"
   ]
  },
  {
   "id": "m05-l1-2",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q numpy pandas statsmodels networkx matplotlib\n",
    "import numpy as np, pandas as pd, networkx as nx\n",
    "import statsmodels.formula.api as smf\n",
    "import matplotlib.pyplot as plt\n",
    "rng = np.random.default_rng(5)"
   ]
  },
  {
   "id": "m05-l1-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Drawing DAGs in Python"
   ]
  },
  {
   "id": "m05-l1-4",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def draw_dag(edges, pos=None, title=\"\", highlight=()):\n",
    "    g = nx.DiGraph(edges)\n",
    "    pos = pos or nx.spring_layout(g, seed=3)\n",
    "    colors = [\"#f59e0b\" if n in highlight else \"#bfdbfe\" for n in g.nodes]\n",
    "    plt.figure(figsize=(7, 4))\n",
    "    nx.draw_networkx(g, pos, node_color=colors, node_size=2600, font_size=9,\n",
    "                     arrowsize=20, edgecolors=\"#334155\")\n",
    "    plt.title(title); plt.axis(\"off\"); plt.show()\n",
    "    return g\n",
    "\n",
    "job_training = draw_dag(\n",
    "    [(\"Motivation\", \"Training\"), (\"Motivation\", \"Earnings\"),\n",
    "     (\"PriorEducation\", \"Training\"), (\"PriorEducation\", \"Earnings\"),\n",
    "     (\"Training\", \"Skills\"), (\"Skills\", \"Earnings\")],\n",
    "    pos={\"Motivation\": (0, 1), \"PriorEducation\": (2, 1), \"Training\": (0, 0),\n",
    "         \"Skills\": (1, -0.6), \"Earnings\": (2, 0)},\n",
    "    title=\"Job training: what causes what?\", highlight=(\"Training\", \"Earnings\"),\n",
    ")"
   ]
  },
  {
   "id": "m05-l1-5",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Chains, forks and colliders by simulation"
   ]
  },
  {
   "id": "m05-l1-6",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** When does controlling for a third variable remove a correlation, and when does it create one?"
   ]
  },
  {
   "id": "m05-l1-7",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "N = 20_000\n",
    "def slope(formula, df, var):\n",
    "    return smf.ols(formula, data=df).fit().params[var]\n",
    "\n",
    "# Chain: X -> M -> Y\n",
    "x = rng.normal(size=N); m = 0.8 * x + rng.normal(size=N); y = 0.8 * m + rng.normal(size=N)\n",
    "chain = pd.DataFrame(dict(x=x, m=m, y=y))\n",
    "\n",
    "# Fork: X <- M -> Y\n",
    "m = rng.normal(size=N); x = 0.8 * m + rng.normal(size=N); y = 0.8 * m + rng.normal(size=N)\n",
    "fork = pd.DataFrame(dict(x=x, m=m, y=y))\n",
    "\n",
    "# Collider: X -> M <- Y\n",
    "x = rng.normal(size=N); y = rng.normal(size=N); m = x + y + rng.normal(size=N)\n",
    "collider = pd.DataFrame(dict(x=x, m=m, y=y))\n",
    "\n",
    "print(f\"{'':10}{'y ~ x':>10}{'y ~ x + m':>12}\")\n",
    "for name, df in [(\"chain\", chain), (\"fork\", fork), (\"collider\", collider)]:\n",
    "    print(f\"{name:10}{slope('y ~ x', df, 'x'):10.3f}{slope('y ~ x + m', df, 'x'):12.3f}\")"
   ]
  },
  {
   "id": "m05-l1-8",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## d-separation: reading independence off a graph"
   ]
  },
  {
   "id": "m05-l1-9",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "d_sep = getattr(nx, \"is_d_separator\", None) or nx.d_separated   # name differs across networkx versions\n",
    "\n",
    "g = nx.DiGraph([(\"X\", \"M\"), (\"Y\", \"M\")])                          # the collider\n",
    "print(\"X ⫫ Y          :\", d_sep(g, {\"X\"}, {\"Y\"}, set()))\n",
    "print(\"X ⫫ Y | M      :\", d_sep(g, {\"X\"}, {\"Y\"}, {\"M\"}))\n",
    "\n",
    "print(\"In the job-training DAG:\")\n",
    "print(\"  PriorEducation ⫫ Motivation        :\", d_sep(job_training, {\"PriorEducation\"}, {\"Motivation\"}, set()))\n",
    "print(\"  PriorEducation ⫫ Motivation | Training:\", d_sep(job_training, {\"PriorEducation\"}, {\"Motivation\"}, {\"Training\"}))"
   ]
  },
  {
   "id": "m05-l1-10",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise"
   ]
  },
  {
   "id": "m05-l1-11",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ✍️ Exercise\n",
    "\n",
    "Write an SCM for the job-training DAG above (choose your own coefficients), simulate 20,000 people, and\n",
    "verify the two d-separation statements: prior education and motivation are uncorrelated overall, but\n",
    "correlated *within* each value of Training."
   ]
  },
  {
   "id": "m05-l1-12",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your answer here"
   ]
  },
  {
   "id": "m05-l1-13",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist and recap"
   ]
  }
 ]
}