Add PyFixest Chapter 27 notebook

pyfixest/Chapter09BridgingFinitePopAndSuperPop.ipynb

@@ -0,0 +1,217 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Chapter 9: Bridging Finite and Super-population Causal Inference"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from joblib import Parallel, delayed\n",
+    "\n",
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "import pyfixest as pf\n",
+    "\n",
+    "np.random.seed(42)\n",
+    "%load_ext autoreload\n",
+    "%autoreload 1\n",
+    "\n",
+    "%load_ext watermark\n",
+    "%watermark --iversions\n",
+    "\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def linestimator(Z, Y, X):\n",
+    "    X = (X - X.mean(axis=0)) / X.std(axis=0)\n",
+    "    n, p = X.shape\n",
+    "    # fully interacted OLS\n",
+    "    covariates = pd.DataFrame(X, columns=[f\"x{i}\" for i in range(p)])\n",
+    "    data = covariates.assign(y=Y, z=Z)\n",
+    "    m = pf.feols(\"y ~ z * (\" + \" + \".join(covariates.columns) + \")\", data=data, vcov=\"HC2\")\n",
+    "    est, vehw = m.coef().loc[\"z\"], m.se().loc[\"z\"] ** 2\n",
+    "    # super-population correction\n",
+    "    inter = m.coef().loc[[f\"z:x{i}\" for i in range(p)]].to_numpy()\n",
+    "    # (β_1 - β_0)' Σ (β_1 - β_0) / n\n",
+    "    superCorr = (inter @ np.cov(X.T) @ inter) / n\n",
+    "    vsuper = vehw + superCorr\n",
+    "    return est, np.sqrt(vehw), np.sqrt(vsuper)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(0.052230404017171474, 0.1475302340448403, 0.1633386978782156)"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "def onerepl(*args):\n",
+    "    n = 500\n",
+    "    X = np.random.normal(0, 1, n * 2).reshape(n, 2)\n",
+    "    Y0 = X[:, 0] + X[:, 0] ** 2 + np.random.uniform(-0.5, 0.5, n)\n",
+    "    Y1 = X[:, 1] + X[:, 1] ** 2 + np.random.uniform(-1, 1, n)\n",
+    "    Z = np.random.binomial(1, 0.6, n)\n",
+    "    Y = Y0 * (1 - Z) + Y1 * Z\n",
+    "    return linestimator(Z, Y, X)\n",
+    "onerepl()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "nrep, k = 2_000, 8\n",
+    "results = Parallel(n_jobs=k)(delayed(onerepl)(i) for i in range(nrep))\n",
+    "simres = np.vstack(results)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(0.0033900784280582142, 0.13559452100782549, 0.15029308662381266)"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# bias, estimated EHW SE, estimated super-population SE\n",
+    "simres[:, 0].mean(), simres[:, 1].mean(), simres[:, 2].mean()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "0.14731850757623555"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# empirical SD\n",
+    "simres[:, 0].std()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "0.926"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# EHW coverage\n",
+    "np.mean(\n",
+    "    (simres[:, 0] - 1.96 * simres[:, 1]) * (simres[:, 0] + 1.96 * simres[:, 1]) <= 0\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "EHW has below nominal coverage for superpopulation."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "0.9515"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# superpop coverage\n",
+    "np.mean(\n",
+    "    (simres[:, 0] - 1.96 * simres[:, 2]) * (simres[:, 0] + 1.96 * simres[:, 2]) <= 0\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Superpopn is above nom coverage for superpopulation."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "metrics",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}