{
 "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-l2-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 2.2 · Law of Large Numbers, Sampling Distributions and Confidence Intervals\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- Show by simulation that sample means converge to the population mean (LLN)\n",
    "- Build a sampling distribution and read the standard error off it\n",
    "- See the Central Limit Theorem at work on a skewed variable\n",
    "- Explain what \"95% confidence\" means by checking the coverage of many intervals\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-l2-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup: the patient population"
   ]
  },
  {
   "id": "m02-l2-2",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q numpy scipy polars matplotlib seaborn\n",
    "import numpy as np, scipy.stats as stats, polars as pl\n",
    "import matplotlib.pyplot as plt, seaborn as sns\n",
    "\n",
    "rng = np.random.default_rng(101)\n",
    "sns.set_theme(style=\"whitegrid\")\n",
    "N = 50_000\n",
    "population = pl.DataFrame({\n",
    "    \"age\": np.clip(rng.normal(45, 12, N), 18, 90),\n",
    "    \"cost\": rng.lognormal(7.5, 1.0, N),          # annual health-care cost, very skewed\n",
    "})\n",
    "TRUE_MEAN_AGE = population[\"age\"].mean()          # the population value we try to recover\n",
    "print(f\"True mean age: {TRUE_MEAN_AGE:.3f}\")"
   ]
  },
  {
   "id": "m02-l2-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Law of Large Numbers"
   ]
  },
  {
   "id": "m02-l2-4",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** How close does the sample mean age get to the truth as the sample grows?"
   ]
  },
  {
   "id": "m02-l2-5",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "sample_sizes = [10, 30, 100, 300, 1_000, 3_000, 10_000]\n",
    "means = [population.sample(n=n, seed=i)[\"age\"].mean() for i, n in enumerate(sample_sizes)]\n",
    "\n",
    "for n, m in zip(sample_sizes, means):\n",
    "    print(f\"N = {n:>6,}: mean age = {m:6.2f}   (error {m - TRUE_MEAN_AGE:+.2f})\")\n",
    "\n",
    "plt.figure(figsize=(9, 4))\n",
    "plt.plot(sample_sizes, means, marker=\"o\", label=\"Sample mean\")\n",
    "plt.axhline(TRUE_MEAN_AGE, color=\"red\", ls=\"--\", label=\"Population mean\")\n",
    "plt.xscale(\"log\"); plt.xlabel(\"Sample size (log scale)\"); plt.ylabel(\"Mean age\")\n",
    "plt.title(\"Law of Large Numbers\"); plt.legend(); plt.show()"
   ]
  },
  {
   "id": "m02-l2-6",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sampling distributions and the standard error"
   ]
  },
  {
   "id": "m02-l2-7",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** If 1,000 survey teams each sampled 100 patients, how much would their averages disagree?"
   ]
  },
  {
   "id": "m02-l2-8",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "ages = population[\"age\"].to_numpy()\n",
    "\n",
    "def sampling_distribution(n, reps=2_000):\n",
    "    return np.array([rng.choice(ages, size=n, replace=False).mean() for _ in range(reps)])\n",
    "\n",
    "for n in [25, 100, 400]:\n",
    "    est = sampling_distribution(n)\n",
    "    formula_se = ages.std() / np.sqrt(n)\n",
    "    print(f\"n={n:>3}: simulated SE = {est.std():.3f} | formula sigma/sqrt(n) = {formula_se:.3f}\")"
   ]
  },
  {
   "id": "m02-l2-9",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Central Limit Theorem"
   ]
  },
  {
   "id": "m02-l2-10",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "costs = population[\"cost\"].to_numpy()\n",
    "fig, axes = plt.subplots(1, 3, figsize=(14, 4))\n",
    "sns.histplot(costs, bins=60, ax=axes[0]).set(title=\"Individual costs (skewed)\")\n",
    "for ax, n in zip(axes[1:], [5, 200]):\n",
    "    m = [rng.choice(costs, n).mean() for _ in range(3_000)]\n",
    "    sns.histplot(m, bins=50, ax=ax).set(title=f\"Means of samples of n={n}\")\n",
    "plt.tight_layout(); plt.show()"
   ]
  },
  {
   "id": "m02-l2-11",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Confidence intervals and coverage"
   ]
  },
  {
   "id": "m02-l2-12",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** Do 95% confidence intervals really contain the true mean 95% of the time?"
   ]
  },
  {
   "id": "m02-l2-13",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def coverage(n, reps=2_000, use_t=False):\n",
    "    hits = 0\n",
    "    for _ in range(reps):\n",
    "        s = rng.choice(ages, n, replace=False)\n",
    "        se = s.std(ddof=1) / np.sqrt(n)\n",
    "        crit = stats.t.ppf(0.975, df=n - 1) if use_t else 1.96\n",
    "        hits += (s.mean() - crit * se <= TRUE_MEAN_AGE <= s.mean() + crit * se)\n",
    "    return hits / reps\n",
    "\n",
    "for n in [5, 30, 200]:\n",
    "    print(f\"n={n:>3}: coverage (z) = {coverage(n):.3f} | coverage (t) = {coverage(n, use_t=True):.3f}\")"
   ]
  },
  {
   "id": "m02-l2-14",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ✍️ Exercise\n",
    "\n",
    "Repeat the coverage check for the **mean cost** (the skewed variable) with n = 10, 50 and 500. Where does\n",
    "the CLT approximation break down, and what does that imply for evaluating a programme on health spending\n",
    "with a small sample?"
   ]
  },
  {
   "id": "m02-l2-15",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your answer here"
   ]
  },
  {
   "id": "m02-l2-16",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist and recap"
   ]
  }
 ]
}