{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "provenance": []
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Day 38: Stratified & Blocked Experiments\n",
    "\n",
    "*Beginner to Ninja: Math for Causal Inference*\n",
    "\n",
    "Randomising within blocks of similar units removes chance imbalance on the blocking variables and can sharply reduce variance.\n",
    "\n",
    "**Learning outcomes**\n",
    "\n",
    "- Compute the stratified (blocked) estimator and its variance.\n",
    "- Prove that it is unbiased and explain when blocking reduces variance.\n",
    "- Show why pooling across blocks with different treatment shares is biased, and what block fixed effects estimate.\n",
    "- Analyse matched-pair designs and choose blocking variables in applied settings."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Theory\n",
    "\n",
    "### The design\n",
    "\n",
    "Partition the $N$ units into blocks $b = 1, \\dots, B$ of sizes $N_b$ (regions, baseline-score bands, pairs of similar villages). Within each block, completely randomise $n_{1b}$ units to treatment. Assignments in different blocks are independent.\n",
    "\n",
    "### Estimation\n",
    "\n",
    "Estimate the effect within each block, then average with block shares as weights:\n",
    "\n",
    "$$\n",
    "\\hat\\tau_{\\text{strat}} = \\sum_{b=1}^B\\frac{N_b}{N}\\,\\hat\\tau_b, \\qquad \\hat\\tau_b = \\bar Y_{1b} - \\bar Y_{0b},\n",
    "$$\n",
    "\n",
    "$$\n",
    "\\operatorname{Var}(\\hat\\tau_{\\text{strat}}) = \\sum_b\\Big(\\frac{N_b}{N}\\Big)^2\\operatorname{Var}(\\hat\\tau_b), \\qquad \\hat V = \\sum_b\\Big(\\frac{N_b}{N}\\Big)^2\\Big(\\frac{s_{1b}^2}{n_{1b}} + \\frac{s_{0b}^2}{n_{0b}}\\Big).\n",
    "$$\n",
    "\n",
    "### Why it helps\n",
    "\n",
    "Under complete randomisation, the variance of $\\hat\\tau$ reflects *all* outcome variation. Blocking removes the **between-block** part, because every block contributes treated and control units in fixed proportions. The gain is large when blocks predict the outcome (baseline test scores for test scores, baseline earnings for earnings), and small otherwise.\n",
    "\n",
    "### Pitfalls\n",
    "\n",
    "- **Unequal treatment shares across blocks:** a pooled difference in means is then biased (it mixes block differences into the comparison). Use $\\hat\\tau_{\\text{strat}}$.\n",
    "- **Block fixed-effects regression** ($Y$ on $T$ and block dummies) weights blocks by $N_bp_b(1 - p_b)$ (Day 28), not $N_b$; it equals $\\hat\\tau_{\\text{strat}}$ only when $p_b$ is the same in every block.\n",
    "- **Matched pairs** (blocks of size 2) cannot estimate within-pair variances; use the variance of pair differences, which is conservative.\n",
    "\n",
    "## Part 1 · Warm-up"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 1 · The stratified estimate\n",
    "\n",
    "A tutoring trial has two blocks: low baseline scores (60 pupils, estimated effect 4 points) and high (40 pupils, effect 1 point). Compute $\\hat\\tau_{\\text{strat}}$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Weights: $60/100 = 0.6$ and $0.4$.\n",
    "\n",
    "**Step 2.** $\\hat\\tau_{\\text{strat}} = 0.6(4) + 0.4(1) = 2.8$.\n",
    "\n",
    "**Answer:** 2.8 points.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 2 · Its variance\n",
    "\n",
    "The within-block variances of $\\hat\\tau_b$ are 1.0 (low) and 0.5 (high). Compute the SE of $\\hat\\tau_{\\text{strat}}$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** $\\operatorname{Var} = 0.6^2(1.0) + 0.4^2(0.5) = 0.36 + 0.08 = 0.44$.\n",
    "\n",
    "**Step 2.** SE $= \\sqrt{0.44} \\approx 0.66$.\n",
    "\n",
    "**Answer:** 0.66 points.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 3 · Matched pairs\n",
    "\n",
    "Five pairs of similar villages each have one treated village. Treated-minus-control differences in household income are 3, 1, 4, 0, 2 (\\$100s). Estimate the effect and its SE.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Mean difference: $(3 + 1 + 4 + 0 + 2)/5 = 2$.\n",
    "\n",
    "**Step 2.** Deviations: $1, -1, 2, -2, 0$; sample variance $= (1 + 1 + 4 + 4 + 0)/4 = 2.5$; sd $\\approx 1.58$.\n",
    "\n",
    "**Step 3.** SE $= 1.58/\\sqrt5 \\approx 0.71$.\n",
    "\n",
    "**Answer:** \\$200 per household (SE ≈ \\$71).\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 4 · Fewer possible assignments\n",
    "\n",
    "20 schools are split into two blocks of 10, with 5 treated per block. How many assignments are possible, versus complete randomisation of 10 of 20?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Blocked: $\\binom{10}{5}^2 = 252^2 = 63{,}504$.\n",
    "\n",
    "**Step 2.** Complete: $\\binom{20}{10} = 184{,}756$.\n",
    "\n",
    "**Answer:** blocking rules out the 121,252 assignments with unbalanced blocks (e.g. all treated schools in one block).\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 5 · Different shares per block\n",
    "\n",
    "In block A half the units are treated; in block B 20% are. Why can't you simply compare all treated with all controls?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** The treated group then over-represents block A (more of its units are treated).\n",
    "\n",
    "**Step 2.** If outcomes differ between blocks, the treated–control comparison partly compares A with B.\n",
    "\n",
    "**Answer:** the pooled difference is biased; estimate within blocks and average (Exercise 12 quantifies the bias).\n",
    "\n",
    "</details>\n",
    "\n",
    "## Part 2 · Proofs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 6 · Unbiasedness\n",
    "\n",
    "Prove that $\\hat\\tau_{\\text{strat}}$ is unbiased for the finite-population ATE.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Within each block, randomisation is complete, so $E[\\hat\\tau_b] = \\tau_b$, the block's ATE (Day 31, Exercise 8).\n",
    "\n",
    "**Step 2.** Linearity: $E[\\hat\\tau_{\\text{strat}}] = \\sum_b\\frac{N_b}{N}\\tau_b$.\n",
    "\n",
    "**Step 3.** $\\tau_b = \\frac{1}{N_b}\\sum_{i \\in b}\\tau_i$, so $\\sum_b\\frac{N_b}{N}\\tau_b = \\frac1N\\sum_i\\tau_i = \\tau_{fp}$. $\\blacksquare$\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 7 · The variance formula\n",
    "\n",
    "Prove $\\operatorname{Var}(\\hat\\tau_{\\text{strat}}) = \\sum_b(N_b/N)^2\\operatorname{Var}(\\hat\\tau_b)$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Assignments in different blocks are independent, so $\\hat\\tau_b$ and $\\hat\\tau_{b'}$ are independent for $b \\ne b'$.\n",
    "\n",
    "**Step 2.** The variance of a weighted sum of independent terms is the sum of squared weights times variances (Day 22): $\\sum_b(N_b/N)^2\\operatorname{Var}(\\hat\\tau_b)$. $\\blacksquare$\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 8 · Where the gain comes from\n",
    "\n",
    "Assume constant effects, two equal blocks of size $N/2$, half treated in each, and $Y_i(0) = \\mu_b + \\varepsilon_i$ with within-block variance $\\sigma_w^2$ and block means $\\mu_1, \\mu_2$. Compare the (approximate) variances under complete randomisation and blocking.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Total variance of $Y(0)$: $\\sigma_w^2 + \\sigma_B^2$, with between-block part $\\sigma_B^2 = (\\mu_1 - \\mu_2)^2/4$.\n",
    "\n",
    "**Step 2.** Complete randomisation (Day 36, constant effects, large $N$): $\\operatorname{Var} \\approx (\\sigma_w^2 + \\sigma_B^2)\\big(\\frac{2}{N} + \\frac{2}{N}\\big) = \\frac4N(\\sigma_w^2 + \\sigma_B^2)$.\n",
    "\n",
    "**Step 3.** Blocking: each block has $N/4$ per arm and variance $\\sigma_w^2$, so $\\operatorname{Var}(\\hat\\tau_b) \\approx \\frac{8}{N}\\sigma_w^2$, and $\\operatorname{Var}(\\hat\\tau_{\\text{strat}}) \\approx 2 \\times \\frac14 \\times \\frac8N\\sigma_w^2 = \\frac4N\\sigma_w^2$.\n",
    "\n",
    "**Answer:** blocking removes the $\\sigma_B^2$ term: the larger the outcome gap between blocks, the larger the gain.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 9 · Bias of pooling with unequal shares\n",
    "\n",
    "With constant effect $\\tau$, block shares $q_b = N_b/N$, treatment shares $p_b$ and block control means $\\mu_b$, show the pooled difference in means equals $\\tau + \\sum_b(w_{1b} - w_{0b})\\mu_b$, where $w_{1b}$ and $w_{0b}$ are the shares of block $b$ among all treated and all controls.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Treated mean (in expectation): $\\sum_bw_{1b}(\\mu_b + \\tau) = \\tau + \\sum_bw_{1b}\\mu_b$, with $w_{1b} = q_bp_b/\\sum_cq_cp_c$.\n",
    "\n",
    "**Step 2.** Control mean: $\\sum_bw_{0b}\\mu_b$, with $w_{0b} = q_b(1 - p_b)/\\sum_cq_c(1 - p_c)$.\n",
    "\n",
    "**Step 3.** Difference: $\\tau + \\sum_b(w_{1b} - w_{0b})\\mu_b$. $\\blacksquare$\n",
    "\n",
    "The bias vanishes if all $p_b$ are equal (then $w_{1b} = w_{0b} = q_b$) or all $\\mu_b$ are equal.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 10 · What block fixed effects estimate\n",
    "\n",
    "Using Day 28's result, explain why regressing $Y$ on $T$ and block dummies gives $\\sum_b\\omega_b\\hat\\tau_b$ with $\\omega_b \\propto N_bp_b(1 - p_b)$, and when this equals $\\hat\\tau_{\\text{strat}}$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** By FWL, the coefficient on $T$ uses $T$ demeaned within blocks; its within-block variance is $p_b(1 - p_b)$.\n",
    "\n",
    "**Step 2.** Within a block, the regression slope on the demeaned treatment is $\\hat\\tau_b$; the overall coefficient weights each block by its share of the residual treatment variation, $N_bp_b(1 - p_b)$.\n",
    "\n",
    "**Step 3.** If $p_b$ is the same in all blocks, $\\omega_b \\propto N_b$: the same as $\\hat\\tau_{\\text{strat}}$.\n",
    "\n",
    "**Answer:** with equal treatment shares, FE regression and the stratified estimator coincide; otherwise FE over-weights blocks with shares near 50%, which matters when effects differ across blocks.\n",
    "\n",
    "</details>\n",
    "\n",
    "## Part 3 · Applied"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 11 · How much does blocking buy?\n",
    "\n",
    "In a tutoring trial with 400 pupils, baseline score terciles explain 60% of the variance of end-line scores. With constant effects and equal shares, by roughly how much does blocking on terciles cut the SE, and what sample size would complete randomisation need to match it?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Blocking removes the between-tercile variance: 60% of the total (Exercise 8).\n",
    "\n",
    "**Step 2.** Variance falls to 40%, so the SE falls by $\\sqrt{0.4} \\approx 0.63$.\n",
    "\n",
    "**Step 3.** Complete randomisation would need $400/0.4 = 1{,}000$ pupils for the same precision.\n",
    "\n",
    "**Answer:** about 37% smaller SEs, equivalent to 600 extra pupils.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 12 · A cash transfer with regional shares\n",
    "\n",
    "North: 1,000 households, 50% treated, control mean consumption \\$100. South: 1,000 households, 20% treated, control mean \\$60. The effect is \\$10 everywhere. Compute the pooled difference in means and the stratified estimate.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Treated: 500 in the North (mean 110) and 200 in the South (mean 70). Treated mean $= (500 \\times 110 + 200 \\times 70)/700 = 69{,}000/700 \\approx 98.6$.\n",
    "\n",
    "**Step 2.** Controls: 500 North (mean 100) and 800 South (mean 60). Control mean $= (50{,}000 + 48{,}000)/1{,}300 \\approx 75.4$.\n",
    "\n",
    "**Step 3.** Pooled difference $\\approx 98.6 - 75.4 = 23.2$.\n",
    "\n",
    "**Step 4.** Stratified: both block effects are 10, so $\\hat\\tau_{\\text{strat}} = 0.5(10) + 0.5(10) = 10$.\n",
    "\n",
    "**Answer:** pooling reports \\$23 instead of \\$10, because the treated are disproportionately from the richer North.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 13 · Matched-pair villages\n",
    "\n",
    "A water-sanitation trial pairs 30 villages into 15 pairs by baseline diarrhoea rates and randomises one per pair. Why is the pair-difference SE conservative, and why is the design still attractive?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Each pair has one treated and one control village, so the within-pair variance of each arm cannot be estimated.\n",
    "\n",
    "**Step 2.** The variance of pair differences also includes genuine differences in effects *between* pairs, which inflates it: conservative (Imai, 2008).\n",
    "\n",
    "**Step 3.** Pairing on baseline diarrhoea removes most outcome variation not due to treatment, so the true variance is much smaller than under complete randomisation.\n",
    "\n",
    "**Answer:** slightly conservative SEs, but a much more precise design.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 14 · Blocks that are too small\n",
    "\n",
    "An analyst creates 50 blocks of 2 pupils each and analyses with $\\hat V = \\sum_b(N_b/N)^2(s_{1b}^2/n_{1b} + s_{0b}^2/n_{0b})$. What goes wrong, and what should be used?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** With one treated and one control per block, $s_{1b}^2$ and $s_{0b}^2$ need at least 2 observations each: they are undefined.\n",
    "\n",
    "**Step 2.** Use the matched-pair variance: the sample variance of the 50 pair differences divided by 50.\n",
    "\n",
    "**Answer:** treat it as a matched-pair design; alternatively, form blocks of at least 4 units with 2 per arm so within-block variances can be estimated.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 15 · Choosing blocking variables\n",
    "\n",
    "A conditional cash transfer evaluation can block on: (a) baseline school enrolment, (b) region, (c) household size, (d) whether the child enrols after the programme starts. Rank them and explain.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** (a) Baseline enrolment predicts the outcome strongly: best.\n",
    "\n",
    "**Step 2.** (b) Region often predicts outcomes and matters for implementation and subgroup analysis: good.\n",
    "\n",
    "**Step 3.** (c) Household size predicts weakly: helpful only if blocks are cheap to form.\n",
    "\n",
    "**Step 4.** (d) Post-treatment: never; it can be affected by the treatment (Day 31, Exercise 10).\n",
    "\n",
    "**Answer:** (a) > (b) > (c); (d) is invalid.\n",
    "\n",
    "</details>\n",
    "\n",
    "## Verify in Python"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "rng = np.random.default_rng(38)\n",
    "N = 400\n",
    "block = np.repeat([0, 1, 2], [134, 133, 133])          # baseline-score terciles\n",
    "y0 = np.array([40.0, 55.0, 70.0])[block] + rng.normal(0, 8, N)\n",
    "y1 = y0 + 5                                            # constant effect\n",
    "cr, bl = [], []\n",
    "for _ in range(10_000):\n",
    "    T = np.zeros(N, bool); T[rng.choice(N, N // 2, replace=False)] = True       # complete\n",
    "    cr.append(y1[T].mean() - y0[~T].mean())\n",
    "    T = np.zeros(N, bool)\n",
    "    for b in range(3):                                                           # blocked\n",
    "        idx = np.where(block == b)[0]; T[rng.choice(idx, len(idx) // 2, replace=False)] = True\n",
    "    est = sum((block == b).mean() * (y1[T & (block == b)].mean() - y0[~T & (block == b)].mean()) for b in range(3))\n",
    "    bl.append(est)\n",
    "print(f\"SD complete {np.std(cr):.3f}, SD blocked {np.std(bl):.3f}, ratio {np.std(bl) / np.std(cr):.2f}\")\n",
    "\n",
    "# Pooling with unequal shares (Exercise 12)\n",
    "treated_mean = (500 * 110 + 200 * 70) / 700; control_mean = (500 * 100 + 800 * 60) / 1300\n",
    "print(f\"pooled {treated_mean - control_mean:.1f} versus stratified 10\")\n",
    "d = np.array([3, 1, 4, 0, 2.]); print(\"pairs: estimate\", d.mean(), \"SE\", round(d.std(ddof=1) / np.sqrt(5), 3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Key takeaways\n",
    "\n",
    "- Randomising within blocks removes chance imbalance on the blocking variables and the between-block variance.\n",
    "- Estimate within blocks and weight by block size; pooling or FE regression mislead when treatment shares differ.\n",
    "- Block on strong pre-treatment predictors of the outcome; pairs are the extreme case."
   ]
  }
 ]
}