{
 "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 57: Selection Bias in DAGs\n",
    "\n",
    "*Beginner to Ninja: Math for Causal Inference*\n",
    "\n",
    "Who ends up in your data is itself a variable. Add a selection node S to the DAG; analysing only S = 1 is conditioning on S, and when S is a collider or a descendant of the outcome, associations, even in randomised trials, are distorted. Graphs say when the bias can be undone and how.\n",
    "\n",
    "**Learning outcomes**\n",
    "\n",
    "- Add a selection node to a DAG and diagnose collider-stratification, outcome-dependent and survivorship bias.\n",
    "- State and prove the conditions under which P(y | x) and odds ratios are recoverable from selected data.\n",
    "- Correct selection bias with external covariate data or inverse-probability-of-selection weights.\n",
    "- Recognise selection bias in case-control studies, testing data, wage regressions and trials with loss to follow-up."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Theory\n",
    "\n",
    "### The selection node\n",
    "\n",
    "Let $S = 1$ indicate that a unit is in the data set. Draw $S$ in the DAG with arrows from whatever determines inclusion. Every analysis of the sample conditions on $S = 1$, so the d-separation rules apply: if $S$ is a collider (or a descendant of one) on a path between $X$ and $Y$, that path is opened.\n",
    "\n",
    "### Typical structures\n",
    "\n",
    "| Structure | Example | Effect |\n",
    "|---|---|---|\n",
    "| $X \\to S \\leftarrow Y$ | hospital patients (Berkson), tested individuals | spurious association |\n",
    "| $X \\to S \\leftarrow U \\to Y$ | trial drop-out due to side effects and poor health | bias even with randomised $X$ |\n",
    "| $Y \\to S$ | wages observed only for the employed, case-control sampling | distorted $P(y \\mid x)$ |\n",
    "| $X \\to M \\to S$, $M \\leftarrow U \\to Y$ | studying survivors, firms that survived | survivorship bias |\n",
    "\n",
    "### Recoverability\n",
    "\n",
    "With only selected data we can estimate $P(\\mathbf v \\mid S = 1)$. A quantity is **recoverable** if it can be written in terms of this (plus any external data).\n",
    "\n",
    "- **Conditional distributions:** $P(y \\mid x)$ is recoverable from the selected sample alone if $(Y \\perp\\!\\!\\!\\perp S \\mid X)$ in the graph; then $P(y \\mid x) = P(y \\mid x, S = 1)$ (Bareinboim and Pearl).\n",
    "- **With external data on covariates:** if $(Y \\perp\\!\\!\\!\\perp S \\mid X, \\mathbf Z)$ and $P(\\mathbf z \\mid x)$ is known from a population source (a census),\n",
    "\n",
    "$$\n",
    "P(y \\mid x) = \\sum_{\\mathbf z} P(y \\mid x, \\mathbf z, S = 1)\\,P(\\mathbf z \\mid x).\n",
    "$$\n",
    "\n",
    "- **Weighting:** equivalently, if $Y \\perp\\!\\!\\!\\perp S \\mid \\mathbf Z$ and $P(S = 1 \\mid \\mathbf z) > 0$ is known or estimable, weight selected units by $1/P(S = 1 \\mid \\mathbf z)$ (inverse probability of selection weighting; Day 14 used the same idea for missing data).\n",
    "- **Odds ratios:** if selection depends only on $Y$ (and on nothing else that relates to $X$), the $X$–$Y$ odds ratio is unchanged by selection. That is why case-control studies estimate odds ratios, not risks.\n",
    "\n",
    "### Causal effects under selection\n",
    "\n",
    "For $P(y \\mid \\operatorname{do}(x))$ we need two things at once: a back-door set $\\mathbf Z$, and a way to undo selection, for example $(Y \\perp\\!\\!\\!\\perp S \\mid X, \\mathbf Z)$ plus population data on $\\mathbf Z$:\n",
    "\n",
    "$$\n",
    "P(y \\mid \\operatorname{do}(x)) = \\sum_{\\mathbf z} P(y \\mid x, \\mathbf z, S = 1)\\,P(\\mathbf z).\n",
    "$$\n",
    "\n",
    "Randomisation removes confounding but **not** selection bias created after randomisation.\n",
    "\n",
    "## Part 1 · Warm-up"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 1 · Is selection a collider?\n",
    "\n",
    "For each, is the $X$–$Y$ association in the selected sample biased? (a) $X \\to Y$, $X \\to S$. (b) $X \\to Y$, $Y \\to S$. (c) $X \\perp\\!\\!\\!\\perp Y$, $X \\to S \\leftarrow Y$. (d) $X \\to Y$, $Z \\to S$, with $Z$ unrelated to $X, Y$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** (a) $Y \\perp\\!\\!\\!\\perp S \\mid X$ (the only path $Y \\leftarrow X \\to S$ is blocked by $X$): $P(y \\mid x)$ unbiased.\n",
    "\n",
    "**Step 2.** (b) $Y$ and $S$ are adjacent: $P(y \\mid x, S = 1) \\ne P(y \\mid x)$ in general.\n",
    "\n",
    "**Step 3.** (c) Collider: a spurious association appears.\n",
    "\n",
    "**Step 4.** (d) $S$ is d-separated from $X, Y$: no bias (just a smaller sample).\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 2 · Check recoverability\n",
    "\n",
    "DAG: $Z \\to X \\to Y$, $Z \\to Y$, $Z \\to S$. Is $P(y \\mid x)$ recoverable from the selected sample alone? Is it recoverable with census data on $P(z \\mid x)$?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Paths $Y$–$S$: $Y \\leftarrow Z \\to S$ and $Y \\leftarrow X \\leftarrow Z \\to S$. Given $X$ alone, the first stays open: $(Y \\perp\\!\\!\\!\\perp S \\mid X)$ fails. Not recoverable from the sample alone.\n",
    "\n",
    "**Step 2.** Given $X$ and $Z$, both paths are blocked: $(Y \\perp\\!\\!\\!\\perp S \\mid X, Z)$.\n",
    "\n",
    "**Answer:** recoverable with external data: $P(y \\mid x) = \\sum_z P(y \\mid x, z, S = 1)P(z \\mid x)$.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 3 · Odds ratios survive outcome selection\n",
    "\n",
    "In a population, $P(Y = 1 \\mid X = 1) = 0.2$ and $P(Y = 1 \\mid X = 0) = 0.1$. A study keeps all units with $Y = 1$ and 10% of units with $Y = 0$. Compute the odds ratio in the population and in the sample.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Population OR: $\\dfrac{0.2/0.8}{0.1/0.9} = \\dfrac{0.25}{0.111} = 2.25$.\n",
    "\n",
    "**Step 2.** In the sample, the odds for $X = 1$ become $\\dfrac{0.2}{0.8 \\times 0.1} = 2.5$ and for $X = 0$, $\\dfrac{0.1}{0.9 \\times 0.1} = 1.111$.\n",
    "\n",
    "**Step 3.** Sample OR: $2.5/1.111 = 2.25$.\n",
    "\n",
    "**Answer:** 2.25 in both; the factor 0.1 cancels.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 4 · Re-weighting to the population\n",
    "\n",
    "In a selected sample, $P(Y = 1 \\mid X = 1, Z = 1, S = 1) = 0.5$ and $P(Y = 1 \\mid X = 1, Z = 0, S = 1) = 0.2$, and $(Y \\perp\\!\\!\\!\\perp S \\mid X, Z)$. A census shows $P(Z = 1 \\mid X = 1) = 0.3$, while in the sample it is 0.6. Recover $P(Y = 1 \\mid X = 1)$ and compare with the naive sample value.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Recovered: $0.3(0.5) + 0.7(0.2) = 0.15 + 0.14 = 0.29$.\n",
    "\n",
    "**Step 2.** Naive (sample mix): $0.6(0.5) + 0.4(0.2) = 0.30 + 0.08 = 0.38$.\n",
    "\n",
    "**Answer:** 0.29; the sample over-represents $Z = 1$.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 5 · Survivors\n",
    "\n",
    "A study of 10-year-old firms finds that firms that started with high debt are as profitable as low-debt firms. Draw a DAG with debt $D$, survival $S$, managerial quality $Q$ (unobserved) and profitability $P$, and explain why the finding may mislead.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** $D \\to S$ (debt raises failure risk), $Q \\to S$, $Q \\to P$, possibly $D \\to P$.\n",
    "\n",
    "**Step 2.** Studying survivors conditions on the collider $S$: among survivors, high-debt firms needed better managers to survive, so they have higher $Q$.\n",
    "\n",
    "**Answer:** the path $D \\to S \\leftarrow Q \\to P$ is opened, pushing high-debt survivors' profits up; the effect of debt is understated.\n",
    "\n",
    "</details>\n",
    "\n",
    "## Part 2 · Proofs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 6 · Selection independent of the outcome\n",
    "\n",
    "Prove that if $(Y \\perp\\!\\!\\!\\perp S \\mid X)$, then $P(y \\mid x, S = 1) = P(y \\mid x)$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** By the global Markov property, the d-separation implies $Y \\perp\\!\\!\\!\\perp S \\mid X$ in the distribution.\n",
    "\n",
    "**Step 2.** So $P(y \\mid x, s) = P(y \\mid x)$ for every $s$ with $P(x, s) > 0$, in particular $s = 1$. $\\blacksquare$\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 7 · Odds-ratio invariance\n",
    "\n",
    "Suppose selection depends only on $Y$: $P(S = 1 \\mid x, y) = s_y$. Prove that the odds ratio of $X$ and $Y$ is the same in the selected sample as in the population.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** $P(x, y \\mid S = 1) = \\dfrac{P(x, y)\\,s_y}{P(S = 1)}$.\n",
    "\n",
    "**Step 2.** The odds ratio is $\\dfrac{P(1, 1)P(0, 0)}{P(1, 0)P(0, 1)}$ (arguments $x, y$).\n",
    "\n",
    "**Step 3.** In the sample, each cell is multiplied by $s_y/P(S = 1)$; numerator and denominator each contain one $y = 1$ and one $y = 0$ cell, so the factors $s_1 s_0/P(S = 1)^2$ cancel. $\\blacksquare$\n",
    "\n",
    "Risk ratios do not cancel this way, so case-control studies cannot estimate absolute risks without external information.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 8 · Recovery with external covariate data\n",
    "\n",
    "Prove: if $(Y \\perp\\!\\!\\!\\perp S \\mid X, \\mathbf Z)$, then $P(y \\mid x) = \\sum_{\\mathbf z} P(y \\mid x, \\mathbf z, S = 1)P(\\mathbf z \\mid x)$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Law of total probability: $P(y \\mid x) = \\sum_{\\mathbf z} P(y \\mid x, \\mathbf z)P(\\mathbf z \\mid x)$.\n",
    "\n",
    "**Step 2.** By the d-separation and Exercise 6's argument, $P(y \\mid x, \\mathbf z) = P(y \\mid x, \\mathbf z, S = 1)$.\n",
    "\n",
    "**Step 3.** Substitute. $\\blacksquare$ The first factor comes from the sample, the second from the census.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 9 · Inverse probability of selection weighting\n",
    "\n",
    "Assume $Y \\perp\\!\\!\\!\\perp S \\mid \\mathbf Z$ and $\\pi(\\mathbf z) = P(S = 1 \\mid \\mathbf z) > 0$. Prove $E[Y] = E\\big[SY/\\pi(\\mathbf Z)\\big]$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** By the LIE, $E[SY/\\pi(\\mathbf Z)] = E\\big[E[SY \\mid \\mathbf Z]/\\pi(\\mathbf Z)\\big]$.\n",
    "\n",
    "**Step 2.** $E[SY \\mid \\mathbf Z] = E[S \\mid \\mathbf Z]E[Y \\mid \\mathbf Z] = \\pi(\\mathbf Z)E[Y \\mid \\mathbf Z]$, by conditional independence.\n",
    "\n",
    "**Step 3.** So $E[SY/\\pi(\\mathbf Z)] = E[E[Y \\mid \\mathbf Z]] = E[Y]$. $\\blacksquare$\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 10 · Randomisation does not protect against drop-out\n",
    "\n",
    "In a trial, $X$ is randomised, $Y = \\tau X + U$, and units stay in the study ($S = 1$) when $X + U + e > 0$ (the programme's support visits encourage staying, and healthier patients stay), with $U, e$ independent normals. Explain with d-separation why $E[Y \\mid X = 1, S = 1] - E[Y \\mid X = 0, S = 1] \\ne \\tau$ in general.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** DAG: $X \\to Y$, $U \\to Y$, $X \\to S$, $U \\to S$.\n",
    "\n",
    "**Step 2.** Conditioning on $S$ opens $X \\to S \\leftarrow U \\to Y$.\n",
    "\n",
    "**Step 3.** Here $X = 1$ makes staying easier, so treated units stay even with low $U$, while controls need a higher $U$ to stay. Hence $E[U \\mid X = 1, S = 1] < E[U \\mid X = 0, S = 1]$.\n",
    "\n",
    "**Step 4.** The completers' contrast equals $\\tau + E[U \\mid X = 1, S = 1] - E[U \\mid X = 0, S = 1]$. $\\blacksquare$\n",
    "\n",
    "</details>\n",
    "\n",
    "## Part 3 · Applied"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 11 · A case-control study of smoking and lung cancer\n",
    "\n",
    "In the population, 1% have lung cancer $D$. Among cases 60% smoke; among non-cases 20% smoke. A case-control study samples cases and an equal number of controls. Compute the odds ratio and the risk ratio in the population and in the study.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Population OR (exposure odds among cases versus controls): $\\dfrac{0.6/0.4}{0.2/0.8} = \\dfrac{1.5}{0.25} = 6$.\n",
    "\n",
    "**Step 2.** Population risks: $P(D \\mid \\text{smoke}) = \\dfrac{0.01(0.6)}{0.01(0.6) + 0.99(0.2)} = \\dfrac{0.006}{0.204} \\approx 0.0294$; $P(D \\mid \\text{no smoke}) = \\dfrac{0.004}{0.004 + 0.792} \\approx 0.00503$. RR $\\approx 5.85$.\n",
    "\n",
    "**Step 3.** In the 1:1 study (per 1 case, 1 control): smokers $0.6$ cases and $0.2$ controls, non-smokers $0.4$ and $0.8$. OR $= (0.6 \\times 0.8)/(0.4 \\times 0.2) = 6$.\n",
    "\n",
    "**Step 4.** Study \"risks\": $0.6/0.8 = 0.75$ and $0.4/1.2 \\approx 0.333$; RR $= 2.25$.\n",
    "\n",
    "**Answer:** the OR (6) is preserved; the risk ratio is not (2.25 in the sample vs 5.85). With a rare disease the OR approximates the RR.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 12 · Does smoking protect against COVID-19?\n",
    "\n",
    "In a population, 20% smoke ($X$) and 5% are infected ($Y$), independently. Only tested people are in the data: $P(\\text{tested} \\mid x, y) = 0.02 + 0.10x + 0.60y$ (smokers cough more; infection causes symptoms). Compute the test positivity for smokers and non-smokers among the tested.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Smokers: tested and infected $0.05 \\times 0.72 = 0.036$; tested and not infected $0.95 \\times 0.12 = 0.114$.\n",
    "\n",
    "**Step 2.** Positivity among tested smokers: $0.036/(0.036 + 0.114) = 0.24$.\n",
    "\n",
    "**Step 3.** Non-smokers: $0.05 \\times 0.62 = 0.031$ and $0.95 \\times 0.02 = 0.019$; positivity $0.031/0.050 = 0.62$.\n",
    "\n",
    "**Answer:** 24% versus 62%: smokers appear protected, although smoking and infection are independent. Testing is a collider ($X \\to S \\leftarrow Y$). Early pandemic studies reported exactly this pattern.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 13 · Women's wages and selection into work\n",
    "\n",
    "Wage offers depend on education: $\\log W^* = 1.0 + 0.10E + \\varepsilon$. Women work ($S = 1$) if the offer exceeds a reservation wage that rises with the number of young children $K$. Wages are observed only for workers. In the simulation below (38% work), what does the regression of observed log wages on education find, and why?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** DAG: $E \\to W^* \\to S \\leftarrow K$. Selection depends on the outcome $W^*$ itself.\n",
    "\n",
    "**Step 2.** $(W^* \\perp\\!\\!\\!\\perp S \\mid E)$ fails ($W^* \\to S$): $P(w \\mid e, S = 1) \\ne P(w \\mid e)$.\n",
    "\n",
    "**Step 3.** Among workers, low-education women are in the sample only if they drew a high $\\varepsilon$, which raises their average wage and flattens the education gradient.\n",
    "\n",
    "**Answer:** the simulation finds a slope of about 0.056 instead of 0.10. Heckman's selection model uses $K$, which affects work but not wage offers, to correct this: the same exclusion idea as an instrument.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 14 · Drop-out in a trial of an exercise programme\n",
    "\n",
    "A trial randomises an exercise programme $X$ for heart patients; outcome $Y$ is fitness at one year. Baseline health $L$ (observed) raises fitness and lowers drop-out; the programme itself raises drop-out (it is demanding). In the simulation below, compare the completers' difference in means with the truth (1.0), and with an inverse-probability-of-completion-weighted (IPCW) analysis.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** DAG: $X \\to Y$, $L \\to Y$, $X \\to S$, $L \\to S$. Completers condition on $S$, opening $X \\to S \\leftarrow L \\to Y$.\n",
    "\n",
    "**Step 2.** Completers' contrast: about 1.35. Programme completers are healthier at baseline than control completers, because less healthy patients drop out of the demanding programme more.\n",
    "\n",
    "**Step 3.** $(Y \\perp\\!\\!\\!\\perp S \\mid X, L)$ holds, so weighting completers by $1/P(S = 1 \\mid X, L)$ recovers the full-trial comparison (Exercise 9).\n",
    "\n",
    "**Answer:** IPCW gives about 1.01; the completers' analysis overstates the benefit by a third.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 15 · Studying successful entrepreneurs\n",
    "\n",
    "A business school surveys founders of firms that reached \\$10 million in revenue and finds that founders without a university degree performed as well as graduates, concluding that \"education doesn't matter for entrepreneurial success\". Draw the DAG and explain the flaw. What data would answer the question?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** DAG: education $E \\to$ success $S$, talent and luck $T \\to S$; the sample is selected on $S$ itself (the outcome).\n",
    "\n",
    "**Step 2.** Within the sample, the outcome no longer varies much, and among successes, founders without degrees must have had more talent or luck: conditioning on the collider opens $E \\to S \\leftarrow T$.\n",
    "\n",
    "**Answer:** selecting on the outcome destroys the comparison. Follow a cohort of all founders (or all potential founders) and compare success rates by education, adjusting for pre-founding confounders.\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: case-control sampling\n",
    "pD, pS_case, pS_ctrl = .01, .6, .2\n",
    "or_pop = (pS_case / (1 - pS_case)) / (pS_ctrl / (1 - pS_ctrl))\n",
    "r1 = pD * pS_case / (pD * pS_case + (1 - pD) * pS_ctrl); r0 = pD * (1 - pS_case) / (pD * (1 - pS_case) + (1 - pD) * (1 - pS_ctrl))\n",
    "print(f\"population OR {or_pop:.2f}, RR {r1 / r0:.2f};  1:1 study OR {(.6 * .8) / (.4 * .2):.2f}, RR {(.6 / .8) / (.4 / 1.2):.2f}\")\n",
    "\n",
    "# Exercise 12: COVID testing\n",
    "for x in (1, 0):\n",
    "    pos = .05 * (.02 + .10 * x + .60); neg = .95 * (.02 + .10 * x)\n",
    "    print(f\"smoker={x}: positivity among tested {pos / (pos + neg):.2f}\")\n",
    "\n",
    "rng = np.random.default_rng(57); n = 1_000_000\n",
    "# Exercise 13: selection into work\n",
    "E = rng.normal(12, 2, n); W = 1.0 + .10 * E + rng.normal(0, .5, n); K = rng.binomial(2, .5, n)\n",
    "S = (W - .3 * K + rng.normal(0, .3, n)) > 2.1\n",
    "print(f\"work share {S.mean():.2f}; slope all {np.polyfit(E, W, 1)[0]:.3f}, workers only {np.polyfit(E[S], W[S], 1)[0]:.3f}\")\n",
    "\n",
    "# Exercise 14: drop-out in a trial and IPCW\n",
    "X = rng.random(n) < .5; L = rng.normal(size=n); Y = 1.0 * X + 2 * L + rng.normal(size=n)\n",
    "pS = 1 / (1 + np.exp(-(1.0 - 1.0 * X + 1.5 * L))); R = rng.random(n) < pS     # L = baseline health\n",
    "w = 1 / pS\n",
    "ipcw = (w * Y * (X & R)).sum() / (w * (X & R)).sum() - (w * Y * (~X & R)).sum() / (w * (~X & R)).sum()\n",
    "print(f\"completion {R.mean():.2f}; full trial {Y[X].mean() - Y[~X].mean():.3f}, completers {Y[X & R].mean() - Y[~X & R].mean():.3f}, IPCW {ipcw:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Key takeaways\n",
    "\n",
    "- Being in the data is a variable; analysing a selected sample conditions on it.\n",
    "- Selection on a collider, on a descendant of the outcome, or on a post-treatment variable biases associations, even in randomised trials.\n",
    "- Recovery needs a d-separation between outcome and selection given what you can measure, plus external data or selection weights; odds ratios survive outcome-only selection."
   ]
  }
 ]
}