docs/Examples/Hernan-Robins-2023.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "629df219",
   "metadata": {},
   "source": [
    "# Hernan & Robins (2023): Causal Inference with Models\n",
    "\n",
    "The following replicates selected examples from the textbook [\"Causal Inference: What If\"](https://www.routledge.com/Causal-Inference-What-If/Hernan-Robins/p/book/9781420076165) by Hernan and Robins. These replications focus on Part II of the textbook. I recommend reading Part I prior to looking through the following code. It is a great, approachable, and freely available resource. \n",
    "\n",
    "Here, we demonstrate application of `delicatessen` for causal inference. Throughout, we use the sandwich variance estimator. This is not described in the textbook, but is an alternative to the bootstrap. Importantly, it is a consistent estimator of the variance that is computationally simpler (in terms of the computers computational effort). `delicatessen` automates the whole procedure, so it is also simpler for us.\n",
    "\n",
    "Broadly, interest is in estimating the average causal effect of stopping smoking (variable name: `qsmk`) on 10-year weight change (variable name: `wt82_71`). If we let $Y^a$ indicate the potential weight change under smoking status $a$, then the average causal effect can be written as\n",
    "$$E[Y^1] - E[Y^0]$$\n",
    "Herafter, we assume that the interest parameter is identified (see the book for what this means). Our focus will be on estimators for this quantity.\n",
    "\n",
    "## Setup\n",
    "\n",
    "Data for these replications is available at the following Harvard School of Public Health [website](https://www.hsph.harvard.edu/miguel-hernan/causal-inference-book/). Data comes from the National Health and Nutrition Examination Survey Data I Epidemiologic Follow-up Study (NHEFS). The NHEFS was jointly initiated by the National Center for Health Statistics and the National Institute on Aging in collaboration with other agencies of the United States Public Health Service. A detailed description of the NHEFS, together with publicly available data sets and documentation, can be found at wwwn.cdc.gov/nchs/nhanes/nhefs/\n",
    "\n",
    "The data set used in the book and this tutorial is a subset of the full NHEFS. First, we will load the data and run some basic variable manipulations. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "23a12509",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Versions\n",
      "NumPy:         1.25.2\n",
      "SciPy:         1.11.2\n",
      "pandas:        1.4.1\n",
      "Delicatessen:  2.1\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import scipy as sp\n",
    "import pandas as pd\n",
    "\n",
    "import delicatessen as deli\n",
    "from delicatessen import MEstimator\n",
    "from delicatessen.estimating_equations import (ee_regression, ee_glm, ee_ipw_msm, \n",
    "                                               ee_gformula, ee_gestimation_snmm)\n",
    "from delicatessen.utilities import inverse_logit, regression_predictions\n",
    "\n",
    "print(\"Versions\")\n",
    "print(\"NumPy:        \", np.__version__)\n",
    "print(\"SciPy:        \", sp.__version__)\n",
    "print(\"pandas:       \", pd.__version__)\n",
    "print(\"Delicatessen: \", deli.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c34711ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv('data/nhefs.csv')\n",
    "df.dropna(subset=['sex', 'age', 'race', 'ht', \n",
    "                  'school', 'alcoholpy', 'smokeintensity'], \n",
    "         inplace=True)\n",
    "\n",
    "# recoding some variables\n",
    "df['inactive'] = np.where(df['active'] == 2, 1, 0)\n",
    "df['no_exercise'] = np.where(df['exercise'] == 2, 1, 0)\n",
    "df['university'] = np.where(df['education'] == 5, 1, 0)\n",
    "\n",
    "# Subsetting only variables of interest\n",
    "df = df[['wt82_71', 'qsmk', 'sex', 'age', 'race', 'wt71', 'wt82', 'ht', \n",
    "         'school', 'alcoholpy', 'smokeintensity', 'smokeyrs', 'smkintensity82_71',\n",
    "         'education', 'exercise', 'active', 'death']]\n",
    "\n",
    "# creating quadratic terms\n",
    "for col in ['age', 'wt71', 'smokeintensity', 'smokeyrs']:\n",
    "    df[col+'_sq'] = df[col] * df[col]\n",
    "\n",
    "df['I'] = 1\n",
    "\n",
    "# Indicator terms\n",
    "df['educ_2'] = np.where(df['education'] == 2, 1, 0)\n",
    "df['educ_3'] = np.where(df['education'] == 3, 1, 0)\n",
    "df['educ_4'] = np.where(df['education'] == 4, 1, 0)\n",
    "df['educ_5'] = np.where(df['education'] == 5, 1, 0)\n",
    "df['exer_1'] = np.where(df['exercise'] == 1, 1, 0)\n",
    "df['exer_2'] = np.where(df['exercise'] == 2, 1, 0)\n",
    "df['active_1'] = np.where(df['active'] == 1, 1, 0)\n",
    "df['active_2'] = np.where(df['active'] == 2, 1, 0)\n",
    "\n",
    "# Interaction terms\n",
    "df['qsmk_smkint'] = df['qsmk'] * df['smokeintensity']\n",
    "\n",
    "# Complete-case data\n",
    "dc = df.dropna(subset=['wt82_71']).copy()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0009679",
   "metadata": {},
   "source": [
    "## Chapter 12: IP weighting and marginal structural models\n",
    "\n",
    "The first estimation approach is inverse probability weighting. Inverse probability weights are defined as \n",
    "$$ \\frac{1}{\\Pr(A=a | W)} $$\n",
    "where $W$ is the set of confounders. We will estimate these weights and then use them to estimate the parameters of a marginal structural model. An example of a marginal structural model is\n",
    "$$ E[Y^a] = \\alpha_0 + \\alpha_1 a $$\n",
    "where $\\alpha$ are the parameters to estimate. Here, $\\alpha_1$ represents the average causal effect. We will estimate the marginal structural model using the observed data and a regression model weighted by the inverse probability weights (see the books for details). \n",
    "\n",
    "### 12.1: The causal question\n",
    "Chapter 12 starts out with estimating the crude association between `qsmk` and `wt82_71` (12.1). We will do this by fitting a linear regression model with `delicatessen`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "278b397c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def psi_ols(theta):\n",
    "    return ee_regression(theta=theta, X=dc[['I', 'qsmk']],\n",
    "                         y=dc['wt82_71'], model='linear')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "00e71d13",
   "metadata": {},
   "outputs": [],
   "source": [
    "estr = MEstimator(psi_ols, init=[0., 0.])\n",
    "estr.estimate(solver='hybr')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "679e5922",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Point: 2.540581454955888\n",
      "95% CI: [1.58620716 3.49495575]\n"
     ]
    }
   ],
   "source": [
    "print(\"Point:\", estr.theta[1])\n",
    "print(\"95% CI:\", estr.confidence_intervals()[1, :])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad503833",
   "metadata": {},
   "source": [
    "The book reports an estimate of 2.5 (95% CI: 1.7, 3.4)\n",
    "\n",
    "You may notice that the confidence interval differs slightly. That is because we are using the sandwich variance (the book does not). While the variance estimators used here and in the book are expected to be asymptotically equivalent in this case, they can produce different results in finite samples as we see here.\n",
    "\n",
    "### 12.2: Estimating inverse probability weights via modeling\n",
    "\n",
    "Now we will estimate the parameters of the marginal structural model using unstabilized inverse probability weights.\n",
    "\n",
    "To do this with `delicatessen`, we will define the corresponding design matrices. Then we will define the stacked estimating equations. Then we will estimate the parameters and covariance using `MEstimator` and present the output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "60e6be04",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Design matrix for the propensity score model\n",
    "W = dc[['I', 'sex', 'race', 'age', 'age_sq', \n",
    "        'educ_2', 'educ_3', 'educ_4', 'educ_5',\n",
    "        'smokeintensity', 'smokeintensity_sq',\n",
    "        'smokeyrs', 'smokeyrs_sq', \n",
    "        'exer_1', 'exer_2', 'active_1', 'active_2',\n",
    "        'wt71', 'wt71_sq']]\n",
    "# Design matrix for the marginal structural model\n",
    "msm = dc[['I', 'qsmk']]\n",
    "# treatment variable\n",
    "a = dc['qsmk']\n",
    "# outcome variable\n",
    "y = dc['wt82_71']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d3143781",
   "metadata": {},
   "outputs": [],
   "source": [
    "def psi_ipw_msm1(theta):\n",
    "    # Dividing parameters into corresponding estimation equations\n",
    "    alpha = theta[0:2]\n",
    "    beta = theta[2:]\n",
    "    \n",
    "    # Estimating the propensity scores\n",
    "    ee_ps = ee_regression(theta=beta,        # Estimate propensity scores\n",
    "                          X=W, y=a,          # ... given observed A,W\n",
    "                          model='logistic')  # ... with logit model\n",
    "    pi = inverse_logit(np.dot(W, beta))      # Get Pr(A = 1 | W)\n",
    "    ipw = 1 / np.where(a == 1, pi, 1-pi)     # Convert to IPW\n",
    "        \n",
    "    # Estimating the MSM using a weighted linear model\n",
    "    ee_msm = ee_regression(theta=alpha,      # MSM parameters \n",
    "                           X=msm, y=y,       # ... observed data  \n",
    "                           model='linear',   # ... with linear model\n",
    "                           weights=ipw)      # ... but weighted by IPW\n",
    "    \n",
    "    # Stacking the estimating equations and returning\n",
    "    return np.vstack([ee_msm, ee_ps])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "574fdabd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# M-estimator\n",
    "init_vals = [0., 0., ] + [0.,]*W.shape[1]\n",
    "estr = MEstimator(psi_ipw_msm1, init=init_vals)\n",
    "estr.estimate(solver='hybr', maxiter=5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4e9976f8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1.77997819 3.44053543]\n",
      "95% CI\n",
      "[[1.35249881 2.48589079]\n",
      " [2.20745757 4.39518006]]\n"
     ]
    }
   ],
   "source": [
    "ci = estr.confidence_intervals()\n",
    "print(estr.theta[0:2])\n",
    "print(\"95% CI\")\n",
    "print(ci[0:2, :].T)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1097cb4d",
   "metadata": {},
   "source": [
    "which are the estimates of the $\\alpha_0$ and $\\alpha_1$ parameters. The book provides the point estimate for `qsmk` ($\\hat{\\alpha}_1$) as 3.4 (95% CI: 2.4, 4.5).\n",
    "\n",
    "The point estimates presented here are the same as the book. However, the confidence intervals differ slightly. The confidence intervals in the book use the 'GEE trick' which provides overly conservative confidence intervals. With `delicatessen` and the sandwich variance, we can estimate the variance *without being overly conservative*. So, the variance reported above would be preferred over the approach described in the book. This highlights the advantage of M-estimators for computation of the variance with nuisance parameters (like the IPW estimator when the propensity score is estimated).\n",
    "\n",
    "Instead of coding this by-hand, we can also use the built-in `ee_ipw_msm` function. This function estimates a marginal structural model using inverse probability weights, as done above. Below is how to apply this functionality"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b8343216",
   "metadata": {},
   "outputs": [],
   "source": [
    "def psi_ipw_msm1a(theta):\n",
    "    # Built-in estimating equation\n",
    "    return ee_ipw_msm(theta, y=y, A=a, W=W, V=msm, \n",
    "                      distribution='normal', \n",
    "                      link='identity')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6326bbd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# M-estimator\n",
    "init_vals = [0., 0., ] + [0.,]*W.shape[1]\n",
    "estr = MEstimator(psi_ipw_msm1a, init=init_vals)\n",
    "estr.estimate(solver='hybr', maxiter=5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a300d352",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1.77997819 3.44053543]\n",
      "95% CI\n",
      "[[1.35249881 2.48589079]\n",
      " [2.20745757 4.39518006]]\n"
     ]
    }
   ],
   "source": [
    "ci = estr.confidence_intervals()\n",
    "print(estr.theta[0:2])\n",
    "print(\"95% CI\")\n",
    "print(ci[0:2, :].T)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac5fb883",
   "metadata": {},
   "source": [
    "As expected, this built-in functionality produces the same results as the by-hand version.\n",
    "\n",
    "### 12.3: Stabilized inverse probability weights\n",
    "\n",
    "Next, we are going to use stabilized weights. The stabilized weights will require us to estimate an additional parameter. We will accomplish this by stacking an estimating equation for that parameter. This extra estimating equation is for an intercept-only model for the probability of `qsmk`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "f6036925",
   "metadata": {},
   "outputs": [],
   "source": [
    "def psi_ipw_msm2(theta):\n",
    "    # Dividing parameters into corresponding estimation equations\n",
    "    alpha = theta[0:2]                 # MSM parameters\n",
    "    gamma = np.array([theta[2], ])     # Numerator parameter\n",
    "    beta = theta[3:]                   # Propensity score parameters\n",
    "    \n",
    "    # Estimating the propensity scores using a logit model\n",
    "    ee_ps = ee_regression(theta=beta,        # Propensity score model\n",
    "                          X=W, y=a,          # ... with observed data\n",
    "                          model='logistic')  # ... and logit model\n",
    "    pi = inverse_logit(np.dot(W, beta))      # Predicted probability of A=1\n",
    "\n",
    "    # Estimating intercept-only for numerator\n",
    "    ee_num = ee_regression(theta=gamma,            # Numerator model\n",
    "                           X=dc[['I']], y=a,       # ... intercept-only\n",
    "                           model='logistic')       # ... logit model\n",
    "    num = inverse_logit(np.dot(dc[['I']], gamma))  # Marginal probability\n",
    "    \n",
    "    # Construct stabilized weights\n",
    "    ipw = np.where(a == 1, num/pi, (1-num)/(1-pi))\n",
    "    \n",
    "    # Estimating the MSM using a weighted linear model\n",
    "    ee_msm = ee_regression(theta=alpha,     # MSM\n",
    "                           X=msm, y=y,      # ... with observed data\n",
    "                           model='linear',  # ... linear \n",
    "                           weights=ipw)     # ... weighted by stabilized\n",
    "    \n",
    "    # Stacking the estimating equations and returning\n",
    "    return np.vstack([ee_msm, ee_num, ee_ps])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "cfb7adc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# M-estimator\n",
    "init_vals = [0., 0., ] + [0., ] + [0.,]*W.shape[1]\n",
    "estr = MEstimator(psi_ipw_msm2, init=init_vals)\n",
    "estr.estimate(solver='hybr', maxiter=5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "eef650ee",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Intercept: 1.77997819053316\n",
      "95% CI: [1.35249875 2.20745763]\n",
      "qsmk: 3.44053542964726\n",
      "95% CI: [2.48589062 4.39518024]\n"
     ]
    }
   ],
   "source": [
    "ci = estr.confidence_intervals()\n",
    "print(\"Intercept:\", estr.theta[0])\n",
    "print(\"95% CI:\", ci[0, :])\n",
    "print(\"qsmk:\", estr.theta[1])\n",
    "print(\"95% CI:\", ci[1, :])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a57fc3a2",
   "metadata": {},
   "source": [
    "The estimate is the same in the previous section. This is expected because as long as the marginal structural model is saturated, the unstabilized and stabilized IPTW should produce the same answer. (note there are some differences out in the smaller decimals places, but this is due to floating point errors, not differences between the estimators).\n",
    "\n",
    "Note: `ee_ipw_msm` does not support the computation of stabilized weights (results are equivalent in this setting, so we can be more computationally efficient by opting for the unstabilized weights).\n",
    "\n",
    "### 12.4: Marginal structural models\n",
    "\n",
    "Now we will consider the IPW estimator for a continuous action. We will look at `smokeintensity` on `wt82_71`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "3420cde5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Restricting data by smoking intensity\n",
    "ds = dc.loc[dc['smokeintensity'] <= 25].copy()\n",
    "\n",
    "# Design matrix for the propensity score model\n",
    "W = ds[['I', 'sex', 'race', 'age', 'age_sq', \n",
    "        'educ_2', 'educ_3', 'educ_4', 'educ_5',\n",
    "        'smokeintensity', 'smokeintensity_sq',\n",
    "        'smokeyrs', 'smokeyrs_sq', \n",
    "        'exer_1', 'exer_2', 'active_1', 'active_2',\n",
    "        'wt71', 'wt71_sq']]\n",
    "# Design matrix for the marginal structural model\n",
    "ds['smkint_sq'] = ds['smkintensity82_71']**2\n",
    "msm = ds[['I', 'smkintensity82_71', 'smkint_sq']]\n",
    "# Treatment array\n",
    "a = ds['smkintensity82_71']\n",
    "# Outcome array\n",
    "y = ds['wt82_71']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "b269e3ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "def psi_ipw_msm3(theta):\n",
    "    # Dividing parameters into corresponding estimation equations\n",
    "    alpha = theta[0:3]               # Marginal structural model\n",
    "    gamma = np.array([theta[3], ])   # Numerator \n",
    "    beta = theta[4:]                 # Propensity score model\n",
    "    div_ps = W.shape[0] - len(beta)  # Divisor for PS SD\n",
    "    div_nm = W.shape[0] - len(gamma) # Divisor for Num SD\n",
    "    \n",
    "    # Estimating the propensity scores using a logit model\n",
    "    ee_ps = ee_regression(theta=beta,           # Generalized PS model \n",
    "                          X=W, y=a,             # ... for observed data\n",
    "                          model='linear')       # ... linear regression\n",
    "    mu = np.dot(W, beta)                        # Predicted values\n",
    "    mu_resid = np.sum((a - mu)**2) / div_ps     # Standard deviation\n",
    "    fAL = sp.stats.norm.pdf(a, mu,              # PDF from normal\n",
    "                            np.sqrt(mu_resid))  # ... with estimates\n",
    "\n",
    "    # Estimating intercept-only for numerator\n",
    "    ee_num = ee_regression(theta=gamma,         # Numerator for stabilized\n",
    "                           X=ds[['I', ]], y=a,  # ... for observed data\n",
    "                           model='linear')      # ... linear regression\n",
    "    num = np.dot(ds[['I', ]], gamma)            # Predicted values\n",
    "    num_resid = np.sum((a - num)**2) / div_nm   # Standard deviation\n",
    "    fA = sp.stats.norm.pdf(a, num,              # PDF from normal\n",
    "                           np.sqrt(num_resid))  # ... with estimates\n",
    "    \n",
    "    # Stabilized weights\n",
    "    ipw = fA / fAL\n",
    "    \n",
    "    # Estimating the MSM using a weighted linear model\n",
    "    ee_msm = ee_regression(theta=alpha,     # Marginal structural model\n",
    "                           X=msm, y=y,      # ... observed data\n",
    "                           model='linear',  # ... linear model\n",
    "                           weights=ipw)     # ... weighted by IPW\n",
    "    \n",
    "    # Stacking the estimating equations and returning\n",
    "    return np.vstack([ee_msm, ee_num, ee_ps])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "2cd2b330",
   "metadata": {},
   "outputs": [],
   "source": [
    "# M-estimator\n",
    "init_vals = [0., 0., 0., ] + [0., ] + [0.,]*W.shape[1]\n",
    "estr = MEstimator(psi_ipw_msm3, init=init_vals)\n",
    "estr.estimate(solver='hybr', tolerance=1e-12, maxiter=5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "8e777bac",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 2.00452474 -0.10898888  0.00269494]\n",
      "95% CI:\n",
      "[[ 1.44058085e+00 -1.66863862e-01 -1.75081112e-03]\n",
      " [ 2.56846862e+00 -5.11138978e-02  7.14069266e-03]]\n"
     ]
    }
   ],
   "source": [
    "ci = estr.confidence_intervals()\n",
    "print(estr.theta[0:3])\n",
    "print(\"95% CI:\")\n",
    "print(ci[0:3, :].T)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2453b34c",
   "metadata": {},
   "source": [
    "The book reports coefficients of: 2.005, −0.109, 0.003. This matches the output shown above.\n",
    "\n",
    "As done in the book, we want to know the weight change for no change in smoking intensity and a +20 in smoking intensity. Rather than add corresponding estimation equations for those parts, we can directly manipulate the point and covariance estimates to get this. The function `regression_predictions` will do this for us given the estimated parameters and the values of interest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "d3e00f31",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No change in smoking\n",
      "2.004524735075991\n",
      "95% CI: [1.44058085 2.56846862]\n",
      "\n",
      "+20 smoking intensity\n",
      "0.9027234481603187\n",
      "95% CI: [-1.49966211  3.305109  ]\n"
     ]
    }
   ],
   "source": [
    "# Creating dataframe for combinations to predict\n",
    "p = pd.DataFrame()\n",
    "p['smkintensity82_71'] = [0, 20]\n",
    "p['smkint_sq'] = [0**2, 20**2]\n",
    "p['I'] = 1\n",
    "vals = np.asarray(p[['I', 'smkintensity82_71', 'smkint_sq']])\n",
    "\n",
    "# Getting predicted values and variance for combinations\n",
    "pred_y = regression_predictions(vals, \n",
    "                                estr.theta[0:3], \n",
    "                                estr.variance[0:3, 0:3])\n",
    "\n",
    "# Displaying results\n",
    "print(\"No change in smoking\")\n",
    "print(pred_y[0, 0])\n",
    "print(\"95% CI:\", pred_y[0, 2:])\n",
    "print()\n",
    "print(\"+20 smoking intensity\")\n",
    "print(pred_y[1, 0])\n",
    "print(\"95% CI:\", pred_y[1, 2:])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e376193d",
   "metadata": {},
   "source": [
    "Again, we get similar results to those reported in the book: 2.0 (95% CI: 1.4, 2.6) and 0.9 (95% CI: −1.7, 3.5). However, our confidence intervals are slightly more narrow. Again this would be expected since we are using a variance estimator that is not overly conservative.\n",
    "\n",
    "Note: `ee_ipw_msm` does not support non-binary treatments. Allowing for generic distributions for the generalized propensity score model is difficult. So, you will need to implement them by-hand (using a similar approach to above).\n",
    "\n",
    "#### Binary outcome\n",
    "\n",
    "We now repeat the process, but for a binary outcome and treatment. We will use a GLM with the binomial distribution and logistic link (this will estimate the causal odds ratio). This is avaible in `ee_glm`. We will also be using `qsmk` again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "5f78c896",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Design matrix for propensity scores\n",
    "W = dc[['I', 'sex', 'race', 'age', 'age_sq', \n",
    "        'educ_2', 'educ_3', 'educ_4', 'educ_5',\n",
    "        'smokeintensity', 'smokeintensity_sq',\n",
    "        'smokeyrs', 'smokeyrs_sq', \n",
    "        'exer_1', 'exer_2', 'active_1', 'active_2',\n",
    "        'wt71', 'wt71_sq']]\n",
    "# Design matrix for marginal structural model\n",
    "msm = dc[['I', 'qsmk']]\n",
    "# Treatment array\n",
    "a = dc['qsmk']\n",
    "# Outcome array\n",
    "y = dc['death']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "a5be66b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def psi_ipw_msm4(theta):\n",
    "    # Dividing parameters into corresponding estimation equations\n",
    "    alpha = theta[0:2]              # Marginal structural model\n",
    "    gamma = np.array([theta[2], ])  # Numerator model\n",
    "    beta = theta[3:]                # Propensity score model\n",
    "    \n",
    "    # Estimating the propensity scores using a logit model\n",
    "    ee_ps = ee_regression(theta=beta,        # Propensity score\n",
    "                          X=W, y=a,          # ... observed data\n",
    "                          model='logistic')  # ... logistic model\n",
    "    pi = inverse_logit(np.dot(W, beta))      # Predicted probability\n",
    "\n",
    "    # Estimating intercept-only for numerator\n",
    "    ee_num = ee_regression(theta=gamma,           # Numerator model\n",
    "                           X=dc[['I']], y=a,      # ... observed data\n",
    "                           model='logistic')      # ... logit model\n",
    "    num = inverse_logit(np.dot(dc[['I']], gamma)) # Predicted prob\n",
    "    \n",
    "    # Stabilized inverse probability weights\n",
    "    ipw = np.where(a == 1, num/pi, (1-num)/(1-pi))\n",
    "    \n",
    "    # Estimating the MSM using a weighted linear model\n",
    "    ee_msm = ee_glm(theta=alpha,               # MSM\n",
    "                    X=msm, y=y,                # ... observed data\n",
    "                    link='logit',              # ... logit link\n",
    "                    distribution='binomial',   # ... binomial dist\n",
    "                    weights=ipw)               # ... weighted\n",
    "    \n",
    "    # Stacking the estimating equations and returning\n",
    "    return np.vstack([ee_msm, ee_num, ee_ps])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "d25c4f4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# M-estimator\n",
    "init_vals = [0., 0., ] + [0., ] + [0.,]*W.shape[1]\n",
    "estr = MEstimator(psi_ipw_msm4, init=init_vals)\n",
    "estr.estimate(solver='hybr', maxiter=5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "5a88a6e6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "qsmk: 1.030577892676402\n",
      "95% CI: [0.78938458 1.34546687]\n"
     ]
    }
   ],
   "source": [
    "ci = estr.confidence_intervals()\n",
    "print(\"qsmk:\", np.exp(estr.theta[1]))\n",
    "print(\"95% CI:\", np.exp(ci[1, :]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "896d2679",
   "metadata": {},
   "source": [
    "The previous results are for the causal odds ratio. Again, they are similar to the book with slight differences in the confidence intervals (i.e., 1.0; 95% CI: 0.8, 1.4).\n",
    "\n",
    "We can replicate this approach using `ee_ipw_msm`. To simplify the process, I am going to use the previous propensity score model parameters as the starting values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "9cae841b",
   "metadata": {},
   "outputs": [],
   "source": [
    "init_ps = list(estr.theta[3:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "e6f83185",
   "metadata": {},
   "outputs": [],
   "source": [
    "def psi_ipw_msm4a(theta):\n",
    "    # Built-in estimating equation\n",
    "    return ee_ipw_msm(theta, y=y, A=a, W=W, V=msm, \n",
    "                      distribution='binomial', \n",
    "                      link='logit')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "3defde1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# M-estimator\n",
    "init_vals = [0., 0., ] + init_ps\n",
    "estr = MEstimator(psi_ipw_msm4a, init=init_vals)\n",
    "estr.estimate(solver='hybr', maxiter=5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "200622ec",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.22527109 1.030578  ]\n",
      "95% CI\n",
      "[[0.19459809 0.78938459]\n",
      " [0.26077884 1.34546713]]\n"
     ]
    }
   ],
   "source": [
    "ci = estr.confidence_intervals()\n",
    "print(np.exp(estr.theta[0:2]))\n",
    "print(\"95% CI\")\n",
    "print(np.exp(ci[0:2, :]).T)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2af90744",
   "metadata": {},
   "source": [
    "Which are the same (exponentiated) coefficients as the previous approach. There is a slight discrepancy you may notice in the confidence intervals, but this is a floating point error resulting from a difference between the stabilized weights (by-hand) and the unstabilized weights (`ee_ipw_msm`).\n",
    "\n",
    "### 12.5: Effect modification and marginal structural models\n",
    "\n",
    "We will now use marginal structural models to study effect measure modification. We will look at effect modification by sex (`sex`) of quitting smoking (`qsmk`) on 10-year weight change (`wt82_71`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "0b317996",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Design matrix for propensity scores\n",
    "W = dc[['I', 'sex', 'race', 'age', 'age_sq', \n",
    "        'educ_2', 'educ_3', 'educ_4', 'educ_5',\n",
    "        'smokeintensity', 'smokeintensity_sq',\n",
    "        'smokeyrs', 'smokeyrs_sq', \n",
    "        'exer_1', 'exer_2', 'active_1', 'active_2',\n",
    "        'wt71', 'wt71_sq']]\n",
    "# Design matrix for marginal structural model\n",
    "dc['qsmk_sex'] = dc['qsmk']*dc['sex']\n",
    "msm = dc[['I', 'qsmk', 'sex', 'qsmk_sex']]\n",
    "# Treatment array\n",
    "a = dc['qsmk']\n",
    "# Outcome array\n",
    "y = dc['wt82_71']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "b94f358c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def psi_ipw_msm5(theta):\n",
    "    # Dividing parameters into corresponding estimation equations\n",
    "    alpha = theta[0:4]              # Marginal structural model\n",
    "    gamma = np.array([theta[4], ])  # Numerator parameter\n",
    "    beta = theta[5:]                # Propensity score\n",
    "    \n",
    "    # Estimating the propensity scores using a logit model\n",
    "    ee_ps = ee_regression(theta=beta,        # Propensity score\n",
    "                          X=W, y=a,          # ... observed data\n",
    "                          model='logistic')  # ... logit model\n",
    "    pi = inverse_logit(np.dot(W, beta))      # Predicted prob\n",
    "\n",
    "    # Estimating intercept-only for numerator\n",
    "    ee_num = ee_regression(theta=gamma,            # Numerator model\n",
    "                           X=dc[['I']], y=a,       # ... observed data\n",
    "                           model='logistic')       # ... logit model\n",
    "    num = inverse_logit(np.dot(dc[['I']], gamma))  # Predicted prob\n",
    "    \n",
    "    # Stabilized inverse probability weights\n",
    "    ipw = np.where(a == 1, num/pi, (1-num)/(1-pi))\n",
    "    \n",
    "    # Estimating the MSM using a weighted linear model\n",
    "    ee_msm = ee_regression(theta=alpha,     # Marginal structural model\n",
    "                           X=msm, y=y,      # ... observed data\n",
    "                           model='linear',  # ... linear model\n",
    "                           weights=ipw)     # ... weighted\n",
    "    \n",
    "    # Stacking the estimating equations and returning\n",
    "    return np.vstack([ee_msm, ee_num, ee_ps])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "8b3d0d83",
   "metadata": {},
   "outputs": [],
   "source": [
    "# M-estimator\n",
    "init_vals = [0., 0., 0., 0., ] + [0., ] + [0.,]*W.shape[1]\n",
    "estr = MEstimator(psi_ipw_msm5, init=init_vals)\n",
    "estr.estimate(solver='hybr', maxiter=5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "ff308a62",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 1.78444688  3.52197763 -0.00872478 -0.15947852]\n",
      "95% CI\n",
      "[[ 1.19225612  2.28240175 -0.88125781 -2.15657223]\n",
      " [ 2.37663763  4.76155352  0.86380824  1.83761518]]\n"
     ]
    }
   ],
   "source": [
    "ci = estr.confidence_intervals()\n",
    "print(estr.theta[0:4])\n",
    "print(\"95% CI\")\n",
    "print(ci[0:4, :].T)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f208c5a9",
   "metadata": {},
   "source": [
    "While not reported in the book, other online references report the folloiwing coefficients: 1.7844,\t3.5220, -0.0087, -0.1595.\n",
    "\n",
    "As before, we can replicate this result using the built-in `ee_ipw_msm` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "b351bfac",
   "metadata": {},
   "outputs": [],
   "source": [
    "def psi_ipw_msm5a(theta):\n",
    "    # Built-in estimating equation\n",
    "    return ee_ipw_msm(theta, y=y, A=a, W=W, V=msm, \n",
    "                      distribution='normal', \n",
    "                      link='identity')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "fd07dd28",
   "metadata": {},
   "outputs": [],
   "source": [
    "# M-estimator\n",
    "init_vals = [0., ]*msm.shape[1] + init_ps\n",
    "estr = MEstimator(psi_ipw_msm5a, init=init_vals)\n",
    "estr.estimate(solver='hybr', maxiter=5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "a5c18154",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 1.78444688  3.52197763 -0.00872478 -0.15947852]\n",
      "95% CI\n",
      "[[ 1.19225614  2.28240246 -0.88125778 -2.15657182]\n",
      " [ 2.37663761  4.7615528   0.86380821  1.83761477]]\n"
     ]
    }
   ],
   "source": [
    "ci = estr.confidence_intervals()\n",
    "print(estr.theta[0:4])\n",
    "print(\"95% CI\")\n",
    "print(ci[0:4, :].T)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f64bd4ee",
   "metadata": {},
   "source": [
    "which again replicates the by-hand results.\n",
    "\n",
    "### 12.6: Censoring and missing data\n",
    "\n",
    "To conclude, we will now consider the missing outcomes that were ignored earlier. To do this, we will use stabilized inverse probability of missingness weights (IPCW in the book). As we have done so many times already, we will use `delicatessen` to stack additional nuisance models together."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "2eef5bcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Design matrix for propensity score model\n",
    "W = df[['I', 'sex', 'race', 'age', 'age_sq', \n",
    "        'educ_2', 'educ_3', 'educ_4', 'educ_5',\n",
    "        'smokeintensity', 'smokeintensity_sq',\n",
    "        'smokeyrs', 'smokeyrs_sq', \n",
    "        'exer_1', 'exer_2', 'active_1', 'active_2',\n",
    "        'wt71', 'wt71_sq']]\n",
    "# Design matrix for missing model\n",
    "X = df[['I', 'qsmk', 'sex', 'race', 'age', 'age_sq', \n",
    "        'educ_2', 'educ_3', 'educ_4', 'educ_5',\n",
    "        'smokeintensity', 'smokeintensity_sq',\n",
    "        'smokeyrs', 'smokeyrs_sq', \n",
    "        'exer_1', 'exer_2', 'active_1', 'active_2',\n",
    "        'wt71', 'wt71_sq']]\n",
    "# Design matrix for marginal structural model\n",
    "msm = df[['I', 'qsmk']]\n",
    "# Treatment array\n",
    "a = df['qsmk']\n",
    "# Outcome array\n",
    "y = df['wt82_71']\n",
    "# Missing indicator\n",
    "r = np.where(df['wt82_71'].isna(), 0, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "7fb0a19b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def psi_ipw_msm6(theta):\n",
    "    # Dividing parameters up for their corresponding estimation equations\n",
    "    alpha = theta[0:2]                # MSM\n",
    "    gamma_n = np.array([theta[2], ])  # Numerator PS\n",
    "    beta_n = np.array([theta[3], ])   # Numerator MW\n",
    "    gamma_d = theta[4:4+W.shape[1]]   # Propensity score\n",
    "    beta_d = theta[4+W.shape[1]:]     # Missing model\n",
    "    \n",
    "    # Estimating the propensity scores using a logit model\n",
    "    ee_ps = ee_regression(theta=gamma_d,      # Propensity score\n",
    "                          X=W, y=a,           # ... observed data\n",
    "                          model='logistic')   # ... logit model\n",
    "    pi_a = inverse_logit(np.dot(W, gamma_d))  # Predicted prob\n",
    "\n",
    "    # Estimating intercept-only for numerator of IPTW\n",
    "    ee_num = ee_regression(theta=gamma_n,     # Numerator\n",
    "                           X=df[['I']], y=a,  # ... observed data\n",
    "                           model='logistic')  # ... logit model\n",
    "    num_a = inverse_logit(np.dot(df[['I']], gamma_n))\n",
    "\n",
    "    # Estimating the missing scores using a logit model\n",
    "    ee_ms = ee_regression(theta=beta_d,      # Missing score\n",
    "                          X=X, y=r,          # ... observed data\n",
    "                          model='logistic')  # ... logit model\n",
    "    pi_m = inverse_logit(np.dot(X, beta_d))  # Predicted prob\n",
    "\n",
    "    # Estimating intercept-only for numerator of IPMW\n",
    "    ee_sms = ee_regression(theta=beta_n,      # Numerator\n",
    "                           X=df[['I']], y=r,  # ... observed data\n",
    "                           model='logistic')  # ... logit model\n",
    "    num_m = inverse_logit(np.dot(df[['I']], beta_n))\n",
    "    \n",
    "    # Stabilized inverse probability weights\n",
    "    iptw = np.where(a == 1, num_a/pi_a, (1-num_a)/(1-pi_a))\n",
    "    ipmw = np.where(r == 1, num_m/pi_m, 0)\n",
    "    ipw = iptw*ipmw\n",
    "    \n",
    "    # Estimating the MSM using a weighted linear model\n",
    "    ee_msm = ee_regression(theta=alpha,     # MSM\n",
    "                           X=msm, y=y,      # ... observed data\n",
    "                           model='linear',  # ... linear model\n",
    "                           weights=ipw)     # ... weighted\n",
    "    # Setting rows with missing Y's as zero (no contribution)\n",
    "    ee_msm = np.nan_to_num(ee_msm, copy=False, nan=0.)\n",
    "    \n",
    "    # Stacking the estimating equations and returning\n",
    "    return np.vstack([ee_msm, ee_num, ee_sms, ee_ps, ee_ms])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "6838a60a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# M-estimator\n",
    "init_vals = [0., 0., ] + [0., 0., ] + [0.,]*W.shape[1] + [0.,]*X.shape[1]\n",
    "estr = MEstimator(psi_ipw_msm6, init=init_vals)\n",
    "estr.estimate(solver='hybr', maxiter=5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "c1914aad",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1.66199003 3.49649333]\n",
      "95% CI\n",
      "[[1.22590804 2.54496958]\n",
      " [2.09807203 4.44801708]]\n"
     ]
    }
   ],
   "source": [
    "ci = estr.confidence_intervals()\n",
    "print(estr.theta[0:2])\n",
    "print(\"95% CI\")\n",
    "print(ci[0:2, :].T)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c9a3721",
   "metadata": {},
   "source": [
    "Here, the book reports 3.5 (95% CI: 2.5, 4.5).\n",
    "\n",
    "Again, we will replicate this using `ee_ipw_msm` instead. To incorporate the missingness weights, we will fit a separate model and then use the optional `weights` argument. Below is how this can be done"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "73faacf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def psi_ipw_msm6a(theta):\n",
    "    # Separating parameters out\n",
    "    alpha = theta[:2+W.shape[1]]  # MSM & PS\n",
    "    gamma = theta[2+W.shape[1]:]  # Missing score\n",
    "    \n",
    "    # Estimating equation for IPMW\n",
    "    ee_ms = ee_regression(theta=gamma,       # Missing score\n",
    "                          X=X, y=r,          # ... observed data\n",
    "                          model='logistic')  # ... logit model\n",
    "    pi_m = inverse_logit(np.dot(X, gamma))   # Predicted prob\n",
    "    ipmw = r / pi_m                          # Missing weights\n",
    "\n",
    "    # Estimating equations for MSM and PS\n",
    "    ee_msm = ee_ipw_msm(theta=alpha, y=y, A=a, W=W, V=msm, \n",
    "                        distribution='normal', \n",
    "                        link='identity',\n",
    "                        weights=ipmw)\n",
    "    # Setting rows with missing Y's as zero (no contribution)\n",
    "    ee_msm = np.nan_to_num(ee_msm, copy=False, nan=0.)\n",
    "    \n",
    "    # Return the stacked estimating equations\n",
    "    return np.vstack([ee_msm, ee_ms])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "3b638089",
   "metadata": {},
   "outputs": [],
   "source": [
    "# M-estimator\n",
    "init_vals = [0., ]*msm.shape[1] + init_ps + [0., ]*X.shape[1]\n",
    "estr = MEstimator(psi_ipw_msm6a, init=init_vals)\n",
    "estr.estimate(solver='hybr', maxiter=5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "696754ca",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1.66199003 3.49649333]\n",
      "95% CI\n",
      "[[1.22590805 2.54496955]\n",
      " [2.09807201 4.4480171 ]]\n"
     ]
    }
   ],
   "source": [
    "ci = estr.confidence_intervals()\n",
    "print(estr.theta[0:2])\n",
    "print(\"95% CI\")\n",
    "print(ci[0:2, :].T)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eae99648",
   "metadata": {},
   "source": [
    "Again, we are able to replicate the by-hand implementation. This concludes chapter 12.\n",
    "\n",
    "## Chapter 13: Standardization and the Parametric G-Formula\n",
    "\n",
    "For Chapter 13, the book reviews the g-formula. Unlike IPW, the g-formula relies on modeling the outcome process. The g-computation algorithm estimator is \n",
    "$$ \\hat{E}[Y^a] = n^{-1} \\sum_{i=1}^n m(a, W_i; \\beta) $$\n",
    "where $m$ is a statistical model for $E[Y | A, W]$ and $\\beta$ are the parameters defining the model. This version is slightly different from the standardization form given in the book, but it is equivalent. Broadly, we apply the g-computation algorithm via (1) estimate an outcome model, (2) predict the outcomes had everyone been assigned $a$, and (3) compute the mean of those predictions. \n",
    "\n",
    "### 13.2: Estimating the mean outcome via modeling\n",
    "\n",
    "First, we will fit a linear model for `wt82_71` conditional on `qsmk` and the set of confounding variables. To begin, we will ignore the missing outcomes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "0f92d4b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Design matrix for outcome model\n",
    "X = dc[['I', 'qsmk', 'qsmk_smkint',\n",
    "        'sex', 'race', 'age', 'age_sq', \n",
    "        'educ_2', 'educ_3', 'educ_4', 'educ_5',\n",
    "        'smokeintensity', 'smokeintensity_sq',\n",
    "        'smokeyrs', 'smokeyrs_sq', \n",
    "        'exer_1', 'exer_2', 'active_1', 'active_2',\n",
    "        'wt71', 'wt71_sq']]\n",
    "# Outcome array\n",
    "y = dc['wt82_71']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "f7b43bc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def psi_regression(theta):    \n",
    "    # Estimating the linear outcome model\n",
    "    return ee_regression(theta=theta, \n",
    "                         X=X, y=y, \n",
    "                         model='linear')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "2c09a41f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# M-estimator\n",
    "init_vals = [0., ]*X.shape[1]\n",
    "estr = MEstimator(psi_regression, init=init_vals)\n",
    "estr.estimate(solver='hybr', maxiter=5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "b7216d18",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Regression coefficients\n",
      "[-1.58816566e+00  2.55959409e+00  4.66628395e-02 -1.43027166e+00\n",
      "  5.60109604e-01  3.59635262e-01 -6.10095538e-03  7.90444028e-01\n",
      "  5.56312403e-01  1.49156952e+00 -1.94977036e-01  4.91364717e-02\n",
      " -9.90651192e-04  1.34368569e-01 -1.86642949e-03  2.95975357e-01\n",
      "  3.53912774e-01 -9.47569487e-01 -2.61377894e-01  4.55017905e-02\n",
      " -9.65325105e-04]\n"
     ]
    }
   ],
   "source": [
    "print(\"Regression coefficients\")\n",
    "print(estr.theta)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ab00df1",
   "metadata": {},
   "source": [
    "To ease application of later steps, we are going to save a list of these optimized values for starting values of the root-finding procedure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "666b5823",
   "metadata": {},
   "outputs": [],
   "source": [
    "init_reg = list(estr.theta)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3a8fcc8",
   "metadata": {},
   "source": [
    "### 13.3: Standardizing the mean outcome to the confounder distribution\n",
    "\n",
    "Now we can apply the g-computation algorithm (a way to evaluate the g-formula). To do this, we are going to create a copy of our data set. In that copy we are going to set `qsmk=1` for all observations. We will then save the design matrix. We will then repeat this process for `qsmk=0`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "b3c91754",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy of the data that we will updated qsmk in\n",
    "dca = dc.copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "0181be8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting qsmk to 1\n",
    "dca['qsmk'] = 1\n",
    "dca['qsmk_smkint'] = dca['qsmk'] * dca['smokeintensity']\n",
    "# Design matrix from qsmk=1 data\n",
    "X1 = dca[['I', 'qsmk', 'qsmk_smkint',\n",
    "          'sex', 'race', 'age', 'age_sq', \n",
    "          'educ_2', 'educ_3', 'educ_4', 'educ_5',\n",
    "          'smokeintensity', 'smokeintensity_sq',\n",
    "          'smokeyrs', 'smokeyrs_sq', \n",
    "          'exer_1', 'exer_2', 'active_1', 'active_2',\n",
    "          'wt71', 'wt71_sq']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "5409cd48",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting qsmk to 0\n",
    "dca['qsmk'] = 0\n",
    "dca['qsmk_smkint'] = dca['qsmk'] * dca['smokeintensity']\n",
    "# Design matrix from qsmk=0 data\n",
    "X0 = dca[['I', 'qsmk', 'qsmk_smkint',\n",
    "          'sex', 'race', 'age', 'age_sq', \n",
    "          'educ_2', 'educ_3', 'educ_4', 'educ_5',\n",
    "          'smokeintensity', 'smokeintensity_sq',\n",
    "          'smokeyrs', 'smokeyrs_sq', \n",
    "          'exer_1', 'exer_2', 'active_1', 'active_2',\n",
    "          'wt71', 'wt71_sq']]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56d5b237",
   "metadata": {},
   "source": [
    "Now, we will use the design matrices to generate predicted outcomes (pseudo-outcomes) for each individual. The average causal effect can then be calculated from those pseudo-outcomes. \n",
    "\n",
    "The following is the corresponding estimating equations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "05d5600d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def psi_gcomp1(theta):    \n",
    "    # Dividing parameters into corresponding estimation equations\n",
    "    rd, r1, r0 = theta[0], theta[1], theta[2]   # Causal means\n",
    "    beta = theta[3:]                            # Outcome model\n",
    "    \n",
    "    # Estimating the linear model\n",
    "    ee_reg = ee_regression(theta=beta,      # Outcome model\n",
    "                           X=X, y=y,        # ... observed data\n",
    "                           model='linear')  # ... linear model\n",
    "    \n",
    "    # Generating pseudo-outcome using the model\n",
    "    y1hat = np.dot(X1, beta)  # Predicted Y when qsmk=1\n",
    "    y0hat = np.dot(X0, beta)  # Predicted Y when qsmk=0\n",
    "    \n",
    "    # Causal means\n",
    "    ee_r1 = y1hat - r1   # Causal mean for qsmk=1\n",
    "    ee_r0 = y0hat - r0   # Causal mean for qsmk=0\n",
    "    \n",
    "    # Average causal effect\n",
    "    ee_rd = np.ones(y.shape[0]) * ((r1 - r0) - rd)\n",
    "    \n",
    "    # Stacking the estimating equations and returning\n",
    "    return np.vstack([ee_rd, ee_r1, ee_r0, ee_reg])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "6b5ed64e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# M-estimator\n",
    "init_vals = [0., 0., 0.,] + init_reg\n",
    "estr = MEstimator(psi_gcomp1, init=init_vals)\n",
    "estr.estimate(solver='hybr', maxiter=5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "24e1fab9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[3.5173742  5.27358732 1.75621312]\n",
      "95% CI\n",
      "[[2.58132997 4.42099134 1.33029619]\n",
      " [4.45341844 6.1261833  2.18213004]]\n"
     ]
    }
   ],
   "source": [
    "ci = estr.confidence_intervals()\n",
    "print(estr.theta[0:3])\n",
    "print(\"95% CI\")\n",
    "print(ci[0:3, :].T)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a53cab09",
   "metadata": {},
   "source": [
    "The first estimate is for the average causal effect (the second are the causal means under all quit smoking and all don't quit smoking). Here, the confidence intervals match the book, but we were able to avoid the computational complexity of the bootstrap (another advantage of the sandwich variance estimator). \n",
    "\n",
    "In the book, they report 3.5 (95% CI: 2.6, 4.5). The confidence interval in the book was generated via the bootstrap, so there is potential for random error in the estimates.\n",
    "\n",
    "Rather than implement g-computation by-hand, we can also use the built-in estimating equations. Here is an example of that"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "0ae94f30",
   "metadata": {},
   "outputs": [],
   "source": [
    "def psi_gcomp1a(theta):\n",
    "    # Built-in g-formula estimating equation\n",
    "    return ee_gformula(theta, \n",
    "                       y=y, X=X,\n",
    "                       X1=X1, X0=X0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "18322f47",
   "metadata": {},
   "outputs": [],
   "source": [
    "# M-estimator\n",
    "init_vals = [0., 0., 0.,] + init_reg\n",
    "estr = MEstimator(psi_gcomp1a, init=init_vals)\n",
    "estr.estimate(solver='hybr', maxiter=5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "c3e0135b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[3.5173742  5.27358732 1.75621312]\n",
      "95% CI\n",
      "[[2.58132997 4.42099134 1.33029619]\n",
      " [4.45341844 6.1261833  2.18213004]]\n"
     ]
    }
   ],
   "source": [
    "ci = estr.confidence_intervals()\n",
    "print(estr.theta[0:3])\n",
    "print(\"95% CI\")\n",
    "print(ci[0:3, :].T)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da42b104",
   "metadata": {},
   "source": [
    "which provides the same answers (as we would expect!)\n",
    "\n",
    "## Chapter 14: G-Estimation of Structural Nested Models\n",
    "\n",
    "G-estimation differs from the previous approaches in what parameter it targets. Rather than the average causal effect, we will estimate the parameters of a structural nested model. The structural nested mean model is\n",
    "$$ E[Y^a - Y^{a=0} | A=a, W] = \\varphi_0 a$$\n",
    "Here, $\\varphi_0$ represents the difference by $a$ (which is equivalent to average causal effect). However, we can also study effect measure modification easily with structural nested models. Consider\n",
    "$$ E[Y^a - Y^{a=0} | A=a, W] = \\varphi_0 a + \\varphi_1 a V$$\n",
    "where $V \\in W$. Therefore, the structural nested model described effect measure modification by $V$. Importantly, structural nested models assume that we have correctly specified this model (hence the difference in interest parameters that occurs between methods).\n",
    "\n",
    "G-estimation is a bit less straightforward to understand compare to the previous methods. Note that we will be using the propensity score for estimation and we still assume the same identification conditions. See the book for a much more in-depth discussion.\n",
    "\n",
    "### 14.5 G-estimation\n",
    "\n",
    "For solve our g-estimator, we are going to use a different approach than the one described in the main text of the book. Instead, we are going to use the approach (the estimating equations) described in Technical Point 14.2. As mentioned there, this approach is equivalent to the process described in the main text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "11ecd3fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Design matrix for propensity scores\n",
    "W = np.asarray(df[['I', 'sex', 'race', 'age', 'age_sq', \n",
    "                   'educ_2', 'educ_3', 'educ_4', 'educ_5',\n",
    "                   'smokeintensity', 'smokeintensity_sq',\n",
    "                   'smokeyrs', 'smokeyrs_sq', \n",
    "                   'exer_1', 'exer_2', 'active_1', 'active_2',\n",
    "                   'wt71', 'wt71_sq']])\n",
    "# Design matrix for missing model\n",
    "X = np.asarray(df[['I', 'qsmk', 'sex', 'race', 'age', 'age_sq', \n",
    "                   'educ_2', 'educ_3', 'educ_4', 'educ_5',\n",
    "                   'smokeintensity', 'smokeintensity_sq',\n",
    "                   'smokeyrs', 'smokeyrs_sq', \n",
    "                   'exer_1', 'exer_2', 'active_1', 'active_2',\n",
    "                   'wt71', 'wt71_sq']])\n",
    "# Design matrix for structural nested model\n",
    "snm = np.asarray(df[['I', ]])\n",
    "# Treatment array\n",
    "a = np.asarray(df['qsmk'])\n",
    "# Outcome array\n",
    "y = np.asarray(df['wt82_71'])\n",
    "# Missing indicator\n",
    "r = np.where(df['wt82_71'].isna(), 0, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "0e563f05",
   "metadata": {},
   "outputs": [],
   "source": [
    "def psi_snm1(theta):\n",
    "    # Dividing parameters into corresponding estimation equations\n",
    "    alpha = theta[:snm.shape[1]]             # SNM parameters\n",
    "    beta = theta[snm.shape[1]:               # Propensity score\n",
    "                 snm.shape[1]+W.shape[1]]\n",
    "    gamma = theta[snm.shape[1]+W.shape[1]:]  # Missing score\n",
    "    \n",
    "    # Estimating equation for IPMW\n",
    "    ee_ms = ee_regression(theta=gamma,       # Missing score\n",
    "                          X=X, y=r,          # ... observed data\n",
    "                          model='logistic')  # ... logit model\n",
    "    pi_m = inverse_logit(np.dot(X, gamma))   # Predicted prob\n",
    "    ipmw = r / pi_m                          # Missing weights\n",
    "\n",
    "    # Estimating equations for PS    \n",
    "    ee_log = ee_regression(theta=beta,        # Propensity score\n",
    "                           X=W, y=a,          # ... observed data\n",
    "                           model='logistic',  # ... logit model\n",
    "                           weights=ipmw)      # ... weighted\n",
    "    pi = inverse_logit(np.dot(W, beta))       # Predicted prob\n",
    "    \n",
    "    # H(psi) equation for linear models\n",
    "    h_psi = y - np.dot(snm*a[:, None], alpha)\n",
    "\n",
    "    # Estimating equation for the structural nested mean model\n",
    "    ee_snm = ipmw*(h_psi[:, None] * (a - pi)[:, None] * snm).T\n",
    "    # Setting rows with missing Y's as zero (no contribution)\n",
    "    ee_snm = np.nan_to_num(ee_snm, copy=False, nan=0.)\n",
    "\n",
    "    return np.vstack([ee_snm, ee_log, ee_ms])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "6ac52feb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# M-estimator\n",
    "init_vals = [0., ] + init_ps + [0., ]*X.shape[1]\n",
    "estr = MEstimator(psi_snm1, init=init_vals)\n",
    "estr.estimate(solver='hybr', maxiter=5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "25663d6a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[3.4458988]\n",
      "95% CI\n",
      "[2.52709776 4.36469984]\n"
     ]
    }
   ],
   "source": [
    "ci = estr.confidence_intervals()\n",
    "print(estr.theta[0:1])\n",
    "print(\"95% CI\")\n",
    "print(ci[0, :])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1a79513",
   "metadata": {},
   "source": [
    "The book reported 3.4 (95% CI: 2.5, 4.5). This confidence interval was generated by inverting the test results. As such, there is expected differences (different estimators are being used). As in the inverse probability weighting examples, these confidence intervals are expected due to be conservative due to the use of the 'GEE trick'. \n",
    "\n",
    "Now consider how we can use the built-in g-estimation functionality, `ee_gestimation_snmm`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "9df3779e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def psi_snm1a(theta):\n",
    "    # Dividing parameters into corresponding estimation equations\n",
    "    alpha = theta[:snm.shape[1]+W.shape[1]]  # SNM and PS\n",
    "    gamma = theta[snm.shape[1]+W.shape[1]:]  # Missing scores\n",
    "    \n",
    "    # Estimating equation for IPMW\n",
    "    ee_ms = ee_regression(theta=gamma,       # Missing score\n",
    "                          X=X, y=r,          # ... observed data\n",
    "                          model='logistic')  # ... logit model \n",
    "    pi_m = inverse_logit(np.dot(X, gamma))   # Predicted prob\n",
    "    ipmw = r / pi_m                          # Missing weights\n",
    "\n",
    "    # Estimating equations for PS    \n",
    "    ee_snm = ee_gestimation_snmm(theta=alpha,\n",
    "                                 y=y, A=a, \n",
    "                                 W=W, V=snm,\n",
    "                                 weights=ipmw)\n",
    "    # Setting rows with missing Y's as zero (no contribution)\n",
    "    ee_snm = np.nan_to_num(ee_snm, copy=False, nan=0.)\n",
    "\n",
    "    return np.vstack([ee_snm, ee_ms])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "64e33d42",
   "metadata": {},
   "outputs": [],
   "source": [
    "# M-estimator\n",
    "init_vals = [0., ] + init_ps + [0., ]*X.shape[1]\n",
    "estr = MEstimator(psi_snm1a, init=init_vals)\n",
    "estr.estimate(solver='hybr', maxiter=5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "d0ebdbfd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[3.4458988]\n",
      "95% CI\n",
      "[2.52709776 4.36469984]\n"
     ]
    }
   ],
   "source": [
    "ci = estr.confidence_intervals()\n",
    "print(estr.theta[0:1])\n",
    "print(\"95% CI\")\n",
    "print(ci[0, :])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "589dfe25",
   "metadata": {},
   "source": [
    "The structural nested model parameters match between implementations, as expected.\n",
    "\n",
    "### 14.6: Structural nested models with two or more parameters\n",
    "\n",
    "We now adapt the previous code to consider more than 2 parameters in the structural nested model. Thankfully, this is pretty straightforward for how we coded the estimating equations since `snm` is a global object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "ba0ca7c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "snm = np.asarray(df[['I', 'smokeintensity']])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "4849b70e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def psi_snm2(theta):\n",
    "    # Dividing parameters into corresponding estimation equations\n",
    "    alpha = theta[:snm.shape[1]]             # SNM parameters\n",
    "    beta = theta[snm.shape[1]:               # Propensity score\n",
    "                 snm.shape[1]+W.shape[1]]\n",
    "    gamma = theta[snm.shape[1]+W.shape[1]:]  # Missing score\n",
    "    \n",
    "    # Estimating equation for IPMW\n",
    "    ee_ms = ee_regression(theta=gamma,       # Missing score\n",
    "                          X=X, y=r,          # ... observed data\n",
    "                          model='logistic')  # ... logit model\n",
    "    pi_m = inverse_logit(np.dot(X, gamma))   # Predicted prob\n",
    "    ipmw = r / pi_m                          # Missing weights\n",
    "\n",
    "    # Estimating equations for PS    \n",
    "    ee_log = ee_regression(theta=beta,        # Propensity score\n",
    "                           X=W, y=a,          # ... observed data\n",
    "                           model='logistic',  # ... logit model\n",
    "                           weights=ipmw)      # ... weighted\n",
    "    pi = inverse_logit(np.dot(W, beta))       # Predicted prob\n",
    "    \n",
    "    # H(psi) equation for linear models\n",
    "    h_psi = y - np.dot(snm*a[:, None], alpha)\n",
    "\n",
    "    # Estimating equation for the structural nested mean model\n",
    "    ee_snm = ipmw*(h_psi[:, None] * (a - pi)[:, None] * snm).T\n",
    "    # Setting rows with missing Y's as zero (no contribution)\n",
    "    ee_snm = np.nan_to_num(ee_snm, copy=False, nan=0.)\n",
    "\n",
    "    return np.vstack([ee_snm, ee_log, ee_ms])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "744e9a11",
   "metadata": {},
   "outputs": [],
   "source": [
    "# M-estimator\n",
    "init_vals = [0., 0., ] + init_ps + [0., ]*X.shape[1]\n",
    "estr = MEstimator(psi_snm2, init=init_vals)\n",
    "estr.estimate(solver='hybr', maxiter=5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "9f097670",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2.85947039 0.03004128]\n",
      "95% CI\n",
      "[[ 1.0291672  -0.05951923]\n",
      " [ 4.68977357  0.11960179]]\n"
     ]
    }
   ],
   "source": [
    "ci = estr.confidence_intervals()\n",
    "print(estr.theta[0:2])\n",
    "print(\"95% CI\")\n",
    "print(ci[:2, :].T)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f413273c",
   "metadata": {},
   "source": [
    "which provides the same results as the book (2.86 and 0.03). Unlike the book, we provide confidence intervals. Inverting the test is not simple in this case, and the book does not use a bootstrapping procedure. Our confidence intervals are computed using the sandwich variance estimator.\n",
    "\n",
    "Again, we can apply the built-in estimating equations for g-estimation. In truth, the new estimating equation is the same as the previous g-estimation version (since `snm` is a global object as called). Regardless, we provide again here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "6092fe6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def psi_snm2a(theta):\n",
    "    # Dividing parameters into corresponding estimation equations\n",
    "    alpha = theta[:snm.shape[1]+W.shape[1]]  # SNM and PS\n",
    "    gamma = theta[snm.shape[1]+W.shape[1]:]  # Missing scores\n",
    "    \n",
    "    # Estimating equation for IPMW\n",
    "    ee_ms = ee_regression(theta=gamma,       # Missing score\n",
    "                          X=X, y=r,          # ... observed data\n",
    "                          model='logistic')  # ... logit model \n",
    "    pi_m = inverse_logit(np.dot(X, gamma))   # Predicted prob\n",
    "    ipmw = r / pi_m                          # Missing weights\n",
    "\n",
    "    # Estimating equations for PS    \n",
    "    ee_snm = ee_gestimation_snmm(theta=alpha,\n",
    "                                 y=y, A=a, \n",
    "                                 W=W, V=snm,\n",
    "                                 weights=ipmw)\n",
    "    # Setting rows with missing Y's as zero (no contribution)\n",
    "    ee_snm = np.nan_to_num(ee_snm, copy=False, nan=0.)\n",
    "\n",
    "    return np.vstack([ee_snm, ee_ms])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "019d5c2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# M-estimator\n",
    "init_vals = [0., 0., ] + init_ps + [0., ]*X.shape[1]\n",
    "estr = MEstimator(psi_snm2a, init=init_vals)\n",
    "estr.estimate(solver='hybr', maxiter=5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "12cb56e3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2.85947039 0.03004128]\n",
      "95% CI\n",
      "[[ 1.0291672  -0.05951923]\n",
      " [ 4.68977357  0.11960179]]\n"
     ]
    }
   ],
   "source": [
    "ci = estr.confidence_intervals()\n",
    "print(estr.theta[0:2])\n",
    "print(\"95% CI\")\n",
    "print(ci[:2, :].T)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8d8ec22",
   "metadata": {},
   "source": [
    "This concludes chapter 14.\n",
    "\n",
    "Replication of the following chapters 15-17 are not yet available."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}