{
 "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 6: Conditional Expectation, the LIE & Total Variance\n",
    "\n",
    "*Beginner to Ninja: Math for Causal Inference*\n",
    "\n",
    "The conditional expectation function, the law of iterated expectations and the adjustment formula.\n",
    "\n",
    "**Learning outcomes**\n",
    "\n",
    "- Compute conditional expectations and use the law of iterated expectations (LIE).\n",
    "- Prove that the conditional expectation function (CEF) is the best predictor of an outcome.\n",
    "- Decompose a variance into within-group and between-group parts.\n",
    "- Use the adjustment formula to compute ATE and ATT from subgroup means."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Theory\n",
    "\n",
    "### The conditional expectation function\n",
    "\n",
    "$E[Y \\mid X = x]$ is the mean of $Y$ among units with $X = x$. As a function of $x$ it is the **conditional expectation function (CEF)**, $m(x)$. Plugging in the random $X$ gives $E[Y \\mid X] = m(X)$, itself a *random variable*.\n",
    "\n",
    "### Three workhorse results\n",
    "\n",
    "**1. Law of iterated expectations (LIE).**\n",
    "\n",
    "$$\n",
    "E[Y] = E\\big[E[Y \\mid X]\\big] = \\sum_x E[Y \\mid X = x]\\,P(X = x).\n",
    "$$\n",
    "\n",
    "The overall mean is the average of the group means, weighted by group size.\n",
    "\n",
    "**2. Taking out what is known.** For any function $g$, $E[g(X)\\,Y \\mid X] = g(X)\\,E[Y \\mid X]$.\n",
    "\n",
    "**3. CEF decomposition.** Write $Y = m(X) + e$ with $e = Y - m(X)$. Then $E[e \\mid X] = 0$, so $e$ is uncorrelated with *every* function of $X$. Moreover, $m(X)$ is the **best predictor** of $Y$: it minimises $E[(Y - g(X))^2]$ over all functions $g$. If the CEF happens to be linear, $m(x) = a + bx$, then $b = \\operatorname{Cov}(X, Y)/\\operatorname{Var}(X)$, the regression slope.\n",
    "\n",
    "### Law of total variance\n",
    "\n",
    "$$\n",
    "\\operatorname{Var}(Y) = \\underbrace{E\\big[\\operatorname{Var}(Y \\mid X)\\big]}_{\\text{within groups}} + \\underbrace{\\operatorname{Var}\\big(E[Y \\mid X]\\big)}_{\\text{between groups}}.\n",
    "$$\n",
    "\n",
    "The share of variance \"explained\" by $X$ is the between part divided by the total: the population $R^2$ of the CEF.\n",
    "\n",
    "### The adjustment formula\n",
    "\n",
    "Suppose treatment is as good as random *within* levels of $X$: $\\big(Y(1), Y(0)\\big) \\perp T \\mid X$ (\"unconfoundedness\"). Then the LIE gives\n",
    "\n",
    "$$\n",
    "E[Y(1)] = \\sum_x E[Y \\mid T = 1, X = x]\\,P(X = x),\n",
    "$$\n",
    "\n",
    "and similarly for $Y(0)$. Weighting by $P(X = x \\mid T = 1)$ instead gives the effect on the treated (ATT). This **adjustment formula** (also called standardisation or the g-formula) is the backbone of matching, regression adjustment and inverse probability weighting.\n",
    "\n",
    "## Part 1 · Warm-up"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 1 · Average wage by education\n",
    "\n",
    "60% of workers have high school (mean wage \\$40k) and 40% college (mean wage \\$70k). Find the average wage.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "By the LIE, $E[W] = 0.6(40) + 0.4(70) = 24 + 28 = 52$ thousand.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 2 · Checking the LIE\n",
    "\n",
    "Using Day 5's table, $E[Y \\mid X = 1] = 0.75$, $E[Y \\mid X = 0] = 0.50$ and $P(X = 1) = 0.4$. Verify that the LIE gives $E[Y] = 0.6$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "$E[Y] = 0.4(0.75) + 0.6(0.50) = 0.30 + 0.30 = 0.60$. ✓\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 3 · Conditional variances\n",
    "\n",
    "In the same table, compute $\\operatorname{Var}(Y \\mid X = 1)$ and $\\operatorname{Var}(Y \\mid X = 0)$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "$Y \\mid X$ is Bernoulli: $\\operatorname{Var}(Y \\mid X = 1) = 0.75 \\times 0.25 = 0.1875$ and $\\operatorname{Var}(Y \\mid X = 0) = 0.5 \\times 0.5 = 0.25$.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 4 · Total variance\n",
    "\n",
    "High-school wages have mean 40 and variance 100; college wages have mean 70 and variance 400; the shares are 0.6 and 0.4. Find $\\operatorname{Var}(W)$ via the law of total variance.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Within-group part, the share-weighted average of group variances:\n",
    "\n",
    "$$\n",
    "E[\\operatorname{Var}(W \\mid X)] = 0.6(100) + 0.4(400) = 60 + 160 = 220.\n",
    "$$\n",
    "\n",
    "**Step 2.** Between-group part, the variance of the group means. $E[W] = 52$ (Exercise 1), so\n",
    "\n",
    "$$\n",
    "\\operatorname{Var}(E[W \\mid X]) = 0.6(40 - 52)^2 + 0.4(70 - 52)^2 = 0.6(144) + 0.4(324) = 86.4 + 129.6 = 216.\n",
    "$$\n",
    "\n",
    "(Shortcut for two groups: $0.6 \\times 0.4 \\times 30^2 = 216$.)\n",
    "\n",
    "**Step 3.** Total: $\\operatorname{Var}(W) = 220 + 216 = 436$.\n",
    "\n",
    "**Answer:** 436 (sd ≈ \\$20.9k).\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 5 · Independence\n",
    "\n",
    "Show that if $X$ and $Y$ are independent, $E[Y \\mid X] = E[Y]$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "Independence means $p(y \\mid x) = p_Y(y)$, so $E[Y \\mid X = x] = \\sum_y y\\,p_Y(y) = E[Y]$ for every $x$. The CEF is flat.\n",
    "\n",
    "</details>\n",
    "\n",
    "## Part 2 · Proofs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 6 · Proving the LIE\n",
    "\n",
    "Prove $E[E[Y \\mid X]] = E[Y]$ for discrete variables.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "$$\n",
    "E\\big[E[Y \\mid X]\\big] = \\sum_x\\Big(\\sum_y y\\,\\frac{p(x, y)}{p_X(x)}\\Big)p_X(x) = \\sum_x\\sum_y y\\,p(x, y) = \\sum_y y\\,p_Y(y) = E[Y].\n",
    "$$\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 7 · The CEF error is orthogonal to functions of X\n",
    "\n",
    "With $e = Y - E[Y \\mid X]$, show that $E[e \\mid X] = 0$ and $E[g(X)\\,e] = 0$ for any $g$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "$E[e \\mid X] = E[Y \\mid X] - E[Y \\mid X] = 0$. Then, by the LIE and \"taking out what is known\",\n",
    "\n",
    "$$\n",
    "E[g(X)e] = E\\big[E[g(X)e \\mid X]\\big] = E\\big[g(X)\\,E[e \\mid X]\\big] = 0.\n",
    "$$\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 8 · The CEF is the best predictor\n",
    "\n",
    "Prove that $E[(Y - g(X))^2] \\ge E[(Y - m(X))^2]$ for any function $g$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Write $Y - g(X) = e + \\big(m(X) - g(X)\\big)$.\n",
    "\n",
    "**Step 2.** Expanding,\n",
    "\n",
    "$$\n",
    "E[(Y - g(X))^2] = E[e^2] + 2E\\big[e\\,(m(X) - g(X))\\big] + E\\big[(m(X) - g(X))^2\\big].\n",
    "$$\n",
    "\n",
    "**Step 3.** The middle term is zero by Exercise 7 (it is $e$ times a function of $X$), and the last is non-negative.\n",
    "\n",
    "**Step 4.** So the error is smallest when $g = m$.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 9 · Proving total variance\n",
    "\n",
    "Prove $\\operatorname{Var}(Y) = E[\\operatorname{Var}(Y \\mid X)] + \\operatorname{Var}(E[Y \\mid X])$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "Write $m(X) = E[Y \\mid X]$.\n",
    "\n",
    "**Step 1.** Conditional variance shortcut: $\\operatorname{Var}(Y \\mid X) = E[Y^2 \\mid X] - m(X)^2$.\n",
    "\n",
    "**Step 2.** Take expectations and use the LIE on the first term:\n",
    "\n",
    "$$\n",
    "E[\\operatorname{Var}(Y \\mid X)] = E[Y^2] - E[m(X)^2].\n",
    "$$\n",
    "\n",
    "**Step 3.** Variance of the CEF, using $E[m(X)] = E[Y]$ (LIE):\n",
    "\n",
    "$$\n",
    "\\operatorname{Var}(m(X)) = E[m(X)^2] - (E[Y])^2.\n",
    "$$\n",
    "\n",
    "**Step 4.** Add steps 2 and 3; the $E[m(X)^2]$ terms cancel:\n",
    "\n",
    "$$\n",
    "E[\\operatorname{Var}(Y \\mid X)] + \\operatorname{Var}(m(X)) = E[Y^2] - (E[Y])^2 = \\operatorname{Var}(Y). \\qquad\\blacksquare\n",
    "$$\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 10 · A linear CEF is the regression line\n",
    "\n",
    "Suppose $E[Y \\mid X] = a + bX$. Prove $b = \\operatorname{Cov}(X, Y)/\\operatorname{Var}(X)$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "By the LIE, $E[XY] = E[X\\,E[Y \\mid X]] = E[X(a + bX)] = aE[X] + bE[X^2]$ and $E[Y] = a + bE[X]$. Then\n",
    "\n",
    "$$\n",
    "\\operatorname{Cov}(X, Y) = aE[X] + bE[X^2] - E[X](a + bE[X]) = b\\operatorname{Var}(X).\n",
    "$$\n",
    "\n",
    "</details>\n",
    "\n",
    "## Part 3 · Applied\n",
    "\n",
    "A hospital's records give mean recovery rates $E[Y \\mid T, S]$ by treatment $T$ and severity $S$. In the population, 70% of cases are mild.\n",
    "\n",
    "| | Mild | Severe |\n",
    "|---|---|---|\n",
    "| Treated ($T = 1$) | 0.90 | 0.60 |\n",
    "| Untreated ($T = 0$) | 0.80 | 0.40 |\n",
    "\n",
    "Assume treatment is as good as random within severity level."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 11 · The adjustment formula\n",
    "\n",
    "Compute $E[Y(1)]$, $E[Y(0)]$ and the ATE.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Unconfoundedness lets us use the treated patients' recovery rate in each severity group as $E[Y(1) \\mid S]$. Average over the *population* case mix (70% mild):\n",
    "\n",
    "$$\n",
    "E[Y(1)] = 0.7(0.90) + 0.3(0.60) = 0.63 + 0.18 = 0.81.\n",
    "$$\n",
    "\n",
    "**Step 2.** Same for the untreated:\n",
    "\n",
    "$$\n",
    "E[Y(0)] = 0.7(0.80) + 0.3(0.40) = 0.56 + 0.12 = 0.68.\n",
    "$$\n",
    "\n",
    "**Step 3.** ATE $= 0.81 - 0.68 = 0.13$.\n",
    "\n",
    "**Answer:** treatment raises the recovery rate by 13 percentage points on average.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 12 · What the naive comparison says\n",
    "\n",
    "Doctors treat the sickest: 60% of treated patients are severe, but only 10% of untreated ones. Compute the naive difference $E[Y \\mid T = 1] - E[Y \\mid T = 0]$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Treated patients' case mix: 40% mild, 60% severe. Their average recovery:\n",
    "\n",
    "$$\n",
    "E[Y \\mid T = 1] = 0.4(0.90) + 0.6(0.60) = 0.36 + 0.36 = 0.72.\n",
    "$$\n",
    "\n",
    "**Step 2.** Untreated case mix: 90% mild, 10% severe:\n",
    "\n",
    "$$\n",
    "E[Y \\mid T = 0] = 0.9(0.80) + 0.1(0.40) = 0.72 + 0.04 = 0.76.\n",
    "$$\n",
    "\n",
    "**Step 3.** Naive difference: $0.72 - 0.76 = -0.04$.\n",
    "\n",
    "**Answer:** the raw comparison says treatment *lowers* recovery by 4 points, although it raises it by 13. Each group is averaged over its own case mix, and the treated are much sicker.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 13 · ATT versus ATE\n",
    "\n",
    "Compute the effect on the treated, $E[Y(1) - Y(0) \\mid T = 1]$. Why does it differ from the ATE?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Effect within each stratum: mild $0.90 - 0.80 = 0.10$; severe $0.60 - 0.40 = 0.20$.\n",
    "\n",
    "**Step 2.** ATT averages these over the *treated* case mix (40% mild, 60% severe):\n",
    "\n",
    "$$\n",
    "\\text{ATT} = 0.4(0.10) + 0.6(0.20) = 0.04 + 0.12 = 0.16.\n",
    "$$\n",
    "\n",
    "**Step 3.** ATE uses the population mix (70% mild): $0.7(0.10) + 0.3(0.20) = 0.13$.\n",
    "\n",
    "**Answer:** ATT = 0.16 > ATE = 0.13, because the effect is larger for severe cases and the treated are disproportionately severe.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 14 · Why a voucher lottery works\n",
    "\n",
    "A city allocates private-school vouchers by lottery: $T = 1$ for winners. Let $Y(1)$ be a child's test score with a voucher. Show that if winning is independent of $Y(1)$, then the winners' average score identifies $E[Y(1)]$, i.e. $E[Y \\mid T = 1] = E[Y(1)]$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Among winners we observe the voucher outcome: $Y = Y(1)$ when $T = 1$, so $E[Y \\mid T = 1] = E[Y(1) \\mid T = 1]$.\n",
    "\n",
    "**Step 2.** The lottery makes $T$ independent of $Y(1)$, so conditioning on $T$ does not change its mean (Exercise 5): $E[Y(1) \\mid T = 1] = E[Y(1)]$.\n",
    "\n",
    "**Step 3.** The same argument for losers gives $E[Y \\mid T = 0] = E[Y(0)]$.\n",
    "\n",
    "**Answer:** the winners' mean minus the losers' mean equals $E[Y(1)] - E[Y(0)]$, the ATE of a voucher. That is why lotteries are prized in education research.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 15 · How much does education explain?\n",
    "\n",
    "Using Exercise 4, what share of wage variance is explained by education group?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** From Exercise 4: between-group variance 216, total 436.\n",
    "\n",
    "**Step 2.** Share explained $= 216/436 \\approx 0.495$.\n",
    "\n",
    "**Answer:** education group explains about half the variance in wages (the population $R^2$ of the CEF); the other half is variation among workers with the same education.\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",
    "mild = 0.7\n",
    "cef = {(1, \"mild\"): 0.90, (1, \"severe\"): 0.60, (0, \"mild\"): 0.80, (0, \"severe\"): 0.40}\n",
    "w_pop = {\"mild\": mild, \"severe\": 1 - mild}\n",
    "w_treated = {\"mild\": 0.4, \"severe\": 0.6}\n",
    "w_untreated = {\"mild\": 0.9, \"severe\": 0.1}\n",
    "\n",
    "def mean(t, w):\n",
    "    return sum(cef[(t, s)] * w[s] for s in w)\n",
    "\n",
    "print(\"ATE  :\", round(mean(1, w_pop) - mean(0, w_pop), 3))\n",
    "print(\"ATT  :\", round(mean(1, w_treated) - mean(0, w_treated), 3))\n",
    "print(\"naive:\", round(mean(1, w_treated) - mean(0, w_untreated), 3))\n",
    "\n",
    "# Law of total variance by simulation (Exercise 4)\n",
    "rng = np.random.default_rng(3)\n",
    "n = 1_000_000\n",
    "college = rng.random(n) < 0.4\n",
    "W = np.where(college, rng.normal(70, 20, n), rng.normal(40, 10, n))\n",
    "within = 0.6 * W[~college].var() + 0.4 * W[college].var()\n",
    "between = 0.6 * 0.4 * (W[college].mean() - W[~college].mean())**2\n",
    "print(f\"Var(W) = {W.var():.1f}; within + between = {within + between:.1f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Key takeaways\n",
    "\n",
    "- The overall mean is a weighted average of conditional means (LIE).\n",
    "- The CEF is the best possible predictor; its error is uncorrelated with every function of $X$.\n",
    "- With unconfoundedness, the adjustment formula turns subgroup comparisons into ATE and ATT."
   ]
  }
 ]
}