Chapter15Matching.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chapter 15: Matching in Observational Studies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "seaborn          : 0.13.0\n",
      "sklearn          : 1.3.2\n",
      "matplotlib       : 3.7.2\n",
      "matplotlib_inline: 0.1.6\n",
      "numpy            : 1.24.3\n",
      "statsmodels      : 0.14.0\n",
      "pandas           : 2.0.3\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from joblib import Parallel, delayed\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import statsmodels.api as sm\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import sklearn as skl\n",
    "\n",
    "font = {'family' : 'IBM Plex Sans Condensed',\n",
    "               'weight' : 'normal',\n",
    "               'size'   : 10}\n",
    "plt.rc('font', **font)\n",
    "plt.rcParams['figure.figsize'] = (6, 6)\n",
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "\n",
    "np.random.seed(42)\n",
    "%load_ext autoreload\n",
    "%autoreload 1\n",
    "\n",
    "%load_ext watermark\n",
    "%watermark --iversions\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## experimental data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>treatment</th>\n",
       "      <th>age</th>\n",
       "      <th>education</th>\n",
       "      <th>black</th>\n",
       "      <th>hispanic</th>\n",
       "      <th>married</th>\n",
       "      <th>nodegree</th>\n",
       "      <th>earnings1974</th>\n",
       "      <th>earnings1975</th>\n",
       "      <th>earnings1978</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1.0</td>\n",
       "      <td>37.0</td>\n",
       "      <td>11.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>9930.0460</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1.0</td>\n",
       "      <td>22.0</td>\n",
       "      <td>9.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>3595.8940</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1.0</td>\n",
       "      <td>30.0</td>\n",
       "      <td>12.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>24909.4500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1.0</td>\n",
       "      <td>27.0</td>\n",
       "      <td>11.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>7506.1460</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1.0</td>\n",
       "      <td>33.0</td>\n",
       "      <td>8.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>289.7899</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   treatment   age  education  black  hispanic  married  nodegree  \\\n",
       "0        1.0  37.0       11.0    1.0       0.0      1.0       1.0   \n",
       "1        1.0  22.0        9.0    0.0       1.0      0.0       1.0   \n",
       "2        1.0  30.0       12.0    1.0       0.0      0.0       0.0   \n",
       "3        1.0  27.0       11.0    1.0       0.0      0.0       1.0   \n",
       "4        1.0  33.0        8.0    1.0       0.0      0.0       1.0   \n",
       "\n",
       "   earnings1974  earnings1975  earnings1978  \n",
       "0           0.0           0.0     9930.0460  \n",
       "1           0.0           0.0     3595.8940  \n",
       "2           0.0           0.0    24909.4500  \n",
       "3           0.0           0.0     7506.1460  \n",
       "4           0.0           0.0      289.7899  "
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import empirical_calibration as ec\n",
    "import empirical_calibration.data.lalonde as lalonde\n",
    "\n",
    "treat, ctrl = lalonde.experimental_treated(), lalonde.experimental_control()\n",
    "lalonde_exp = pd.concat([treat, ctrl])\n",
    "lalonde_exp.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "y, z = lalonde_exp.earnings1978.values, lalonde_exp.treatment.values\n",
    "X = lalonde_exp.drop(columns=[\"earnings1978\", \"treatment\"]).values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>coef</th>\n",
       "      <th>se</th>\n",
       "      <th>t</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>neyman</th>\n",
       "      <td>1794.342404</td>\n",
       "      <td>670.996546</td>\n",
       "      <td>2.674146</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>fisher</th>\n",
       "      <td>1676.342644</td>\n",
       "      <td>677.049148</td>\n",
       "      <td>2.475954</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>lin</th>\n",
       "      <td>1621.583101</td>\n",
       "      <td>694.721574</td>\n",
       "      <td>2.334148</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "               coef          se         t\n",
       "neyman  1794.342404  670.996546  2.674146\n",
       "fisher  1676.342644  677.049148  2.475954\n",
       "lin     1621.583101  694.721574  2.334148"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def reg_adjust(z, y, X):\n",
    "    tau_n_fit = sm.OLS(y, sm.add_constant(z)).fit(cov_type=\"HC2\")\n",
    "    tau_f_fit = sm.OLS(y, np.c_[sm.add_constant(z), X]).fit(cov_type=\"HC2\")\n",
    "    tau_l_fit = sm.OLS(\n",
    "        y,\n",
    "        np.c_[\n",
    "            sm.add_constant(z),\n",
    "            X - X.mean(axis=0),\n",
    "            z.reshape(-1, 1) * (X - X.mean(axis=0)),\n",
    "        ],\n",
    "    ).fit(cov_type=\"HC2\")\n",
    "    resmat = np.r_[\n",
    "        np.c_[tau_n_fit.params[1], tau_n_fit.bse[1], tau_n_fit.tvalues[1]],\n",
    "        np.c_[tau_f_fit.params[1], tau_f_fit.bse[1], tau_f_fit.tvalues[1]],\n",
    "        np.c_[tau_l_fit.params[1], tau_l_fit.bse[1], tau_l_fit.tvalues[1]],\n",
    "    ]\n",
    "    return pd.DataFrame(\n",
    "        resmat, index=[\"neyman\", \"fisher\", \"lin\"], columns=[\"coef\", \"se\", \"t\"]\n",
    "    )\n",
    "\n",
    "\n",
    "reg_adjust(z, y, X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Matching"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "def one_nn_att(X, z, y, k=1, bias_corr_mod=None):\n",
    "    \"\"\"Matching estimator for the ATT using 1-NN.\n",
    "\n",
    "    Args:\n",
    "        X (2d Array): Covariate Matrix\n",
    "        z (1d Array): Treatment Indicator\n",
    "        y (1d Array): Outcome vector\n",
    "        k (int, optional): Number of neighbours. Defaults to 1.\n",
    "        bias_corr_mod (sklearn model, optional): Model used for bias correction. Defaults to None.\n",
    "\n",
    "    Returns:\n",
    "        estimate, ai SE, and optionally Otsu and Rai (2017) SE\n",
    "    \"\"\"\n",
    "    from sklearn.neighbors import KNeighborsRegressor\n",
    "\n",
    "    n, n1 = len(z), z.sum()\n",
    "    mod = KNeighborsRegressor(n_neighbors=k)\n",
    "    treat_nn_mod = mod.fit(X[z == 0, :], y[z == 0])\n",
    "    Y0hat = treat_nn_mod.predict(X[z == 1, :])\n",
    "    point_est = y[z == 1].mean() - Y0hat.mean()\n",
    "    # store neighbour index for each treated\n",
    "    _, neighbours = mod.kneighbors(X[z == 1, :])\n",
    "    if bias_corr_mod:  # bias correction with outcome model\n",
    "        # outcome model\n",
    "        muhat = bias_corr_mod.fit(X[z == 0, :], y[z == 0])\n",
    "        # bias correction term is μ^0(x_i) - μ^0(x_j) for each\n",
    "        bias_corr_term = muhat.predict(X[z == 1, :]) - muhat.predict(\n",
    "            X[z == 0, :][neighbours.flatten(), :]\n",
    "        )\n",
    "        point_est = y[z == 1].mean() - Y0hat.mean() - bias_corr_term.mean()\n",
    "    ######################################################################\n",
    "    # abadie/imbens variance\n",
    "    ######################################################################\n",
    "    # 1/N1^2 ∑ (Y_i - \\hat{Y}_i - \\hat{τ})^2\n",
    "    first_term = 1 / (n1**2) * np.sum(y - treat_nn_mod.predict(X) - point_est) ** 2\n",
    "    # second term\n",
    "    # 1/N1^2 ∑ (K_i^2 - K_sq) σ^2\n",
    "    mod2 = KNeighborsRegressor(n_neighbors=k)\n",
    "    ctrl_nn_mod = mod2.fit(X[z == 1, :], y[z == 1])\n",
    "    Yhat_all = np.zeros(n)\n",
    "    Yhat_all[z == 1], Yhat_all[z == 0] = Y0hat, ctrl_nn_mod.predict(X[z == 0, :])\n",
    "    sigma2 = 1 / 2 * np.sum((y - Yhat_all) ** 2)\n",
    "    # K terms: counts are the number of times each ctrl obs is matched with a treated obs\n",
    "    ctrl_index, counts = np.unique(neighbours, return_counts=True)\n",
    "    K, Ksq = np.zeros(n), np.zeros(n)\n",
    "    K[ctrl_index], Ksq[ctrl_index] = counts / k, counts / (k**2)\n",
    "    # conditional variance of Y given X, W\n",
    "    second_term = 1 / (n1**2) * np.sum((K**2 - Ksq) * sigma2)\n",
    "    v_ai = first_term + second_term\n",
    "    ######################################################################\n",
    "    # otsu and rai\n",
    "    ######################################################################\n",
    "    if bias_corr_mod:\n",
    "        psi = z * (y - muhat.predict(X)) - (1 - z) * (K / k) * (y - muhat.predict(X))\n",
    "        v_or = (1 / (n1**2)) * np.sum((psi - point_est * n1 / n) ** 2)\n",
    "        return point_est, np.sqrt(v_ai / n), np.sqrt(v_or)\n",
    "    return point_est, np.sqrt(v_ai / n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# wrapper function that bootstraps the standard error\n",
    "def nn_att(X, z, y, k=1, bias_corr_mod=None, n_boot=1e3):\n",
    "    analytic_est = one_nn_att(X, z, y, k=k, bias_corr_mod=bias_corr_mod)\n",
    "    # bootstrap (invalid?)\n",
    "    n = len(z)\n",
    "\n",
    "    def bootfn(*args):\n",
    "        # draw indices\n",
    "        ids = np.random.choice(np.arange(n), size=n, replace=True)\n",
    "        return one_nn_att(X[ids,], z[ids], y[ids], k=1, bias_corr_mod=bias_corr_mod)\n",
    "\n",
    "    boot_est = Parallel(n_jobs=-1)(delayed(bootfn)(i) for i in range(int(n_boot)))\n",
    "    boot_est = np.stack(boot_est)\n",
    "    boot_se = boot_est[:, 0].std()\n",
    "    if bias_corr_mod:\n",
    "        return pd.DataFrame(\n",
    "            np.c_[*analytic_est, boot_se], columns=[\"est\", \"ai_se\", \"or_se\", \"boot_se\"]\n",
    "        )\n",
    "    return pd.DataFrame(\n",
    "        np.c_[*analytic_est, boot_se], columns=[\"est\", \"ai_se\", \"boot_se\"]\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "No bias correction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>est</th>\n",
       "      <th>ai_se</th>\n",
       "      <th>boot_se</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2011.153187</td>\n",
       "      <td>585.758733</td>\n",
       "      <td>853.85284</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           est       ai_se    boot_se\n",
       "0  2011.153187  585.758733  853.85284"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nn_att(X, z, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "bias correction with OLS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>est</th>\n",
       "      <th>ai_se</th>\n",
       "      <th>or_se</th>\n",
       "      <th>boot_se</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1922.564125</td>\n",
       "      <td>583.441871</td>\n",
       "      <td>668.21205</td>\n",
       "      <td>846.011315</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           est       ai_se      or_se     boot_se\n",
       "0  1922.564125  583.441871  668.21205  846.011315"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from sklearn.linear_model import LinearRegression\n",
    "\n",
    "nn_att(X, z, y, bias_corr_mod=LinearRegression())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Bootstrap SE is more in-line with the OLS SE, while AI SE looks too small."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### observational"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read CPS data\n",
    "dat = pd.read_table(\"cps1re74.csv\", delimiter=\" \")\n",
    "dat[\"u74\"], dat[\"u75\"] = 1 * (dat.re74 == 0), 1 * (dat.re75 == 0)\n",
    "z, y = dat.treat.values, dat.re78.values\n",
    "X = dat.drop(columns=[\"treat\", \"re78\"]).values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Regression is bad."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>coef</th>\n",
       "      <th>se</th>\n",
       "      <th>t</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>neyman</th>\n",
       "      <td>-8506.495361</td>\n",
       "      <td>583.442609</td>\n",
       "      <td>-14.579832</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>fisher</th>\n",
       "      <td>1067.546135</td>\n",
       "      <td>628.438879</td>\n",
       "      <td>1.698727</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>lin</th>\n",
       "      <td>-4265.800513</td>\n",
       "      <td>3211.771843</td>\n",
       "      <td>-1.328177</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "               coef           se          t\n",
       "neyman -8506.495361   583.442609 -14.579832\n",
       "fisher  1067.546135   628.438879   1.698727\n",
       "lin    -4265.800513  3211.771843  -1.328177"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reg_adjust(z, y, X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "No bias correction : Much better than OLS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>est</th>\n",
       "      <th>ai_se</th>\n",
       "      <th>boot_se</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1521.376503</td>\n",
       "      <td>1244.056155</td>\n",
       "      <td>862.221261</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           est        ai_se     boot_se\n",
       "0  1521.376503  1244.056155  862.221261"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nn_att(X, z, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "bias correction with OLS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>est</th>\n",
       "      <th>ai_se</th>\n",
       "      <th>or_se</th>\n",
       "      <th>boot_se</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1929.466879</td>\n",
       "      <td>1465.478449</td>\n",
       "      <td>887.049183</td>\n",
       "      <td>845.413394</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           est        ai_se       or_se     boot_se\n",
       "0  1929.466879  1465.478449  887.049183  845.413394"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nn_att(X, z, y, bias_corr_mod=LinearRegression())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "AI SE is now much bigger than boot, potentially because the second term (conditional variance) dominates. Otsu and Rai (2017) SEs look reasonable throughout."
   ]
  }
 ],
 "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
}