{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "cells": [
  {
   "id": "m04-l3-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 4.3 · Inference for Experiments: Permutation Tests, Clusters and Power\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- Run a randomisation (permutation) test that relies only on the design\n",
    "- Analyse a cluster-randomised experiment with clustered standard errors\n",
    "- Estimate the effect of incentives on learning HIV status (Thornton, 2008)\n",
    "- Compute statistical power and a minimum detectable effect by simulation\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": "m04-l3-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Randomisation inference"
   ]
  },
  {
   "id": "m04-l3-2",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q numpy pandas statsmodels matplotlib seaborn causaldata\n",
    "import numpy as np, pandas as pd\n",
    "import statsmodels.formula.api as smf\n",
    "import matplotlib.pyplot as plt, seaborn as sns\n",
    "rng = np.random.default_rng(8)\n",
    "sns.set_theme(style=\"whitegrid\")"
   ]
  },
  {
   "id": "m04-l3-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** With only 40 pupils, could a 4-point voucher effect be a fluke?"
   ]
  },
  {
   "id": "m04-l3-4",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "n = 40\n",
    "d = rng.permutation(np.repeat([0, 1], n // 2))\n",
    "y = 60 + 4 * d + rng.normal(0, 8, n)\n",
    "observed = y[d == 1].mean() - y[d == 0].mean()\n",
    "\n",
    "perm = np.array([\n",
    "    (lambda s: y[s == 1].mean() - y[s == 0].mean())(rng.permutation(d)) for _ in range(10_000)\n",
    "])\n",
    "p_perm = np.mean(np.abs(perm) >= abs(observed))\n",
    "p_ols = smf.ols(\"y ~ d\", data=pd.DataFrame({\"y\": y, \"d\": d})).fit(cov_type=\"HC1\").pvalues[\"d\"]\n",
    "print(f\"Estimate {observed:.2f} | permutation p = {p_perm:.3f} | OLS p = {p_ols:.3f}\")\n",
    "\n",
    "plt.hist(perm, bins=60, alpha=.7); plt.axvline(observed, color=\"red\", lw=2)\n",
    "plt.title(\"Distribution of estimates under the sharp null\"); plt.show()"
   ]
  },
  {
   "id": "m04-l3-5",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A cluster-randomised field experiment"
   ]
  },
  {
   "id": "m04-l3-6",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** Do small cash incentives make people more likely to learn their HIV status?"
   ]
  },
  {
   "id": "m04-l3-7",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "from causaldata import thornton_hiv\n",
    "hiv = thornton_hiv.load_pandas().data.dropna(subset=[\"got\", \"any\", \"tinc\", \"villnum\"])\n",
    "\n",
    "m_any = smf.ols(\"got ~ any\", data=hiv).fit(cov_type=\"cluster\", cov_kwds={\"groups\": hiv[\"villnum\"]})\n",
    "print(f\"Share learning status without incentive: {hiv.loc[hiv['any'] == 0, 'got'].mean():.3f}\")\n",
    "print(f\"Effect of any incentive: {m_any.params['any']:.3f} (clustered SE {m_any.bse['any']:.3f})\")\n",
    "\n",
    "m_amt = smf.ols(\"got ~ any + tinc + distvct\", data=hiv).fit(cov_type=\"cluster\", cov_kwds={\"groups\": hiv[\"villnum\"]})\n",
    "print(m_amt.params.round(4))"
   ]
  },
  {
   "id": "m04-l3-8",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Power and the minimum detectable effect"
   ]
  },
  {
   "id": "m04-l3-9",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** How many pupils do I need to detect a 2-point voucher effect with 80% power?"
   ]
  },
  {
   "id": "m04-l3-10",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def power(n, effect, sd=10, reps=600, alpha=0.05):\n",
    "    hits = 0\n",
    "    for _ in range(reps):\n",
    "        d = rng.permutation(np.repeat([0, 1], n // 2))\n",
    "        y = 60 + effect * d + rng.normal(0, sd, n)\n",
    "        hits += smf.ols(\"y ~ d\", data=pd.DataFrame({\"y\": y, \"d\": d})).fit().pvalues[\"d\"] < alpha\n",
    "    return hits / reps\n",
    "\n",
    "sizes = [100, 200, 400, 800, 1_600]\n",
    "curve = [power(n, effect=2.0) for n in sizes]\n",
    "for n, p in zip(sizes, curve):\n",
    "    print(f\"n = {n:>5}: power = {p:.2f}\")\n",
    "\n",
    "plt.plot(sizes, curve, marker=\"o\"); plt.axhline(0.8, ls=\"--\", color=\"grey\")\n",
    "plt.xscale(\"log\"); plt.xlabel(\"Total sample size\"); plt.ylabel(\"Power\"); plt.title(\"Power to detect +2 points\"); plt.show()"
   ]
  },
  {
   "id": "m04-l3-11",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "for n in [400, 1_600]:\n",
    "    print(f\"n = {n:>5}: minimum detectable effect ≈ {2.8 * 10 * np.sqrt(4 / n):.2f} points\")"
   ]
  },
  {
   "id": "m04-l3-12",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise"
   ]
  },
  {
   "id": "m04-l3-13",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ✍️ Exercise\n",
    "\n",
    "A ministry can afford 2,000 pupils but must randomise **whole schools** of 50 pupils (40 schools), with a\n",
    "school-level shock of SD 4 on top of pupil noise of SD 10. Simulate the power to detect a 2-point effect\n",
    "with clustered standard errors. Compare with 2,000 individually randomised pupils. What would you advise?"
   ]
  },
  {
   "id": "m04-l3-14",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your answer here"
   ]
  },
  {
   "id": "m04-l3-15",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist and recap"
   ]
  }
 ]
}