{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Day 58: Mediation: Direct & Indirect Effects\n",
    "\n",
    "*Beginner to Ninja: Math for Causal Inference*\n",
    "\n",
    "How much of an effect runs through a particular mechanism? Controlled and natural direct effects, natural indirect effects, the mediation formula that identifies them, and the strong no-confounding assumptions they need, including the case where regression \"controlling for the mediator\" goes badly wrong.\n",
    "\n",
    "**Learning outcomes**\n",
    "\n",
    "- Define total, controlled direct, natural direct and natural indirect effects with potential outcomes.\n",
    "- Compute them with the mediation formula for binary mediators and with closed-form expressions in linear models with interaction.\n",
    "- State the identification assumptions and show why mediator–outcome confounding biases the \"controlled regression\" approach.\n",
    "- Choose the right mediation quantity for a policy question."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Theory\n",
    "\n",
    "### Nested potential outcomes\n",
    "\n",
    "Let $X$ be the treatment, $M$ a mediator and $Y$ the outcome. Write $M(x)$ for the mediator under treatment $x$, and $Y(x, m)$ for the outcome when $X$ is set to $x$ and $M$ to $m$. The observed outcome under treatment $x$ is $Y(x) = Y(x, M(x))$.\n",
    "\n",
    "### Four effects (binary $X$)\n",
    "\n",
    "- **Total effect:** $TE = E[Y(1, M(1)) - Y(0, M(0))]$.\n",
    "- **Controlled direct effect** at level $m$: $CDE(m) = E[Y(1, m) - Y(0, m)]$. Fix the mediator for everyone.\n",
    "- **Natural direct effect:** $NDE = E[Y(1, M(0)) - Y(0, M(0))]$. Change the treatment but keep each unit's mediator at the value it *would have had* without treatment.\n",
    "- **Natural indirect effect:** $NIE = E[Y(1, M(1)) - Y(1, M(0))]$. Keep the treatment on and shift only the mediator.\n",
    "\n",
    "They add up: $TE = NDE + NIE$. The **proportion mediated** is $NIE/TE$.\n",
    "\n",
    "### Identification: the mediation formula\n",
    "\n",
    "Assume, possibly within levels of pre-treatment covariates $C$: no unmeasured confounding of (1) $X$ and $Y$, (2) $X$ and $M$, (3) $M$ and $Y$, and (4) no confounder of $M$ and $Y$ that is itself affected by $X$. Then (Pearl's **mediation formula**)\n",
    "\n",
    "$$\n",
    "NDE = \\sum_m \\big(E[Y \\mid X = 1, m] - E[Y \\mid X = 0, m]\\big)\\,P(m \\mid X = 0),\n",
    "$$\n",
    "\n",
    "$$\n",
    "NIE = \\sum_m E[Y \\mid X = 1, m]\\,\\big(P(m \\mid X = 1) - P(m \\mid X = 0)\\big).\n",
    "$$\n",
    "\n",
    "$CDE(m) = E[Y \\mid X = 1, m] - E[Y \\mid X = 0, m]$ needs only assumptions (1) and (3) (and can be identified even when (4) fails, by the g-formula).\n",
    "\n",
    "### Linear models with interaction\n",
    "\n",
    "If $E[M \\mid x] = \\beta_0 + \\beta_1 x$ and $E[Y \\mid x, m] = \\theta_0 + \\theta_1 x + \\theta_2 m + \\theta_3 xm$:\n",
    "\n",
    "$$\n",
    "CDE(m) = \\theta_1 + \\theta_3 m, \\qquad NDE = \\theta_1 + \\theta_3\\beta_0, \\qquad NIE = (\\theta_2 + \\theta_3)\\beta_1.\n",
    "$$\n",
    "\n",
    "Without interaction ($\\theta_3 = 0$): $NDE = CDE = \\theta_1$ and $NIE = \\theta_2\\beta_1$, the \"product of coefficients\" of Baron and Kenny.\n",
    "\n",
    "### The trap\n",
    "\n",
    "\"Regress $Y$ on $X$ and $M$ and call the $X$ coefficient the direct effect\" is right only under assumption (3). If something unobserved affects both $M$ and $Y$, $M$ is a collider on $X \\to M \\leftarrow U \\to Y$, and the coefficient is biased (Day 53, Exercise 9); it can even flip sign.\n",
    "\n",
    "## Part 1 · Warm-up"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 1 · Read the potential outcomes\n",
    "\n",
    "In words, what is $Y(1, M(0))$ for a jobseeker in a training programme where $M$ is the number of job applications sent?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** First argument: the jobseeker receives training.\n",
    "\n",
    "**Step 2.** Second argument: applications are held at the number they would have sent *without* training.\n",
    "\n",
    "**Answer:** earnings with training but without training's effect on job search, so only the skill (direct) channel operates. It is a cross-world quantity: never observed for anyone.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 2 · Linear, no interaction\n",
    "\n",
    "$M = 3 + 0.5X + e_M$ and $Y = 2 + 1.0X + 2.0M + e_Y$, with no confounding. Find $NDE$, $NIE$ and $TE$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** $NDE = \\theta_1 = 1.0$.\n",
    "\n",
    "**Step 2.** $NIE = \\theta_2\\beta_1 = 2.0 \\times 0.5 = 1.0$.\n",
    "\n",
    "**Step 3.** $TE = 1.0 + 1.0 = 2.0$.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 3 · Linear with interaction\n",
    "\n",
    "$E[M \\mid x] = 1 + 0.5x$ and $E[Y \\mid x, m] = \\theta_0 + 1.0x + 2.0m + 0.5xm$. Compute $CDE(0)$, $CDE(2)$, $NDE$, $NIE$ and $TE$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** $CDE(0) = 1.0$; $CDE(2) = 1.0 + 0.5(2) = 2.0$.\n",
    "\n",
    "**Step 2.** $NDE = \\theta_1 + \\theta_3\\beta_0 = 1.0 + 0.5(1) = 1.5$.\n",
    "\n",
    "**Step 3.** $NIE = (\\theta_2 + \\theta_3)\\beta_1 = 2.5 \\times 0.5 = 1.25$.\n",
    "\n",
    "**Step 4.** $TE = 1.5 + 1.25 = 2.75$. Check directly: $E[Y \\mid \\operatorname{do}(1)] - E[Y \\mid \\operatorname{do}(0)] = [1.0 + 2.5(1.5)] - [2.0(1)] = 4.75 - 2 = 2.75$.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 4 · Proportion mediated\n",
    "\n",
    "In Exercise 2, what share of the effect runs through $M$? In Exercise 3?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Exercise 2: $1.0/2.0 = 50\\%$.\n",
    "\n",
    "**Step 2.** Exercise 3: $1.25/2.75 \\approx 45\\%$.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 5 · Which effect answers which question?\n",
    "\n",
    "Match: (a) \"What would the effect of training be if we also guaranteed every trainee exactly 10 applications?\" (b) \"How much of training's effect is due to its effect on job search?\" (c) \"Does training work at all?\"\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** (a) $CDE(10)$: the mediator is fixed by policy.\n",
    "\n",
    "**Step 2.** (b) $NIE$ (or the proportion mediated): a mechanism question.\n",
    "\n",
    "**Step 3.** (c) $TE$.\n",
    "\n",
    "</details>\n",
    "\n",
    "## Part 2 · Proofs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 6 · The decomposition\n",
    "\n",
    "Prove $TE = NDE + NIE$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Add and subtract $Y(1, M(0))$: $Y(1, M(1)) - Y(0, M(0)) = [Y(1, M(1)) - Y(1, M(0))] + [Y(1, M(0)) - Y(0, M(0))]$.\n",
    "\n",
    "**Step 2.** Take expectations: the first bracket is $NIE$, the second $NDE$. $\\blacksquare$\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 7 · Deriving the mediation formula for the NDE\n",
    "\n",
    "Assume (a) $Y(x, m) \\perp\\!\\!\\!\\perp X$, (b) $M(x) \\perp\\!\\!\\!\\perp X$, (c) $Y(x, m) \\perp\\!\\!\\!\\perp M \\mid X$, and (d) $Y(x, m) \\perp\\!\\!\\!\\perp M(x')$ (the cross-world assumption). Prove the NDE formula.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** $E[Y(1, M(0))] = \\sum_m E[Y(1, m) \\mid M(0) = m]\\,P(M(0) = m)$.\n",
    "\n",
    "**Step 2.** By (d), $E[Y(1, m) \\mid M(0) = m] = E[Y(1, m)]$.\n",
    "\n",
    "**Step 3.** By (a) and (c), $E[Y(1, m)] = E[Y(1, m) \\mid X = 1] = E[Y(1, m) \\mid X = 1, M = m] = E[Y \\mid X = 1, M = m]$ (consistency).\n",
    "\n",
    "**Step 4.** By (b), $P(M(0) = m) = P(M(0) = m \\mid X = 0) = P(m \\mid X = 0)$.\n",
    "\n",
    "**Step 5.** So $E[Y(1, M(0))] = \\sum_m E[Y \\mid 1, m]P(m \\mid 0)$; similarly $E[Y(0, M(0))] = \\sum_m E[Y \\mid 0, m]P(m \\mid 0)$. Subtract. $\\blacksquare$\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 8 · Linear formulas\n",
    "\n",
    "With $E[M \\mid x] = \\beta_0 + \\beta_1 x$ and $E[Y \\mid x, m] = \\theta_0 + \\theta_1 x + \\theta_2 m + \\theta_3 xm$, derive $NDE$ and $NIE$ from the mediation formula.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** $E[Y \\mid 1, m] - E[Y \\mid 0, m] = \\theta_1 + \\theta_3 m$.\n",
    "\n",
    "**Step 2.** $NDE = \\sum_m(\\theta_1 + \\theta_3 m)P(m \\mid 0) = \\theta_1 + \\theta_3 E[M \\mid X = 0] = \\theta_1 + \\theta_3\\beta_0$.\n",
    "\n",
    "**Step 3.** $NIE = \\sum_m(\\theta_0 + \\theta_1 + (\\theta_2 + \\theta_3)m)\\big(P(m \\mid 1) - P(m \\mid 0)\\big)$.\n",
    "\n",
    "**Step 4.** Constants sum to zero against a difference of probabilities; the $m$ term gives $(\\theta_2 + \\theta_3)(E[M \\mid 1] - E[M \\mid 0]) = (\\theta_2 + \\theta_3)\\beta_1$. $\\blacksquare$\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 9 · No interaction, one direct effect\n",
    "\n",
    "Show that if $E[Y \\mid x, m]$ has no $x \\times m$ interaction, then $CDE(m) = NDE$ for every $m$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** No interaction: $E[Y \\mid 1, m] - E[Y \\mid 0, m] = c$, a constant.\n",
    "\n",
    "**Step 2.** $CDE(m) = c$ for all $m$.\n",
    "\n",
    "**Step 3.** $NDE = \\sum_m c\\,P(m \\mid 0) = c$. $\\blacksquare$\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 10 · A treatment-induced mediator–outcome confounder\n",
    "\n",
    "Let $X \\to L \\to M \\to Y$, $L \\to Y$, $X \\to M$, $X \\to Y$ (all observed). Explain why $CDE(m)$ is identified but the $NDE$ is not, without further assumptions.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** $CDE(m) = E[Y \\mid \\operatorname{do}(x, m)]$ is a joint intervention; with $X$ unconfounded and $\\{X, L\\}$ blocking the back-door paths of $M$, the g-formula gives $E[Y \\mid \\operatorname{do}(x, m)] = \\sum_l E[Y \\mid x, l, m]P(l \\mid x)$ (Day 55, Exercise 13).\n",
    "\n",
    "**Step 2.** The NDE needs $E[Y(1, M(0))]$, which involves $L(1)$ (through the direct path $X \\to L \\to Y$) jointly with $M(0)$, which depends on $L(0)$.\n",
    "\n",
    "**Step 3.** The joint distribution of $L(1)$ and $L(0)$ for the same unit is never observed and is not fixed by any observable distribution (a \"recanting witness\").\n",
    "\n",
    "**Answer:** CDEs are identified; natural effects need extra assumptions or are reported with bounds. $\\blacksquare$\n",
    "\n",
    "</details>\n",
    "\n",
    "## Part 3 · Applied"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 11 · Training, job finding and earnings\n",
    "\n",
    "A training programme $X$ may raise monthly earnings $Y$ (\\$100s) directly (skills) and by speeding up job finding $M$ (found a job within 3 months). With no confounding: $P(M = 1 \\mid X = 0) = 0.4$, $P(M = 1 \\mid X = 1) = 0.6$; $E[Y \\mid X, M]$: $(0, 0)$: 10, $(0, 1)$: 20, $(1, 0)$: 12, $(1, 1)$: 24. Compute $TE$, $NDE$, $NIE$, $CDE(0)$, $CDE(1)$ and the proportion mediated.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** $E[Y \\mid X = 1] = 0.4(12) + 0.6(24) = 19.2$; $E[Y \\mid X = 0] = 0.6(10) + 0.4(20) = 14$. $TE = 5.2$.\n",
    "\n",
    "**Step 2.** $NDE = (12 - 10)(0.6) + (24 - 20)(0.4) = 1.2 + 1.6 = 2.8$.\n",
    "\n",
    "**Step 3.** $NIE = 24(0.6 - 0.4) + 12(0.4 - 0.6) = 4.8 - 2.4 = 2.4$.\n",
    "\n",
    "**Step 4.** Check: $2.8 + 2.4 = 5.2$. $CDE(0) = 2$, $CDE(1) = 4$.\n",
    "\n",
    "**Answer:** about 46% of the effect ($2.4/5.2$) runs through faster job finding. The direct effect depends on the mediator level (2 or 4), so $NDE$ is a weighted average of the CDEs.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 12 · Cash transfers, food spending and child growth\n",
    "\n",
    "A cash transfer $X$ raises monthly food spending $M$ (\\$): $E[M \\mid x] = 50 + 20x$. Child height-for-age z-scores follow $E[Y \\mid x, m] = \\theta_0 + 0.10x + 0.004m + 0.002xm$. Assume no unmeasured confounding. Compute the effects and the proportion mediated by food spending.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** $NDE = 0.10 + 0.002 \\times 50 = 0.20$.\n",
    "\n",
    "**Step 2.** $NIE = (0.004 + 0.002) \\times 20 = 0.12$.\n",
    "\n",
    "**Step 3.** $TE = 0.32$; proportion mediated $0.12/0.32 = 37.5\\%$.\n",
    "\n",
    "**Answer:** most of the effect runs through other channels (health care, schooling, reduced stress); the interaction means that food spending matters more for treated households.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 13 · Smoking, birth weight and infant mortality\n",
    "\n",
    "Return to Day 48's paradox in mediation form. In a linear model, $M$ = birth weight (standardised), $U$ = birth defects (unobserved): $M = -X - U + e_M$, $Y = 0.2X - 0.5M + 1.0U + e_Y$, with $\\operatorname{Var}(U) = \\operatorname{Var}(e_M) = 1$. What is the true direct effect of smoking, and what does the regression of mortality on smoking and birth weight report?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** True direct effect: $\\theta_1 = 0.2$ (harmful).\n",
    "\n",
    "**Step 2.** Given $X$ and $M$, $-U + e_M = M + X$ is known, so $E[U \\mid X, M] = -k(M + X)$ with $k = 1/(1 + 1) = 0.5$.\n",
    "\n",
    "**Step 3.** $E[Y \\mid X, M] = 0.2X - 0.5M + 1.0 \\times (-0.5)(M + X) = -0.3X - 1.0M$.\n",
    "\n",
    "**Answer:** the regression reports $-0.3$: smoking looks *protective* holding birth weight fixed, because birth weight is a collider of smoking and defects. The simulation confirms −0.3.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 14 · Occupation and the gender pay gap\n",
    "\n",
    "A study regresses wages on gender, \"controlling for occupation\", and calls the remaining gap \"discrimination\" (a direct effect). Draw the DAG with gender $G$, occupation $O$, wage $W$ and unobserved ability or preferences $A$ that affect both occupation and wages. What is wrong, and which quantity could be defended?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** $G \\to O \\to W$, $G \\to W$, $A \\to O$, $A \\to W$.\n",
    "\n",
    "**Step 2.** Occupation is a mediator *and* a collider ($G \\to O \\leftarrow A$). Conditioning on it opens $G \\to O \\leftarrow A \\to W$, so the \"within-occupation gap\" mixes discrimination with ability differences induced by selection into occupations.\n",
    "\n",
    "**Step 3.** The total gap $TE$ is identified (gender is not caused by the other variables); the CDE given occupation needs $A$ measured.\n",
    "\n",
    "**Answer:** report the total gap, and treat within-occupation gaps as descriptive unless the occupation–wage confounders are measured. Note also that occupational sorting can itself reflect discrimination, so the \"indirect\" part is not necessarily benign.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 15 · Obesity, hypertension and stroke\n",
    "\n",
    "A health ministry can fund either a weight-loss programme or a hypertension-control programme. Researchers estimate the effect of obesity $X$ on stroke $Y$ and the part mediated by hypertension $M$. Which mediation quantity answers \"how many strokes would remain among obese people if everyone's blood pressure were controlled?\" What assumptions does it need?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Controlling blood pressure fixes the mediator at a normal level $m_0$ for everyone: $CDE(m_0)$ is the effect of obesity that remains.\n",
    "\n",
    "**Step 2.** Strokes prevented by hypertension control among the obese: $E[Y(1)] - E[Y(1, m_0)]$, which needs the same assumptions.\n",
    "\n",
    "**Step 3.** Assumptions: no unmeasured confounding of obesity and stroke, and of hypertension and stroke given obesity and covariates (e.g. diet, smoking, kidney function); no cross-world assumption is needed.\n",
    "\n",
    "**Answer:** the CDE answers the policy question; natural effects answer the mechanism question and need stronger assumptions.\n",
    "\n",
    "</details>\n",
    "\n",
    "## Verify in Python"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "# Exercise 11: mediation formula with a binary mediator\n",
    "pM = {0: .4, 1: .6}; EY = {(0, 0): 10, (0, 1): 20, (1, 0): 12, (1, 1): 24}\n",
    "P = lambda m, x: pM[x] if m else 1 - pM[x]\n",
    "TE = sum(EY[1, m] * P(m, 1) for m in (0, 1)) - sum(EY[0, m] * P(m, 0) for m in (0, 1))\n",
    "NDE = sum((EY[1, m] - EY[0, m]) * P(m, 0) for m in (0, 1))\n",
    "NIE = sum(EY[1, m] * (P(m, 1) - P(m, 0)) for m in (0, 1))\n",
    "print(f\"TE {TE:.2f}  NDE {NDE:.2f}  NIE {NIE:.2f}  proportion mediated {NIE / TE:.2f}\")\n",
    "\n",
    "# Exercises 3 and 12: linear formulas, checked by simulating nested potential outcomes\n",
    "def effects(b0, b1, t1, t2, t3, n=1_000_000, seed=58):\n",
    "    rng = np.random.default_rng(seed)\n",
    "    eM = rng.normal(size=n); eY = rng.normal(size=n)\n",
    "    M = lambda x: b0 + b1 * x + eM\n",
    "    Y = lambda x, m: t1 * x + t2 * m + t3 * x * m + eY\n",
    "    return (Y(1, M(0)) - Y(0, M(0))).mean(), (Y(1, M(1)) - Y(1, M(0))).mean()\n",
    "for args in [(1, .5, 1, 2, .5), (50, 20, .10, .004, .002)]:\n",
    "    b0, b1, t1, t2, t3 = args\n",
    "    nde, nie = effects(*args)\n",
    "    print(f\"simulated NDE {nde:.3f} (formula {t1 + t3 * b0:.3f}), NIE {nie:.3f} (formula {(t2 + t3) * b1:.3f})\")\n",
    "\n",
    "# Exercise 13: regression 'direct effect' with a confounded mediator\n",
    "rng = np.random.default_rng(13); n = 1_000_000\n",
    "X = rng.normal(size=n); U = rng.normal(size=n)\n",
    "M = -X - U + rng.normal(size=n); Y = .2 * X - .5 * M + 1.0 * U + rng.normal(size=n)\n",
    "coef = np.linalg.lstsq(np.column_stack([np.ones(n), X, M]), Y, rcond=None)[0]\n",
    "print(f\"coefficient on smoking given birth weight: {coef[1]:.3f} (true direct effect 0.2)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Key takeaways\n",
    "\n",
    "- $TE = NDE + NIE$; controlled direct effects fix the mediator by policy, natural effects describe mechanisms.\n",
    "- The mediation formula identifies natural effects under no confounding of $X$–$Y$, $X$–$M$ and $M$–$Y$, plus a cross-world assumption that fails when treatment affects a mediator–outcome confounder.\n",
    "- \"Controlling for the mediator\" gives a direct effect only if the mediator is unconfounded with the outcome; otherwise it is collider bias."
   ]
  }
 ]
}