{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "cells": [
  {
   "id": "m13-l2-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 13.2 · Cluster Randomisation, Partial Interference and Saturation Designs\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- Explain partial interference: spillovers within clusters but not between them\n",
    "- Analyse a randomised saturation design with pure-control clusters\n",
    "- Estimate the direct effect, the spillover on the untreated, and the total effect at each saturation\n",
    "- Choose between individual and cluster randomisation for a pilot\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": "m13-l2-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Partial interference"
   ]
  },
  {
   "id": "m13-l2-2",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q numpy pandas statsmodels matplotlib\n",
    "import numpy as np, pandas as pd\n",
    "import statsmodels.formula.api as smf\n",
    "import matplotlib.pyplot as plt\n",
    "rng = np.random.default_rng(13)"
   ]
  },
  {
   "id": "m13-l2-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A two-stage randomised vaccination design"
   ]
  },
  {
   "id": "m13-l2-4",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** How much does vaccination protect the vaccinated, and how much does village coverage protect the unvaccinated?"
   ]
  },
  {
   "id": "m13-l2-5",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "V, M = 120, 50\n",
    "saturation = rng.permutation(np.repeat([0.0, 0.3, 0.7], V // 3))\n",
    "rows = []\n",
    "for v in range(V):\n",
    "    s = saturation[v]\n",
    "    d = rng.permutation(np.r_[np.ones(int(s * M)), np.zeros(M - int(s * M))])\n",
    "    village_risk = rng.normal(0, 0.03)\n",
    "    for i in range(M):\n",
    "        p = 0.30 + village_risk - 0.15 * d[i] - 0.20 * s       # direct -0.15; herd effect -0.20 x coverage\n",
    "        rows.append((v, s, int(d[i]), rng.binomial(1, np.clip(p, 0.01, 0.99))))\n",
    "df = pd.DataFrame(rows, columns=[\"village\", \"saturation\", \"vaccinated\", \"infected\"])\n",
    "print(df.groupby([\"saturation\", \"vaccinated\"]).infected.agg([\"mean\", \"size\"]).round(3))"
   ]
  },
  {
   "id": "m13-l2-6",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Direct, spillover and total effects"
   ]
  },
  {
   "id": "m13-l2-7",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** What would a naive within-village comparison miss?"
   ]
  },
  {
   "id": "m13-l2-8",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def fit(formula, data):\n",
    "    return smf.ols(formula, data=data).fit(cov_type=\"cluster\", cov_kwds={\"groups\": data[\"village\"]})\n",
    "\n",
    "df[\"sat30\"] = (df.saturation == 0.3).astype(int)\n",
    "df[\"sat70\"] = (df.saturation == 0.7).astype(int)\n",
    "m = fit(\"infected ~ vaccinated + sat30 + sat70 + vaccinated:sat70\", df)\n",
    "print(m.params.round(3))\n",
    "\n",
    "pure = df[df.saturation == 0].infected.mean()\n",
    "for s in (0.3, 0.7):\n",
    "    sub = df[df.saturation == s]\n",
    "    spill = sub[sub.vaccinated == 0].infected.mean() - pure\n",
    "    direct = sub[sub.vaccinated == 1].infected.mean() - sub[sub.vaccinated == 0].infected.mean()\n",
    "    total = sub[sub.vaccinated == 1].infected.mean() - pure\n",
    "    print(f\"saturation {s:.0%}: spillover on unvaccinated {spill:+.3f} | direct {direct:+.3f} | total on vaccinated {total:+.3f}\")"
   ]
  },
  {
   "id": "m13-l2-9",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Design choices for your pilot"
   ]
  },
  {
   "id": "m13-l2-10",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ✍️ Exercise\n",
    "\n",
    "Simulate the same study with only 30 villages (10 per arm). Repeat it 200 times and record how often the\n",
    "spillover at 70% coverage is significant at 5%. What does that say about the number of clusters such a\n",
    "design needs?"
   ]
  },
  {
   "id": "m13-l2-11",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your answer here"
   ]
  },
  {
   "id": "m13-l2-12",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist and recap"
   ]
  }
 ]
}