{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "cells": [
  {
   "id": "m03-l1-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 3.1 · OLS Mechanics: the Line of Best Fit\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- Explain OLS as the line that minimises the sum of squared residuals\n",
    "- Recover the same coefficients three ways: numerical optimisation, matrix formula and statsmodels\n",
    "- Interpret slopes, intercepts, residuals and R² in a Mincer earnings equation\n",
    "- Use log outcomes and read coefficients as approximate percentage effects\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": "m03-l1-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Why regression is still the workhorse"
   ]
  },
  {
   "id": "m03-l1-2",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q numpy scipy polars statsmodels matplotlib seaborn\n",
    "import numpy as np, polars as pl\n",
    "import statsmodels.formula.api as smf\n",
    "from scipy.optimize import minimize\n",
    "import matplotlib.pyplot as plt, seaborn as sns\n",
    "\n",
    "rng = np.random.default_rng(11)\n",
    "sns.set_theme(style=\"darkgrid\")"
   ]
  },
  {
   "id": "m03-l1-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulating a Mincer earnings world"
   ]
  },
  {
   "id": "m03-l1-4",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "N = 1_000\n",
    "TRUE_RETURN = 2_500\n",
    "education = np.clip(np.round(rng.normal(13, 2, N)), 8, 20)\n",
    "experience = rng.uniform(0, 30, N)\n",
    "wage = 20_000 + TRUE_RETURN * education + 600 * experience + rng.normal(0, 6_000, N)\n",
    "\n",
    "labor = pl.DataFrame({\"education\": education, \"experience\": experience, \"wage\": wage})\n",
    "print(labor.describe())"
   ]
  },
  {
   "id": "m03-l1-5",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Three ways to get the same line"
   ]
  },
  {
   "id": "m03-l1-6",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** Which straight line fits the wage–education cloud best, and what does \"best\" mean?"
   ]
  },
  {
   "id": "m03-l1-7",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "x, y = labor[\"education\"].to_numpy(), labor[\"wage\"].to_numpy()\n",
    "\n",
    "def ssr(beta):\n",
    "    return np.sum((y - (beta[0] + beta[1] * x)) ** 2)\n",
    "\n",
    "opt = minimize(ssr, x0=[0.0, 0.0], method=\"BFGS\")\n",
    "print(\"Optimiser:   intercept = {:,.0f}, slope = {:,.1f}\".format(*opt.x))"
   ]
  },
  {
   "id": "m03-l1-8",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "X = np.column_stack([np.ones(N), x])\n",
    "beta_hat = np.linalg.solve(X.T @ X, X.T @ y)\n",
    "print(\"Matrix:      intercept = {:,.0f}, slope = {:,.1f}\".format(*beta_hat))"
   ]
  },
  {
   "id": "m03-l1-9",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "df = labor.to_pandas()\n",
    "simple = smf.ols(\"wage ~ education\", data=df).fit()\n",
    "print(\"statsmodels: intercept = {:,.0f}, slope = {:,.1f}\".format(*simple.params))"
   ]
  },
  {
   "id": "m03-l1-10",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reading the output: coefficients, residuals, R²"
   ]
  },
  {
   "id": "m03-l1-11",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "print(simple.summary().tables[1])\n",
    "print(f\"R-squared: {simple.rsquared:.3f}\")\n",
    "\n",
    "plt.figure(figsize=(10, 5))\n",
    "sns.scatterplot(x=\"education\", y=\"wage\", data=df, alpha=0.4, label=\"Workers\")\n",
    "grid = np.linspace(8, 20, 50)\n",
    "plt.plot(grid, simple.params[\"Intercept\"] + simple.params[\"education\"] * grid, color=\"red\", lw=3, label=\"OLS fit\")\n",
    "plt.title(\"Wage vs years of education\"); plt.legend(); plt.show()"
   ]
  },
  {
   "id": "m03-l1-12",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "print(f\"Mean residual: {simple.resid.mean():.2e} | corr(resid, education): {np.corrcoef(simple.resid, x)[0, 1]:.2e}\")"
   ]
  },
  {
   "id": "m03-l1-13",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Logs and percentage effects"
   ]
  },
  {
   "id": "m03-l1-14",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** By roughly what percentage does one more year of schooling raise wages?"
   ]
  },
  {
   "id": "m03-l1-15",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "df[\"log_wage\"] = np.log(df[\"wage\"])\n",
    "log_model = smf.ols(\"log_wage ~ education + experience\", data=df).fit()\n",
    "b = log_model.params[\"education\"]\n",
    "print(f\"Coefficient {b:.4f}  ->  about {100 * b:.1f}% per year (exact: {100 * (np.exp(b) - 1):.1f}%)\")"
   ]
  },
  {
   "id": "m03-l1-16",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ✍️ Exercise\n",
    "\n",
    "Add a quadratic in experience (`experience + I(experience**2)`) to the log-wage model, the classic Mincer\n",
    "specification. Use the coefficients to compute the experience level at which predicted wages peak. Does a\n",
    "peak exist in our simulated world? Why or why not?"
   ]
  },
  {
   "id": "m03-l1-17",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your answer here"
   ]
  },
  {
   "id": "m03-l1-18",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist and recap"
   ]
  }
 ]
}