{
 "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 28: Frisch–Waugh–Lovell & Partialling Out\n",
    "\n",
    "*Beginner to Ninja: Math for Causal Inference*\n",
    "\n",
    "What \"controlling for\" really does. Every multiple-regression coefficient is a simple regression on residualised data.\n",
    "\n",
    "**Learning outcomes**\n",
    "\n",
    "- State and prove the Frisch–Waugh–Lovell (FWL) theorem.\n",
    "- Compute a multiple-regression coefficient by partialling out, and read it from an added-variable plot.\n",
    "- Show that fixed-effects (within) estimators are FWL applied to group dummies.\n",
    "- Explain how OLS with controls weights heterogeneous treatment effects, and preview partialling out with machine learning."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Theory\n",
    "\n",
    "### The theorem\n",
    "\n",
    "Split the regressors: $y = X_1\\beta_1 + X_2\\beta_2 + \\varepsilon$. Let $M_2 = I - X_2(X_2^{\\top}X_2)^{-1}X_2^{\\top}$, the residual-maker for $X_2$ (Day 20). Define the residualised variables\n",
    "\n",
    "$$\n",
    "\\tilde X_1 = M_2X_1, \\qquad \\tilde y = M_2y.\n",
    "$$\n",
    "\n",
    "**FWL theorem.** The OLS coefficient on $X_1$ in the full regression equals the coefficient from regressing $\\tilde y$ on $\\tilde X_1$:\n",
    "\n",
    "$$\n",
    "\\hat\\beta_1 = (\\tilde X_1^{\\top}\\tilde X_1)^{-1}\\tilde X_1^{\\top}\\tilde y = (\\tilde X_1^{\\top}\\tilde X_1)^{-1}\\tilde X_1^{\\top}y,\n",
    "$$\n",
    "\n",
    "and the residuals of the two regressions are identical.\n",
    "\n",
    "### What it means\n",
    "\n",
    "\"Controlling for $X_2$\" means: strip out of both $y$ and $X_1$ everything linearly predictable from $X_2$, then relate what is left. For a single regressor this is the **regression anatomy** formula\n",
    "\n",
    "$$\n",
    "\\hat\\beta_1 = \\frac{\\widehat{\\operatorname{Cov}}(y, \\tilde x_1)}{\\widehat{\\operatorname{Var}}(\\tilde x_1)}.\n",
    "$$\n",
    "\n",
    "Consequences you will use repeatedly:\n",
    "\n",
    "- **Fixed effects:** with $X_2$ = group dummies, $M_2$ subtracts group means, so the within (demeaned) regression gives the same slope as including the dummies.\n",
    "- **Precision:** $\\operatorname{Var}(\\hat\\beta_1) = \\sigma^2/\\tilde x_1^{\\top}\\tilde x_1$, and $\\tilde x_1^{\\top}\\tilde x_1 = (1 - R_1^2)\\sum(x_1 - \\bar x_1)^2$, giving the VIF of Day 18.\n",
    "- **Added-variable plot:** scatter $\\tilde y$ against $\\tilde x_1$; the fitted slope is $\\hat\\beta_1$.\n",
    "- **Double machine learning** (Day 90) replaces the linear projections with flexible machine-learning predictions: residualise $y$ and the treatment on $X$, then regress residual on residual.\n",
    "\n",
    "## Part 1 · Warm-up"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 1 · What M₂ does\n",
    "\n",
    "In words, what is $M_2X_1$?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "The residuals from regressing each column of $X_1$ on $X_2$: the part of $X_1$ that $X_2$ cannot linearly explain.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 2 · The intercept is partialled out too\n",
    "\n",
    "Show that the slope in $y = \\beta_0 + \\beta_1x$ equals the slope from regressing demeaned $y$ on demeaned $x$ without an intercept.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "Take $X_2 = \\mathbf 1$: $M_2$ is the centring matrix, so $\\tilde x = x - \\bar x$ and $\\tilde y = y - \\bar y$. FWL gives $\\hat\\beta_1 = \\tilde x^{\\top}\\tilde y/\\tilde x^{\\top}\\tilde x$, the Cov/Var formula.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 3 · Partialling out by hand\n",
    "\n",
    "Four workers: $x_2 = (0, 0, 1, 1)$ marks urban residence, $x_1 = (1, 2, 3, 5)$ is years of post-secondary schooling and $y = (2, 3, 6, 9)$ is hourly earnings (in \\$10). Regress earnings on an intercept, schooling and the urban dummy. Compute the schooling coefficient $\\hat\\beta_1$ by FWL.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Partialling out $(1, x_2)$ means subtracting the rural and urban means (FWL with group dummies).\n",
    "\n",
    "**Step 2.** Schooling means: rural $(1 + 2)/2 = 1.5$, urban $(3 + 5)/2 = 4$. So $\\tilde x_1 = (-0.5, 0.5, -1, 1)$.\n",
    "\n",
    "**Step 3.** Earnings means: rural $2.5$, urban $7.5$. So $\\tilde y = (-0.5, 0.5, -1.5, 1.5)$.\n",
    "\n",
    "**Step 4.** Residual-on-residual slope:\n",
    "\n",
    "$$\n",
    "\\hat\\beta_1 = \\frac{\\tilde x_1^{\\top}\\tilde y}{\\tilde x_1^{\\top}\\tilde x_1} = \\frac{0.25 + 0.25 + 1.5 + 1.5}{0.25 + 0.25 + 1 + 1} = \\frac{3.5}{2.5} = 1.4.\n",
    "$$\n",
    "\n",
    "**Answer:** each extra year of schooling goes with \\$14/hour more earnings *within* the same area.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 4 · The short regression\n",
    "\n",
    "For the same data, compute the slope from regressing $y$ on $x_1$ alone (with an intercept).\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Means: $\\bar x_1 = 11/4 = 2.75$ and $\\bar y = 20/4 = 5$.\n",
    "\n",
    "**Step 2.** Deviations: $(-1.75, -0.75, 0.25, 2.25)$ and $(-3, -2, 1, 4)$.\n",
    "\n",
    "**Step 3.** Cross-product $5.25 + 1.5 + 0.25 + 9 = 16$; squared deviations $3.0625 + 0.5625 + 0.0625 + 5.0625 = 8.75$.\n",
    "\n",
    "**Step 4.** Slope $= 16/8.75 \\approx 1.83$.\n",
    "\n",
    "**Answer:** 1.83 without the urban control versus 1.40 with it: part of schooling's raw association with earnings is the urban premium.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 5 · Only one side needs residualising\n",
    "\n",
    "Show that $\\tilde X_1^{\\top}y = \\tilde X_1^{\\top}\\tilde y$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** $\\tilde X_1^{\\top}\\tilde y = X_1^{\\top}M_2^{\\top}M_2y = X_1^{\\top}M_2y = \\tilde X_1^{\\top}y$, since $M_2$ is symmetric and idempotent.\n",
    "\n",
    "**Step 2.** So regressing raw $y$ on $\\tilde X_1$ gives the same coefficient (but not the same residuals or SEs).\n",
    "\n",
    "</details>\n",
    "\n",
    "## Part 2 · Proofs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 6 · Proving FWL\n",
    "\n",
    "Prove the FWL theorem.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Write the full OLS fit: $y = X_1\\hat\\beta_1 + X_2\\hat\\beta_2 + e$, with $X_1^{\\top}e = 0$ and $X_2^{\\top}e = 0$ (Day 20).\n",
    "\n",
    "**Step 2.** Multiply through by $M_2$. $M_2X_2 = 0$, and $M_2e = e$ because $e$ is already orthogonal to $X_2$ ($P_2e = 0$):\n",
    "\n",
    "$$\n",
    "M_2y = M_2X_1\\hat\\beta_1 + e, \\quad\\text{i.e.}\\quad \\tilde y = \\tilde X_1\\hat\\beta_1 + e.\n",
    "$$\n",
    "\n",
    "**Step 3.** Pre-multiply by $\\tilde X_1^{\\top} = X_1^{\\top}M_2$: $\\tilde X_1^{\\top}\\tilde y = \\tilde X_1^{\\top}\\tilde X_1\\hat\\beta_1 + X_1^{\\top}M_2e$.\n",
    "\n",
    "**Step 4.** $X_1^{\\top}M_2e = X_1^{\\top}e = 0$ (Step 1).\n",
    "\n",
    "**Step 5.** So $\\hat\\beta_1 = (\\tilde X_1^{\\top}\\tilde X_1)^{-1}\\tilde X_1^{\\top}\\tilde y$. $\\blacksquare$\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 7 · Identical residuals\n",
    "\n",
    "Show that the residuals from regressing $\\tilde y$ on $\\tilde X_1$ equal the full-regression residuals $e$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "From Exercise 6, $\\tilde y - \\tilde X_1\\hat\\beta_1 = e$. So $e^{\\top}e$ is the same, but the residualised regression should still use $n - k$ (all regressors), not $n - k_1$, when computing $s^2$.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 8 · Regression anatomy\n",
    "\n",
    "For a single regressor $x_1$, show $\\hat\\beta_1 = \\widehat{\\operatorname{Cov}}(y, \\tilde x_1)/\\widehat{\\operatorname{Var}}(\\tilde x_1)$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** By FWL and Exercise 5, $\\hat\\beta_1 = \\tilde x_1^{\\top}y/\\tilde x_1^{\\top}\\tilde x_1$.\n",
    "\n",
    "**Step 2.** If $X_2$ includes an intercept, $\\tilde x_1$ has mean zero, so $\\tilde x_1^{\\top}y = \\sum\\tilde x_{1i}(y_i - \\bar y)$, i.e. $n - 1$ times the sample covariance; the denominator is $n - 1$ times the variance.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 9 · Orthogonal controls change nothing\n",
    "\n",
    "Prove that if $X_1^{\\top}X_2 = 0$, adding $X_2$ leaves $\\hat\\beta_1$ unchanged.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** $\\tilde X_1 = X_1 - X_2(X_2^{\\top}X_2)^{-1}X_2^{\\top}X_1 = X_1$.\n",
    "\n",
    "**Step 2.** FWL then gives $(X_1^{\\top}X_1)^{-1}X_1^{\\top}y$, the short-regression coefficient.\n",
    "\n",
    "**Step 3.** Randomised treatments are (approximately) orthogonal to pre-treatment covariates, which is why adding covariates changes the point estimate little while improving precision (Day 27).\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 10 · Fixed effects are FWL\n",
    "\n",
    "Show that including a full set of group dummies $D$ is equivalent to demeaning $y$ and $X_1$ by group.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** With $X_2 = D$ (no separate intercept), $D^{\\top}D = \\operatorname{diag}(n_g)$ and $D^{\\top}y$ stacks group sums, so $P_Dy$ replaces each $y_i$ by its group mean.\n",
    "\n",
    "**Step 2.** Hence $M_Dy = y_i - \\bar y_{g(i)}$ and likewise for $X_1$; FWL gives the within estimator.\n",
    "\n",
    "</details>\n",
    "\n",
    "## Part 3 · Applied"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 11 · Interpreting the change\n",
    "\n",
    "In Exercises 3–4 the coefficient on $x_1$ falls from 1.83 to 1.40 when $x_2$ is added. Explain using the OVB formula.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Urban workers have more schooling (group means 4 versus 1.5).\n",
    "\n",
    "**Step 2.** Urban residence raises earnings at a given schooling level (the full regression gives a coefficient of 1.5 on the dummy, i.e. \\$15/hour).\n",
    "\n",
    "**Step 3.** Omitting it loads part of the urban premium onto schooling: OVB = (effect of urban) × (slope of urban on schooling) > 0.\n",
    "\n",
    "**Answer:** the drop from 1.83 to 1.40 removes that bias, *provided* urban residence is a confounder. If schooling caused people to move to cities (a mediator), controlling for it would instead remove part of schooling's effect.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 12 · Two-way fixed effects\n",
    "\n",
    "In a balanced panel, show that partialling out unit and time dummies transforms $x_{it}$ into $x_{it} - \\bar x_{i\\cdot} - \\bar x_{\\cdot t} + \\bar x_{\\cdot\\cdot}$.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Projecting on unit dummies replaces $x_{it}$ by the unit mean $\\bar x_{i\\cdot}$; projecting on time dummies replaces it by the period mean $\\bar x_{\\cdot t}$.\n",
    "\n",
    "**Step 2.** In a balanced panel the two sets of dummies are orthogonal once the overall mean is removed, so the projection onto both adds the two parts and subtracts the grand mean once (otherwise it would be counted twice): $\\bar x_{i\\cdot} + \\bar x_{\\cdot t} - \\bar x_{\\cdot\\cdot}$.\n",
    "\n",
    "**Step 3.** The residual is $x_{it} - \\bar x_{i\\cdot} - \\bar x_{\\cdot t} + \\bar x_{\\cdot\\cdot}$.\n",
    "\n",
    "**Answer:** two-way FE = double demeaning. With unbalanced panels or staggered treatment this simple form fails, part of the story of Day 70.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 13 · OLS weights heterogeneous effects by treatment variance\n",
    "\n",
    "Two strata of equal size have treatment effects 1 and 3. The treatment probability is 0.5 in stratum A and 0.1 in stratum B. Regressing $y$ on $T$ and stratum dummies gives approximately $\\sum_sw_s\\tau_s$ with $w_s \\propto P(s)\\,p_s(1 - p_s)$. Compute it and compare with the ATE.\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** Conditional treatment variances: stratum A $0.5 \\times 0.5 = 0.25$; stratum B $0.1 \\times 0.9 = 0.09$.\n",
    "\n",
    "**Step 2.** Weights (equal stratum sizes): $w_A = 0.25/(0.25 + 0.09) \\approx 0.735$, $w_B \\approx 0.265$.\n",
    "\n",
    "**Step 3.** Regression estimate: $0.735(1) + 0.265(3) \\approx 1.53$.\n",
    "\n",
    "**Step 4.** ATE: $0.5(1) + 0.5(3) = 2$.\n",
    "\n",
    "**Answer:** 1.53 versus 2. With heterogeneous effects, OLS with controls estimates a *variance-weighted* average (Angrist, 1998), favouring strata where treatment is closest to 50/50.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 14 · Partialling out with machine learning: 401(k) eligibility\n",
    "\n",
    "To estimate the effect of 401(k) eligibility $D$ on household savings $y$, controlling flexibly for income, age and education $X$, double ML assumes $y = \\theta D + g(X) + \\varepsilon$ and regresses $y - \\hat E[y \\mid X]$ on $D - \\hat E[D \\mid X]$, with both predictions from machine learning. Why is this FWL?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** FWL removes from $y$ and $D$ everything *linearly* predictable from the controls, then regresses residual on residual.\n",
    "\n",
    "**Step 2.** Here the linear projection is replaced by the conditional expectation $E[\\cdot \\mid X]$, estimated flexibly (e.g. random forests), so $g(X)$ need not be linear.\n",
    "\n",
    "**Step 3.** The residualised $D$ is the variation in eligibility not explained by income, age and education; its relation to residualised savings identifies $\\theta$.\n",
    "\n",
    "**Answer:** it is FWL with machine-learning projections (Robinson, 1988; Day 90).\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Your work:*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 15 · An added-variable plot\n",
    "\n",
    "Why can an added-variable plot reveal an influential observation that a raw scatter of $y$ against $x_1$ hides?\n",
    "\n",
    "<details><summary>Show solution</summary>\n",
    "\n",
    "**Step 1.** By FWL, $\\hat\\beta_1$ depends only on the residualised data $(\\tilde x_1, \\tilde y)$.\n",
    "\n",
    "**Step 2.** A unit can look ordinary in raw $x_1$ yet be unusual *given* $X_2$ (a large $\\tilde x_1$), which gives it high leverage on $\\hat\\beta_1$.\n",
    "\n",
    "**Step 3.** A raw scatter of $y$ on $x_1$ hides this; the added-variable plot of $\\tilde y$ against $\\tilde x_1$ shows it.\n",
    "\n",
    "**Answer:** the added-variable plot shows the data the coefficient actually uses (see the FWL simulator).\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",
    "x2 = np.array([0, 0, 1, 1.]); x1 = np.array([1, 2, 3, 5.]); y = np.array([2, 3, 6, 9.])\n",
    "X = np.column_stack([np.ones(4), x1, x2])\n",
    "b_full = np.linalg.lstsq(X, y, rcond=None)[0]\n",
    "X2 = np.column_stack([np.ones(4), x2]); M2 = np.eye(4) - X2 @ np.linalg.pinv(X2)\n",
    "x1t, yt = M2 @ x1, M2 @ y\n",
    "print(\"full regression:\", b_full.round(4))\n",
    "print(\"FWL slope:\", round(x1t @ yt / (x1t @ x1t), 4), \" raw y on residualised x1:\", round(x1t @ y / (x1t @ x1t), 4))\n",
    "print(\"short regression slope:\", round(np.polyfit(x1, y, 1)[0], 4))\n",
    "e_full = y - X @ b_full; e_fwl = yt - x1t * (x1t @ yt / (x1t @ x1t))\n",
    "print(\"same residuals:\", np.allclose(e_full, e_fwl))\n",
    "\n",
    "# Variance-weighted effects (Exercise 13)\n",
    "rng = np.random.default_rng(28); n = 400_000\n",
    "s = rng.random(n) < 0.5                                       # stratum B\n",
    "p = np.where(s, 0.1, 0.5); T = rng.random(n) < p\n",
    "Y = np.where(s, 5, 0) + np.where(s, 3, 1) * T + rng.normal(size=n)\n",
    "Z = np.column_stack([np.ones(n), T, s])\n",
    "print(\"regression estimate:\", np.linalg.lstsq(Z, Y, rcond=None)[0][1].round(3), \" ATE: 2\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Key takeaways\n",
    "\n",
    "- A multiple-regression coefficient is a simple regression of residualised $y$ on residualised $x$.\n",
    "- Fixed effects, demeaning and \"controlling for\" are all FWL.\n",
    "- With heterogeneous effects, OLS with controls weights strata by the conditional variance of treatment, not by population share."
   ]
  }
 ]
}