{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "cells": [
  {
   "id": "m01-l3-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 1.3 · Aggregation, Windows and the Analysis-Ready Panel\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- Compute grouped summaries and rolling statistics with Polars expressions\n",
    "- Build pre- and post-period baselines, the raw material of a before/after comparison\n",
    "- Use Polars lazy queries and save an analysis-ready panel to Parquet\n",
    "- Hand a tidy panel to pandas and statsmodels for estimation\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": "m01-l3-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Rebuild the panel"
   ]
  },
  {
   "id": "m01-l3-2",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q polars duckdb pandas pyarrow statsmodels\n",
    "import polars as pl, duckdb, pandas as pd, numpy as np\n",
    "\n",
    "np.random.seed(7)\n",
    "dates = pd.date_range(\"2018-01-01\", periods=72, freq=\"ME\")\n",
    "provinces = [\"Ontario\", \"Quebec\", \"Alberta\", \"Manitoba\"]\n",
    "\n",
    "rows = []\n",
    "for p_i, p in enumerate(provinces):\n",
    "    level = 5 + p_i * 0.4\n",
    "    for t, d in enumerate(dates):\n",
    "        # A commodity slump hits every province from 2021 onwards\n",
    "        shock = 1.5 if d.year in (2021, 2022) else (0.8 if d.year == 2023 else 0.0)\n",
    "        # Alberta's hiring subsidy (from January 2021) lowers its unemployment by 0.6 points\n",
    "        subsidy = -0.6 if (p == \"Alberta\" and d.year >= 2021) else 0.0\n",
    "        rows.append((d, p, level + shock + subsidy + np.random.normal(0, 0.3),\n",
    "                     2 + 0.03 * t + np.random.normal(0, 0.2)))\n",
    "panel = pl.DataFrame(rows, schema=[\"date\", \"province\", \"unemp_rate\", \"inflation_rate\"], orient=\"row\")\n",
    "print(panel.head())"
   ]
  },
  {
   "id": "m01-l3-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Grouped aggregation with expressions"
   ]
  },
  {
   "id": "m01-l3-4",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** Which province had the highest average unemployment, and how volatile was inflation in each?"
   ]
  },
  {
   "id": "m01-l3-5",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "summary = (\n",
    "    panel.group_by(\"province\")\n",
    "    .agg(\n",
    "        pl.col(\"unemp_rate\").mean().round(2).alias(\"avg_unemp\"),\n",
    "        pl.col(\"unemp_rate\").max().round(2).alias(\"peak_unemp\"),\n",
    "        pl.col(\"inflation_rate\").std().round(3).alias(\"inflation_volatility\"),\n",
    "        pl.len().alias(\"months\"),\n",
    "    )\n",
    "    .sort(\"avg_unemp\", descending=True)\n",
    ")\n",
    "print(summary)"
   ]
  },
  {
   "id": "m01-l3-6",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Rolling windows inside each unit"
   ]
  },
  {
   "id": "m01-l3-7",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** What is the 3-month moving average of unemployment in each province?"
   ]
  },
  {
   "id": "m01-l3-8",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "panel = panel.sort(\"province\", \"date\").with_columns(\n",
    "    pl.col(\"unemp_rate\").rolling_mean(window_size=3).over(\"province\").alias(\"unemp_ma3\"),\n",
    "    (pl.col(\"inflation_rate\") - pl.col(\"inflation_rate\").shift(12).over(\"province\"))\n",
    "        .alias(\"inflation_yoy_change\"),\n",
    ")\n",
    "print(panel.filter(pl.col(\"province\") == \"Quebec\").select(\"date\", \"unemp_rate\", \"unemp_ma3\").head(5))"
   ]
  },
  {
   "id": "m01-l3-9",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pre/post baselines: a first naive comparison"
   ]
  },
  {
   "id": "m01-l3-10",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** How much did Alberta's unemployment change after the January 2021 subsidy?"
   ]
  },
  {
   "id": "m01-l3-11",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "POLICY = pd.Timestamp(\"2021-01-31\")\n",
    "before_after = (\n",
    "    panel.with_columns((pl.col(\"date\") >= POLICY).alias(\"post\"))\n",
    "    .group_by(\"province\", \"post\")\n",
    "    .agg(pl.col(\"unemp_rate\").mean().round(2).alias(\"mean_unemp\"))\n",
    "    .pivot(on=\"post\", index=\"province\", values=\"mean_unemp\")\n",
    "    .rename({\"false\": \"pre\", \"true\": \"post\"})\n",
    "    .with_columns((pl.col(\"post\") - pl.col(\"pre\")).round(2).alias(\"change\"))\n",
    "    .select(\"province\", \"pre\", \"post\", \"change\")\n",
    "    .sort(\"province\")\n",
    ")\n",
    "print(before_after)"
   ]
  },
  {
   "id": "m01-l3-12",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Lazy queries, Parquet and the hand-off to pandas"
   ]
  },
  {
   "id": "m01-l3-13",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "panel.write_parquet(\"macro_panel.parquet\")          # columnar, typed, compressed\n",
    "\n",
    "result = (\n",
    "    pl.scan_parquet(\"macro_panel.parquet\")           # lazy: nothing is read yet\n",
    "    .filter(pl.col(\"date\") >= pd.Timestamp(\"2019-01-01\"))\n",
    "    .group_by(\"province\")\n",
    "    .agg(pl.col(\"unemp_rate\").mean().alias(\"avg_unemp\"))\n",
    "    .sort(\"province\")\n",
    ")\n",
    "print(result.explain())                                # the optimised plan\n",
    "print(result.collect())"
   ]
  },
  {
   "id": "m01-l3-14",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import statsmodels.formula.api as smf\n",
    "df = panel.drop_nulls(\"unemp_ma3\").to_pandas()\n",
    "model = smf.ols(\"unemp_rate ~ inflation_rate + C(province)\", data=df).fit()\n",
    "print(model.params.round(3))"
   ]
  },
  {
   "id": "m01-l3-15",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise and checklist"
   ]
  },
  {
   "id": "m01-l3-16",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ✍️ Exercise\n",
    "\n",
    "Using a window function, compute for each province the **deviation of unemployment from that province's\n",
    "pre-2020 average**. Then report the average deviation in 2021 by province. Which province looks\n",
    "\"most affected\", and why can't this table alone tell you whether the Alberta subsidy worked?"
   ]
  },
  {
   "id": "m01-l3-17",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your answer here"
   ]
  }
 ]
}