{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "cells": [
  {
   "id": "m02-l3-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 2.3 · Conditional Probability: the Precursor to Confounding\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- Compute conditional probabilities from data with grouped aggregations\n",
    "- Test whether two variables are independent\n",
    "- Simulate a world where age drives both treatment and outcome\n",
    "- See how a naive comparison mixes the treatment effect with confounding, and how stratifying fixes it\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": "m02-l3-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conditional probability in plain words"
   ]
  },
  {
   "id": "m02-l3-2",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q numpy polars matplotlib seaborn\n",
    "import numpy as np, polars as pl\n",
    "import matplotlib.pyplot as plt, seaborn as sns\n",
    "rng = np.random.default_rng(101)\n",
    "sns.set_theme(style=\"whitegrid\")"
   ]
  },
  {
   "id": "m02-l3-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A DGP where risk depends on age"
   ]
  },
  {
   "id": "m02-l3-4",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** How much more likely is a comorbidity for patients over 60?"
   ]
  },
  {
   "id": "m02-l3-5",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "N = 50_000\n",
    "age = np.clip(rng.normal(45, 12, N), 18, 90)\n",
    "log_odds = -4.0 + 0.07 * age\n",
    "p_comorbidity = 1 / (1 + np.exp(-log_odds))           # sigmoid\n",
    "comorbidity = rng.binomial(1, p_comorbidity)\n",
    "\n",
    "df = pl.DataFrame({\"age\": age, \"comorbidity\": comorbidity})\n",
    "risk = (\n",
    "    df.with_columns((pl.col(\"age\") > 60).alias(\"over_60\"))\n",
    "    .group_by(\"over_60\")\n",
    "    .agg(pl.col(\"comorbidity\").mean().round(3).alias(\"P(comorbidity | group)\"),\n",
    "         pl.len().alias(\"patients\"))\n",
    "    .sort(\"over_60\")\n",
    ")\n",
    "print(risk)\n",
    "\n",
    "xs = np.linspace(18, 90, 200)\n",
    "plt.plot(xs, 1 / (1 + np.exp(-(-4 + 0.07 * xs))))\n",
    "plt.xlabel(\"Age\"); plt.ylabel(\"P(comorbidity)\"); plt.title(\"Risk rises with age (logistic DGP)\"); plt.show()"
   ]
  },
  {
   "id": "m02-l3-6",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Independence and how to check it"
   ]
  },
  {
   "id": "m02-l3-7",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "from scipy.stats import chi2_contingency\n",
    "age_group = (age > 60).astype(int)\n",
    "table = np.array([[np.sum((age_group == g) & (comorbidity == c)) for c in (0, 1)] for g in (0, 1)])\n",
    "chi2, pval, *_ = chi2_contingency(table)\n",
    "print(table)\n",
    "print(f\"Chi-square = {chi2:.1f}, p-value = {pval:.2g}  -> age and comorbidity are NOT independent\")"
   ]
  },
  {
   "id": "m02-l3-8",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Your first confounder"
   ]
  },
  {
   "id": "m02-l3-9",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** Does the drug reduce hospital days, and why does a simple comparison say the opposite?"
   ]
  },
  {
   "id": "m02-l3-10",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "TRUE_EFFECT = -2.0\n",
    "p_treat = 1 / (1 + np.exp(-(-5 + 0.11 * age)))         # older -> more likely treated\n",
    "drug = rng.binomial(1, p_treat)\n",
    "hospital_days = 2 + 0.3 * age + TRUE_EFFECT * drug + rng.normal(0, 2, N)\n",
    "\n",
    "d = pl.DataFrame({\"age\": age, \"drug\": drug, \"days\": hospital_days})\n",
    "naive = d.filter(pl.col(\"drug\") == 1)[\"days\"].mean() - d.filter(pl.col(\"drug\") == 0)[\"days\"].mean()\n",
    "print(f\"True effect: {TRUE_EFFECT:+.2f} days | Naive difference in means: {naive:+.2f} days\")"
   ]
  },
  {
   "id": "m02-l3-11",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "strata = (\n",
    "    d.with_columns((pl.col(\"age\") // 5 * 5).alias(\"age_band\"))\n",
    "    .group_by(\"age_band\")\n",
    "    .agg(\n",
    "        (pl.col(\"days\").filter(pl.col(\"drug\") == 1).mean()\n",
    "         - pl.col(\"days\").filter(pl.col(\"drug\") == 0).mean()).alias(\"diff\"),\n",
    "        pl.len().alias(\"n\"),\n",
    "    )\n",
    "    .drop_nans().drop_nulls()\n",
    ")\n",
    "stratified = (strata[\"diff\"] * strata[\"n\"]).sum() / strata[\"n\"].sum()\n",
    "print(f\"Stratified (age-adjusted) estimate: {stratified:+.2f} days\")"
   ]
  },
  {
   "id": "m02-l3-12",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise: Simpson's paradox"
   ]
  },
  {
   "id": "m02-l3-13",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ✍️ Exercise\n",
    "\n",
    "Build a dataset with two hospitals. Hospital A treats mostly severe cases and has a higher death rate\n",
    "overall, yet within each severity group it has a *lower* death rate than Hospital B. Show both the pooled\n",
    "and the within-severity comparisons."
   ]
  },
  {
   "id": "m02-l3-14",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your answer here"
   ]
  },
  {
   "id": "m02-l3-15",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist and recap"
   ]
  }
 ]
}