Chapter26principalStratification.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chapter 26: Principal Stratification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "matplotlib       : 3.7.2\n",
      "pandas           : 2.0.3\n",
      "matplotlib_inline: 0.1.6\n",
      "statsmodels      : 0.14.0\n",
      "numpy            : 1.24.3\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import statsmodels.api as sm\n",
    "\n",
    "np.random.seed(42)\n",
    "\n",
    "%load_ext watermark\n",
    "%watermark --iversions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Partial ID\n",
    "\n",
    "### SACE with grouped data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.03698057577458183, 0.19404122726930068)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## Chapter 26.4.2\n",
    "## truncation by death example\n",
    "## data from Yang and Small (2016)\n",
    "pi11 = 277 / (277 + 152)\n",
    "pi00 = 109 / (109 + 322)\n",
    "pi10 = 1 - pi11 - pi00\n",
    "\n",
    "## observed means\n",
    "mu11 = 54 / 322\n",
    "mu01 = 59 / 277\n",
    "\n",
    "## bounds on the treatment potential outcomes\n",
    "lb = ((pi11 + pi10) * mu11 - pi10) / pi11\n",
    "ub = ((pi11 + pi10) * mu11) / pi11\n",
    "lb, ub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(-0.17601581411711492, -0.018955162622396077)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## bounds on the sace\n",
    "lb - mu01, ub - mu01"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### SACE with microdata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "## function for SACE with a binary outcome\n",
    "def SACE01_fit(Z, M, Y):\n",
    "    ## summary statistics\n",
    "    pM1 = np.mean(M[Z == 1])\n",
    "    pM0 = np.mean(M[Z == 0])\n",
    "    mu11 = np.mean(Y[(Z == 1) & (M == 1)])\n",
    "    mu01 = np.mean(Y[(Z == 0) & (M == 1)])\n",
    "    ## proporitions of the strata\n",
    "    pi11 = pM0\n",
    "    pi00 = 1 - pM1\n",
    "    pi10 = 1 - pi11 - pi00\n",
    "    ## bounds on the treatment potential outcomes\n",
    "    lb = ((pi11 + pi10) * mu11 - pi10) / pi11\n",
    "    ub = ((pi11 + pi10) * mu11) / pi11\n",
    "    ## bounds on the SACE\n",
    "    return np.array([lb - mu01, ub - mu01])\n",
    "\n",
    "\n",
    "def SACE01(Z, M, Y, nboot=1e3):\n",
    "    from joblib import Parallel, delayed\n",
    "\n",
    "    bounds = SACE01_fit(Z, M, Y)\n",
    "    n = len(Z)\n",
    "\n",
    "    def bootfn(*args):\n",
    "        idx = np.random.choice(n, n, replace=True)\n",
    "        return SACE01_fit(Z[idx], M[idx], Y[idx])\n",
    "\n",
    "    res = Parallel(n_jobs=-1, verbose=0)(delayed(bootfn)() for _ in range(int(nboot)))\n",
    "    res = np.vstack(res)\n",
    "    b_se = np.std(res, axis=0)\n",
    "    # IM CI\n",
    "    l_ci, u_ci = bounds[0] - 1.96 * b_se[0], bounds[1] + 1.96 * b_se[1]\n",
    "    res = pd.DataFrame(\n",
    "        np.c_[bounds, b_se, np.array([l_ci, u_ci])],\n",
    "        columns=[\"est\", \"se\", \"ci\"],\n",
    "        index=[\"lb\", \"ub\"],\n",
    "    )\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "## truncation by death example\n",
    "## data from Yang and Small (2016)\n",
    "Z = np.r_[np.repeat(1, 322 + 109), np.repeat(0, 277 + 152)]\n",
    "\n",
    "M = np.r_[np.repeat(1, 322), np.repeat(0, 109), np.repeat(1, 277), np.repeat(0, 152)]\n",
    "Y = np.r_[\n",
    "    np.repeat(1, 54),\n",
    "    np.repeat(0, 268),\n",
    "    np.repeat(np.nan, 109),\n",
    "    np.repeat(1, 59),\n",
    "    np.repeat(0, 218),\n",
    "    np.repeat(np.nan, 152),\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "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>se</th>\n",
       "      <th>ci</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>lb</th>\n",
       "      <td>-0.176016</td>\n",
       "      <td>0.055382</td>\n",
       "      <td>-0.284565</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>ub</th>\n",
       "      <td>-0.018955</td>\n",
       "      <td>0.036026</td>\n",
       "      <td>0.051657</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "         est        se        ci\n",
       "lb -0.176016  0.055382 -0.284565\n",
       "ub -0.018955  0.036026  0.051657"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "yangsmall = SACE01(Z, M, Y)\n",
    "yangsmall"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Principal Score Methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.linear_model import LogisticRegression\n",
    "from sklearn.utils._testing import ignore_warnings\n",
    "from sklearn.exceptions import ConvergenceWarning\n",
    "\n",
    "\n",
    "@ignore_warnings(category=ConvergenceWarning)\n",
    "def psw(Z, M, Y, X):\n",
    "    pi_10 = np.mean(M[Z == 1])\n",
    "    pi_00 = 1 - pi_10\n",
    "    # ps_10 = sm.Logit(M[Z==1], X[Z == 1,:]).fit(disp = 0).predict(X)\n",
    "    ps_10 = LogisticRegression().fit(X[Z == 1, :], M[Z == 1]).predict_proba(X)[:, 1]\n",
    "    ps_00 = 1 - ps_10\n",
    "    # PCEs 10 and 00\n",
    "    tau_10 = (\n",
    "        np.mean(Y[(Z == 1) & (M == 1)]) - np.mean(Y[Z == 0] * ps_10[Z == 0]) / pi_10\n",
    "    )\n",
    "    tau_00 = (\n",
    "        np.mean(Y[(Z == 1) & (M == 0)]) - np.mean(Y[Z == 0] * ps_00[Z == 0]) / pi_00\n",
    "    )\n",
    "    return np.r_[tau_10, tau_00]\n",
    "\n",
    "\n",
    "def psw_boot(Z, M, Y, X, n_boot=1e3):\n",
    "    from joblib import Parallel, delayed\n",
    "\n",
    "    point_est = psw(Z, M, Y, X)\n",
    "    n = len(Z)\n",
    "\n",
    "    def bootfn(*args):\n",
    "        idx = np.random.choice(n, n, replace=True)\n",
    "        return psw(Z[idx], M[idx], Y[idx], X[idx, :])\n",
    "\n",
    "    res = Parallel(n_jobs=-1, verbose=0)(delayed(bootfn)() for _ in range(int(n_boot)))\n",
    "    res = np.vstack(res)\n",
    "    boot_se = np.std(res, axis=0)\n",
    "    # results\n",
    "    res = pd.DataFrame(\n",
    "        np.c_[point_est, boot_se], columns=[\"est\", \"se\"], index=[\"tau10\", \"tau00\"]\n",
    "    )\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "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>col_0</th>\n",
       "      <th>0.0</th>\n",
       "      <th>1.0</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>row_0</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0.0</th>\n",
       "      <td>299</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1.0</th>\n",
       "      <td>228</td>\n",
       "      <td>372</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "col_0  0.0  1.0\n",
       "row_0          \n",
       "0.0    299    0\n",
       "1.0    228  372"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import formulaic as fm\n",
    "\n",
    "jobsdata = pd.read_csv(\"jobsdata.csv\")\n",
    "X = (\n",
    "    fm.Formula(\"~ sex + age + marital + nonwhite + educ + income -1\")\n",
    "    .get_model_matrix(data=jobsdata)\n",
    "    .values\n",
    ")\n",
    "Z, M, Y = jobsdata[[\"treat\", \"comply\", \"job_seek\"]].values.T\n",
    "pd.crosstab(Z, M)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "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>se</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>tau10</th>\n",
       "      <td>0.167204</td>\n",
       "      <td>0.101708</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>tau00</th>\n",
       "      <td>-0.095306</td>\n",
       "      <td>0.151533</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "            est        se\n",
       "tau10  0.167204  0.101708\n",
       "tau00 -0.095306  0.151533"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "psw_boot(Z, M, Y, X)"
   ]
  }
 ],
 "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
}