{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "cells": [
  {
   "id": "m14-l1-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 14.1 · Partial Identification: Manski Bounds and Lee Bounds\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- Explain partial identification: reporting the set of effects consistent with the data and weak assumptions\n",
    "- Compute worst-case (Manski) bounds for a job-training programme with 30% attrition\n",
    "- Tighten bounds with monotonicity assumptions, and see what each assumption buys\n",
    "- Compute Lee (trimming) bounds for differential attrition in an experiment\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": "m14-l1-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## From point estimates to honest ranges"
   ]
  },
  {
   "id": "m14-l1-2",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q numpy pandas matplotlib\n",
    "import numpy as np, pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "rng = np.random.default_rng(14)"
   ]
  },
  {
   "id": "m14-l1-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Worst-case bounds with missing outcomes"
   ]
  },
  {
   "id": "m14-l1-4",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** What can we say about the programme's effect on employment when 30% of outcomes are missing?"
   ]
  },
  {
   "id": "m14-l1-5",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "N = 2_000\n",
    "treat = rng.binomial(1, 0.5, N)\n",
    "employed = rng.binomial(1, np.where(treat == 1, 0.6, 0.4))\n",
    "p_missing = 0.2 + 0.2 * employed                       # employed people are harder to reach\n",
    "missing = rng.binomial(1, p_missing).astype(bool)\n",
    "y_obs = np.where(missing, np.nan, employed)\n",
    "df = pd.DataFrame({\"treat\": treat, \"y\": y_obs, \"missing\": missing})\n",
    "\n",
    "def group(t):\n",
    "    g = df[df.treat == t]\n",
    "    return g.y.mean(), g.missing.mean()                 # mean among observed, share missing\n",
    "\n",
    "(m1, q1), (m0, q0) = group(1), group(0)\n",
    "naive = m1 - m0\n",
    "# Bounds on each potential-outcome mean: missing outcomes could all be 0 or all be 1\n",
    "lo1, hi1 = m1 * (1 - q1), m1 * (1 - q1) + q1\n",
    "lo0, hi0 = m0 * (1 - q0), m0 * (1 - q0) + q0\n",
    "manski = (lo1 - hi0, hi1 - lo0)\n",
    "print(f\"Missing: treated {q1:.1%}, control {q0:.1%}\")\n",
    "print(f\"Naive (complete cases): {naive:+.3f} | truth +0.200\")\n",
    "print(f\"Manski bounds on the ATE: [{manski[0]:+.3f}, {manski[1]:+.3f}]\")"
   ]
  },
  {
   "id": "m14-l1-6",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Buying tighter bounds with assumptions"
   ]
  },
  {
   "id": "m14-l1-7",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Assumption: missing people are at least as employed as observed people in the same arm\n",
    "lo1_b, lo0_b = m1, m0                                   # lower bounds rise to the observed means\n",
    "bounds_b = (lo1_b - hi0, hi1 - lo0_b)\n",
    "bounds_mtr = (max(0, manski[0]), manski[1])\n",
    "table = pd.DataFrame({\n",
    "    \"lower\": [manski[0], bounds_mtr[0], max(0, bounds_b[0])],\n",
    "    \"upper\": [manski[1], bounds_mtr[1], bounds_b[1]],\n",
    "}, index=[\"Worst case (no assumptions)\", \"+ monotone response (ATE ≥ 0)\", \"+ missing ≥ observed, and MTR\"]).round(3)\n",
    "table[\"width\"] = table.upper - table.lower\n",
    "print(table)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(8, 3))\n",
    "for i, (name, r) in enumerate(table.iterrows()):\n",
    "    ax.plot([r.lower, r.upper], [i, i], lw=8, solid_capstyle=\"butt\")\n",
    "ax.axvline(0.2, color=\"red\", ls=\"--\", label=\"Truth\"); ax.axvline(0, color=\"grey\")\n",
    "ax.set_yticks(range(len(table))); ax.set_yticklabels(table.index); ax.legend(); ax.set_title(\"What each assumption buys\")\n",
    "plt.show()"
   ]
  },
  {
   "id": "m14-l1-8",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Lee bounds for attrition in experiments"
   ]
  },
  {
   "id": "m14-l1-9",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** What is the effect of training on wages when training also changes who we observe with a wage?"
   ]
  },
  {
   "id": "m14-l1-10",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "N = 4_000\n",
    "treat = rng.binomial(1, 0.5, N)\n",
    "ability = rng.normal(size=N)\n",
    "works = rng.binomial(1, 1 / (1 + np.exp(-(0.0 + 0.8 * treat + ability))))   # training raises employment\n",
    "wage = 20 + 2 * treat + 3 * ability + rng.normal(0, 2, N)                    # true wage effect = 2\n",
    "wage_obs = np.where(works == 1, wage, np.nan)                                 # wages only seen for workers\n",
    "d = pd.DataFrame({\"treat\": treat, \"wage\": wage_obs})\n",
    "\n",
    "def lee_bounds(data):\n",
    "    obs1 = data[(data.treat == 1)].wage.notna().mean(); obs0 = data[(data.treat == 0)].wage.notna().mean()\n",
    "    w1 = np.sort(data[(data.treat == 1)].wage.dropna().to_numpy()); w0 = data[(data.treat == 0)].wage.dropna()\n",
    "    p = (obs1 - obs0) / obs1                           # share of treated observations to trim\n",
    "    k = int(round(p * len(w1)))\n",
    "    lower = w1[: len(w1) - k].mean() - w0.mean()       # drop the top k\n",
    "    upper = w1[k:].mean() - w0.mean()                  # drop the bottom k\n",
    "    return p, lower, upper\n",
    "\n",
    "p, lower, upper = lee_bounds(d)\n",
    "naive_w = d[d.treat == 1].wage.mean() - d[d.treat == 0].wage.mean()\n",
    "print(f\"Observed wage: treated {d[d.treat == 1].wage.notna().mean():.2f}, control {d[d.treat == 0].wage.notna().mean():.2f} -> trim {p:.1%}\")\n",
    "print(f\"Naive wage difference (workers only): {naive_w:+.2f}\")\n",
    "print(f\"Lee bounds: [{lower:+.2f}, {upper:+.2f}]  (true effect for always-workers: +2)\")"
   ]
  },
  {
   "id": "m14-l1-11",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ✍️ Exercise\n",
    "\n",
    "Bootstrap the Lee bounds (200 resamples) and report a 95% interval for the identified set (use the 2.5th\n",
    "percentile of the lower bound and the 97.5th percentile of the upper bound). Is zero excluded?"
   ]
  },
  {
   "id": "m14-l1-12",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your answer here"
   ]
  },
  {
   "id": "m14-l1-13",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist and recap"
   ]
  }
 ]
}