{
 "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-l2-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 4.2 · Randomised Experiments: the Voucher Lottery\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- Explain why random assignment removes selection bias\n",
    "- Estimate an ATE by difference in means and by regression\n",
    "- Check covariate balance and interpret a balance table\n",
    "- Use covariate adjustment to gain precision without introducing bias\n",
    "- Analyse a real experiment (LaLonde's National Supported Work programme)\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-l2-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Why a lottery solves the problem"
   ]
  },
  {
   "id": "m04-l2-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(12)\n",
    "sns.set_theme(style=\"whitegrid\")"
   ]
  },
  {
   "id": "m04-l2-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulating an oversubscribed voucher lottery"
   ]
  },
  {
   "id": "m04-l2-4",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** How much does winning a voucher lottery raise test scores?"
   ]
  },
  {
   "id": "m04-l2-5",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "N = 4_000\n",
    "applicants = pd.DataFrame({\n",
    "    \"motivation\": rng.normal(0, 1, N),\n",
    "    \"baseline_score\": rng.normal(60, 10, N),\n",
    "    \"low_income\": rng.binomial(1, 0.4, N),\n",
    "})\n",
    "applicants[\"won\"] = rng.permutation(np.repeat([0, 1], N // 2))          # the lottery\n",
    "TRUE_ATE = 4.0\n",
    "applicants[\"score\"] = (\n",
    "    10 + 0.8 * applicants.baseline_score + 3 * applicants.motivation\n",
    "    - 2 * applicants.low_income + TRUE_ATE * applicants.won + rng.normal(0, 6, N)\n",
    ")\n",
    "\n",
    "diff = applicants.groupby(\"won\")[\"score\"].mean()\n",
    "print(f\"Difference in means: {diff[1] - diff[0]:.2f} (true {TRUE_ATE})\")\n",
    "\n",
    "ols = smf.ols(\"score ~ won\", data=applicants).fit(cov_type=\"HC1\")\n",
    "print(ols.summary().tables[1])"
   ]
  },
  {
   "id": "m04-l2-6",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Balance checks"
   ]
  },
  {
   "id": "m04-l2-7",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** Did the lottery produce comparable winners and losers?"
   ]
  },
  {
   "id": "m04-l2-8",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def balance_table(df, treat, covariates):\n",
    "    rows = []\n",
    "    for c in covariates:\n",
    "        t, u = df.loc[df[treat] == 1, c], df.loc[df[treat] == 0, c]\n",
    "        std_diff = (t.mean() - u.mean()) / np.sqrt((t.var() + u.var()) / 2)\n",
    "        p = smf.ols(f\"{c} ~ {treat}\", data=df).fit(cov_type=\"HC1\").pvalues[treat]\n",
    "        rows.append((c, t.mean(), u.mean(), std_diff, p))\n",
    "    return pd.DataFrame(rows, columns=[\"covariate\", \"treated\", \"control\", \"std_diff\", \"p_value\"]).round(3)\n",
    "\n",
    "print(balance_table(applicants, \"won\", [\"motivation\", \"baseline_score\", \"low_income\"]))"
   ]
  },
  {
   "id": "m04-l2-9",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Covariate adjustment for precision"
   ]
  },
  {
   "id": "m04-l2-10",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** Can baseline test scores make the voucher estimate more precise?"
   ]
  },
  {
   "id": "m04-l2-11",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "a = applicants.copy()\n",
    "for c in [\"baseline_score\", \"low_income\"]:\n",
    "    a[c + \"_c\"] = a[c] - a[c].mean()\n",
    "adjusted = smf.ols(\"score ~ won * (baseline_score_c + low_income_c)\", data=a).fit(cov_type=\"HC1\")\n",
    "\n",
    "print(f\"Unadjusted: {ols.params['won']:.2f} (SE {ols.bse['won']:.3f})\")\n",
    "print(f\"Adjusted:   {adjusted.params['won']:.2f} (SE {adjusted.bse['won']:.3f})\")\n",
    "print(f\"Equivalent to a sample {((ols.bse['won'] / adjusted.bse['won'])**2):.1f}x larger\")"
   ]
  },
  {
   "id": "m04-l2-12",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Real data: the NSW job-training experiment"
   ]
  },
  {
   "id": "m04-l2-13",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "from causaldata import nsw_mixtape\n",
    "nsw = nsw_mixtape.load_pandas().data\n",
    "print(balance_table(nsw, \"treat\", [\"age\", \"educ\", \"black\", \"hisp\", \"marr\", \"nodegree\", \"re74\", \"re75\"]))\n",
    "\n",
    "nsw_ate = smf.ols(\"re78 ~ treat\", data=nsw).fit(cov_type=\"HC1\")\n",
    "nsw_adj = smf.ols(\"re78 ~ treat + age + educ + black + hisp + marr + nodegree + re74 + re75\", data=nsw).fit(cov_type=\"HC1\")\n",
    "print(f\"Experimental ATE: ${nsw_ate.params['treat']:,.0f} (SE {nsw_ate.bse['treat']:,.0f})\")\n",
    "print(f\"Adjusted:         ${nsw_adj.params['treat']:,.0f} (SE {nsw_adj.bse['treat']:,.0f})\")"
   ]
  },
  {
   "id": "m04-l2-14",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ✍️ Exercise\n",
    "\n",
    "Simulate a *broken* lottery: let families with high motivation be twice as likely to win. Re-run the\n",
    "balance table and the difference in means. Which line of the balance table warns you?"
   ]
  },
  {
   "id": "m04-l2-15",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your answer here"
   ]
  },
  {
   "id": "m04-l2-16",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist and recap"
   ]
  }
 ]
}