{
 "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-l1-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 4.1 · Potential Outcomes: ATE, ATT, ATU and Selection Bias\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- Write each unit's two potential outcomes and the fundamental problem of causal inference\n",
    "- Define and compute the ATE, ATT and ATU in a simulation where both outcomes are visible\n",
    "- Decompose a naive difference in means into the ATT plus selection bias\n",
    "- State SUTVA and give examples of when it fails\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-l1-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Two worlds per person"
   ]
  },
  {
   "id": "m04-l1-2",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q numpy pandas statsmodels matplotlib seaborn\n",
    "import numpy as np, pandas as pd\n",
    "import matplotlib.pyplot as plt, seaborn as sns\n",
    "rng = np.random.default_rng(4)\n",
    "sns.set_theme(style=\"whitegrid\")"
   ]
  },
  {
   "id": "m04-l1-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A voucher world where we can see both outcomes"
   ]
  },
  {
   "id": "m04-l1-4",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** What is the average effect of a private-school voucher on test scores, for everyone and for users?"
   ]
  },
  {
   "id": "m04-l1-5",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "N = 10_000\n",
    "motivation = rng.normal(0, 1, N)\n",
    "y0 = 60 + 6 * motivation + rng.normal(0, 5, N)          # score without a voucher\n",
    "effect = 5 - 2 * motivation                              # heterogeneous effect\n",
    "y1 = y0 + effect                                         # score with a voucher\n",
    "\n",
    "# Self-selection: motivated families are more likely to use a voucher\n",
    "d = rng.binomial(1, 1 / (1 + np.exp(-1.5 * motivation)))\n",
    "y = np.where(d == 1, y1, y0)                             # what we actually observe\n",
    "\n",
    "pupils = pd.DataFrame({\"motivation\": motivation, \"y0\": y0, \"y1\": y1, \"d\": d, \"y\": y})\n",
    "ATE = (pupils.y1 - pupils.y0).mean()\n",
    "ATT = (pupils.y1 - pupils.y0)[pupils.d == 1].mean()\n",
    "ATU = (pupils.y1 - pupils.y0)[pupils.d == 0].mean()\n",
    "print(f\"ATE = {ATE:.2f} | ATT = {ATT:.2f} | ATU = {ATU:.2f}\")\n",
    "print(pupils.head())"
   ]
  },
  {
   "id": "m04-l1-6",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Naive comparison = effect + selection bias"
   ]
  },
  {
   "id": "m04-l1-7",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** How much of the raw score gap between voucher users and non-users is the voucher, and how much is who chose it?"
   ]
  },
  {
   "id": "m04-l1-8",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "naive = pupils.y[pupils.d == 1].mean() - pupils.y[pupils.d == 0].mean()\n",
    "selection_bias = pupils.y0[pupils.d == 1].mean() - pupils.y0[pupils.d == 0].mean()\n",
    "print(f\"Naive difference : {naive:6.2f}\")\n",
    "print(f\"ATT              : {ATT:6.2f}\")\n",
    "print(f\"Selection bias   : {selection_bias:6.2f}\")\n",
    "print(f\"ATT + bias       : {ATT + selection_bias:6.2f}\")\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(8, 4))\n",
    "sns.kdeplot(data=pupils, x=\"y0\", hue=\"d\", common_norm=False, ax=ax)\n",
    "ax.set(title=\"Y(0) differs between users and non-users: selection bias\", xlabel=\"Score without voucher\")\n",
    "plt.show()"
   ]
  },
  {
   "id": "m04-l1-9",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SUTVA: the assumption hiding in the notation"
   ]
  },
  {
   "id": "m04-l1-10",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise"
   ]
  },
  {
   "id": "m04-l1-11",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ✍️ Exercise\n",
    "\n",
    "Change the DGP so that the voucher effect is the **same** for everyone (`effect = 5`). Recompute the ATE,\n",
    "ATT, ATU and the selection bias. Which quantities change, which do not, and is the naive estimate still\n",
    "biased?"
   ]
  },
  {
   "id": "m04-l1-12",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your answer here"
   ]
  },
  {
   "id": "m04-l1-13",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist and recap"
   ]
  }
 ]
}