{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2ceaed98-3128-48e5-ae81-1f67899f9aa9",
   "metadata": {},
   "source": [
    "# Coffman et al. (2021): Causal Mediation Analysis\n",
    "\n",
    "For the `twangMediation` package in R, Coffman et al. provided a vignette on causal mediation analysis in the context of health disparity research. Specifically, the vignette examined possible mediating pathways for substance use disparities among sexual minority women. Data for this analysis comes from the National Survey of Drug Use and Health (NSDUH) and is available from the R package.\n",
    "\n",
    "Drawing from this vignette, we show how `delicatessen` can be used to conduct causal mediation analyses. Note that the provided data consists of more than 40,000 observations, so fitting these models might take a second.\n",
    "\n",
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "71cbce16-1843-4929-88a9-b0dfc2e8171d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Versions\n",
      "NumPy:         2.3.5\n",
      "SciPy:         1.16.3\n",
      "pandas:        2.3.3\n",
      "Delicatessen:  4.1\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import scipy as sp\n",
    "import pandas as pd\n",
    "from formulaic import model_matrix\n",
    "\n",
    "import delicatessen as deli\n",
    "from delicatessen import MEstimator\n",
    "from delicatessen.estimating_equations import ee_regression, ee_ipw\n",
    "from delicatessen.utilities import inverse_logit, spline, 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": "efaf3b4e-1e63-4118-946f-2a477c25a808",
   "metadata": {},
   "outputs": [],
   "source": [
    "d = pd.read_csv(\"data/nsduh.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "313f3423-f6e8-464b-9994-df23d86c8c6e",
   "metadata": {},
   "source": [
    "In the context of the vignette, interest was in the disparities of adult smoking status (`cigmon`) by sexual minority (e.g., lesbian, gay, bisexual; `lgb_flag`) as mediated by early smoking initiation (prior to age 15; `cig15`). The adjustment set is age (`age`), race (`race`), education (`educ`), income (`income`), and employment (`employ`). \n",
    "\n",
    "## Average Causal Effect\n",
    "\n",
    "To start, we will estimate the average causal effect of sexual minority on adult smoking status as done in the vignette. We will do this using an inverse probability of treatment weighting estimator. The weights are defined as\n",
    "$$ IPTW = \\frac{A_i}{\\Pr(A_i = 1 \\mid W_i)} + \\frac{1 - A_i}{1 - \\Pr(A_i = 1 \\mid W_i)} $$\n",
    "and then can be used to compute the weighted mean. Here, we use the Hajek estimator due to its improved finite-sample performance relative to the Horvitz-Thompson estimator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "819ae36b-a9f4-4f6e-ad3d-029b19f39968",
   "metadata": {},
   "outputs": [],
   "source": [
    "W_matrix = np.asarray(model_matrix(\"C(age) + C(race) + C(educ) + C(income) + C(employ)\", d))\n",
    "a = np.asarray(d['lgb_flag'])\n",
    "y = np.asarray(d['cigmon'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "14582e16-2683-4e2b-8fe7-a75a7df425d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def psi_ace(theta):\n",
    "    ace, mu1, mu0 = theta[:3]\n",
    "    alpha = theta[3:]\n",
    "\n",
    "    # Propensity score model\n",
    "    ee_ps = ee_regression(theta=alpha, X=W_matrix, y=a, model='logistic')\n",
    "    ps = inverse_logit(np.dot(W_matrix, alpha))\n",
    "    iptw = a/ps + (1-a)/(1-ps)\n",
    "\n",
    "    # Hajek estimator for the causal risks\n",
    "    ee_r1 = iptw * a * (y - mu1)\n",
    "    ee_r0 = iptw * (1-a) * (y - mu0)\n",
    "    ee_ace = np.ones(y.shape[0]) * (mu1 - mu0 - ace)\n",
    "\n",
    "    return np.vstack([ee_ace, ee_r1, ee_r0, ee_ps])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4b2137ac-cde0-4893-9e19-d8756a923590",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ACE:    0.12460267719839338\n",
      "95% CI: [0.10793198 0.14127338]\n"
     ]
    }
   ],
   "source": [
    "init_vals = [0., 0.5, 0.5, ] + [0., ]*W_matrix.shape[1]\n",
    "estr = MEstimator(psi_ace, init=init_vals)\n",
    "estr.estimate()\n",
    "ci = estr.confidence_intervals()\n",
    "\n",
    "print(\"ACE:   \", estr.theta[0])\n",
    "print(\"95% CI:\", ci[0, :])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d0a4781-b2c3-4dfc-b1f9-217618bafe30",
   "metadata": {},
   "source": [
    "These results suggest a disparity by sexual minority status in adult smoking behaviors that is not explained by the adjustment set. These results are similar to `twangMediation` but do differ as a result of how the propensity scores are being estimated. In the vignette, XGBoost (an ensemble of decision trees) is used. Here, a simple parameteric model is used instead. It would be better practice to specify a more flexible model for the propensity score (e.g., interactions). \n",
    "\n",
    "While it may seem like `twangMediation` has an advantage here, the corresponding statisitcal inference is invalid as XGBoost is not generally valid without the use of cross-fitting and doubly robust estimators. See Zivich & Breskin (2021) for details on why the `twangMediation` approach is not statistically valid. The approach based on parametric models used here with `delicatessen` is statistically valid.\n",
    "\n",
    "## Controlled Direct Effect\n",
    "\n",
    "The controlled direct effect is one parameter we can consider in mediation analysis. Letting $Y(a,m)$ denote the potential outcome under action $a$ and mediator $m$, the controlled direct effect is \n",
    "$$ E[Y(a,m)] - E[Y(a',m)] $$\n",
    "so we are considering the effect of changing $a$ when $m$ is held at a fixed value. \n",
    "\n",
    "For the controlled direct effect, we need to estimate a pair of weights (Coffman & Zhong 2012). First, we estimate the inverse probability of treatment weights as before. The second weight needed is the inverse probability of mediator weights, defined as the following\n",
    "$$ IPMW = \\frac{M_i}{\\Pr(M_i = 1 \\mid A_i, W_i)} + \\frac{1 - M_i}{1 - \\Pr(M_i = 1 \\mid A_i, W_i)} $$\n",
    "which we can jointly estimate all parameters using a stacked extension of the previous estimating functions with an additional logistic regression model. Here, the same adjustment set is used for both nuisance models. However, it can be the case that there are additional confounding variables for mediators. These can be dealt with by extending the conditional probability for $IPMW$.\n",
    "\n",
    "For estimation, we again use the Hajek estimator with the *product* of the two weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "12d473d7-7e55-4a37-aa46-8e20c5f831bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "W_matrix = np.asarray(model_matrix(\"C(age) + C(race) + C(educ) + C(income) + C(employ)\", d))\n",
    "X_matrix = np.asarray(model_matrix(\"lgb_flag + C(age) + C(race) + C(educ) + C(income) + C(employ)\", d))\n",
    "a = np.asarray(d['lgb_flag'])\n",
    "m = np.asarray(d['cig15'])\n",
    "y = np.asarray(d['cigmon'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c901b8a5-e37d-446b-9f0f-ddb242ebce30",
   "metadata": {},
   "outputs": [],
   "source": [
    "def psi_cde(theta):\n",
    "    ndim_W = W_matrix.shape[1]\n",
    "    ace1, ace0, mu11, mu01, mu10, mu00 = theta[:6]\n",
    "    alpha = theta[6: 6+ndim_W]\n",
    "    gamma = theta[6+ndim_W: ]\n",
    "\n",
    "    # Propensity score model for A\n",
    "    ee_ps = ee_regression(theta=alpha, X=W_matrix, y=a, model='logistic')\n",
    "    ps = inverse_logit(np.dot(W_matrix, alpha))\n",
    "    iptw = a/ps + (1-a)/(1-ps)\n",
    "\n",
    "    # Mediation score model for M\n",
    "    ee_ms = ee_regression(theta=gamma, X=X_matrix, y=m, model='logistic')\n",
    "    ms = inverse_logit(np.dot(X_matrix, gamma))\n",
    "    ipmw = m/ms + (1-m)/(1-ms)\n",
    "\n",
    "    # Hajek estimator for the causal risks\n",
    "    ipw = iptw * ipmw\n",
    "    ee_r11 = ipw * a * m * (y - mu11)\n",
    "    ee_r01 = ipw * (1-a) * m * (y - mu01)\n",
    "    ee_r10= ipw * a * (1-m) * (y - mu10)\n",
    "    ee_r00 = ipw * (1-a) * (1-m) * (y - mu00)\n",
    "    ee_ace1 = np.ones(y.shape[0]) * (mu11 - mu01 - ace1)\n",
    "    ee_ace0 = np.ones(y.shape[0]) * (mu10 - mu00 - ace0)\n",
    "\n",
    "    return np.vstack([ee_ace1, ee_ace0, ee_r11, ee_r01, ee_r10, ee_r00, ee_ps, ee_ms])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e162ba96-74b8-4d0f-a3cf-1c68af02d5af",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CDE(1): 0.10231467457422143\n",
      "95% CI: [0.06129448 0.14333486]\n",
      "CDE(0): 0.09821972928863482\n",
      "95% CI: [0.07878279 0.11765667]\n"
     ]
    }
   ],
   "source": [
    "init_vals = ([0., 0., 0.5, 0.5, 0.5, 0.5] \n",
    "             + [0., ]*W_matrix.shape[1]\n",
    "             + [0., ]*X_matrix.shape[1])\n",
    "estr = MEstimator(psi_cde, init=init_vals)\n",
    "estr.estimate()\n",
    "ci = estr.confidence_intervals()\n",
    "\n",
    "print(\"CDE(1):\", estr.theta[0])\n",
    "print(\"95% CI:\", ci[0, :])\n",
    "print(\"CDE(0):\", estr.theta[1])\n",
    "print(\"95% CI:\", ci[1, :])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bfca1707-385a-41d7-9947-cdd3e914b709",
   "metadata": {},
   "source": [
    "Note that the `twangMediation` documentation does not list the controlled direct effects, so no comparison can be made here. However, these numbers do happen to be similar to the *pure* direct effect estimates, which we will examine in the next section.\n",
    "\n",
    "As we have estimated a pair of parameters (two controlled direct effects), confidence bands are more appropriate to report here (since they provide *simultaneous* coverage for both paramters). sup-t confidence bands can be easily computed using `delicatessen` with the following functionality. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2f986e59-f246-47c8-85f5-b0c8b862918f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "95% CB - CDE(1): [0.05552766 0.14910169]\n",
      "95% CB - CDE(0): [0.07605025 0.12038921]\n"
     ]
    }
   ],
   "source": [
    "cb = estr.confidence_bands(subset=[0, 1], seed=4281014)\n",
    "print(\"95% CB - CDE(1):\", cb[0, :])\n",
    "print(\"95% CB - CDE(0):\", cb[1, :])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e07e8eb-8442-4b8b-bfa1-1199d570b000",
   "metadata": {},
   "source": [
    "As shown here, these are slightly wider than the confidence intervals. This is to be expected, as we are interested in multiple parameters but still would like to have simultaneous coverage. However, the sup-t correction is not overly conservative like the Bonferroni correction.\n",
    "\n",
    "## Pure Direct Effect\n",
    "\n",
    "The next parameter we could be interested in is the pure direct effect (also known as the natural direct effect). Using the potential outcomes from before, the pure direct effect is \n",
    "$$ E[Y(a, M(a'))] -  E[Y(a', M(a'))]$$\n",
    "where $M(a)$ is the potential mediator under $a$. Note that there is the seemingly odd combination of $Y$ under $a$ but the mediator is the one under $a'$. Such a pattern does not occur in nature. Identification of the pure direct effect thus relies on a 'cross-world' assumption for identification. This assumption is what makes these effect 'unnatural' (to me).\n",
    "\n",
    "For the right quantity, note that $E[Y(a')]$. Thus, that part can be estimated using the IPW estimator for the never-act part of the average causal effect estimator. For the left quantity, we use 'strategy 3' of Tchetgen Tchetgen & Shpitser (2012). The described IPW estimator is\n",
    "$$ n^{-1} \\sum_{i=1}^{n} Y_i \\frac{I(A_i = 1)}{\\Pr(A=1 \\mid W_i)} \\times \\frac{\\Pr(M_i = m \\mid A = 0, W_i)}{\\Pr(M_i = m \\mid A_i, W_i)} $$\n",
    "where the last probability fractions are the probability of the *observed* value of $M$ for that individual."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b14f99e6-a9ee-40c0-a835-569a8f97319c",
   "metadata": {},
   "outputs": [],
   "source": [
    "W_matrix = np.asarray(model_matrix(\"C(age) + C(race) + C(educ) + C(income) + C(employ)\", d))\n",
    "X_matrix = np.asarray(model_matrix(\"lgb_flag + C(age) + C(race) + C(educ) + C(income) + C(employ)\", d))\n",
    "d0 = d.copy()\n",
    "d0['lgb_flag'] = 0\n",
    "X0_matrix = np.asarray(model_matrix(\"lgb_flag + C(age) + C(race) + C(educ) + C(income) + C(employ)\", d0))\n",
    "a = np.asarray(d['lgb_flag'])\n",
    "m = np.asarray(d['cig15'])\n",
    "y = np.asarray(d['cigmon'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5a2f69cc-cbc6-4a5d-82cd-8bb215df67d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def psi_pde(theta):\n",
    "    ndim_W = W_matrix.shape[1]\n",
    "    pde, mu1m, mu0 = theta[:3]\n",
    "    alpha = theta[3: 3+ndim_W]\n",
    "    gamma = theta[3+ndim_W: ]\n",
    "\n",
    "    # Propensity score model for A\n",
    "    ee_ps = ee_regression(theta=alpha, X=W_matrix, y=a, model='logistic')\n",
    "    ps = inverse_logit(np.dot(W_matrix, alpha))\n",
    "    iptw = a/ps + (1-a)/(1-ps)\n",
    "\n",
    "    # Mediation score model for M\n",
    "    ee_ms = ee_regression(theta=gamma, X=X_matrix, y=m, model='logistic')\n",
    "    ms = inverse_logit(np.dot(X_matrix, gamma))\n",
    "    m0s = inverse_logit(np.dot(X0_matrix, gamma))\n",
    "    ipmw = m*(m0s / ms) + (1-m)*(1-m0s)/(1-ms)\n",
    "\n",
    "    # Hajek estimator for the causal risks\n",
    "    ee_r1m = iptw * a * ipmw * (y - mu1m)\n",
    "    ee_r0 = iptw * (1-a) * (y - mu0)\n",
    "    ee_pde = np.ones(y.shape[0]) * (mu1m - mu0 - pde)\n",
    "\n",
    "    return np.vstack([ee_pde, ee_r1m, ee_r0, ee_ps, ee_ms])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "cdff4eee-c5ed-4a25-b7ae-6486004d92e1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PDE:    0.09552324822589883\n",
      "95% CI: [0.07882165 0.11222485]\n"
     ]
    }
   ],
   "source": [
    "init_vals = ([0., 0.5, 0.5] \n",
    "             + [0., ]*W_matrix.shape[1]\n",
    "             + [0., ]*X_matrix.shape[1])\n",
    "estr = MEstimator(psi_pde, init=init_vals)\n",
    "estr.estimate()\n",
    "ci = estr.confidence_intervals()\n",
    "\n",
    "print(\"PDE:   \", estr.theta[0])\n",
    "print(\"95% CI:\", ci[0, :])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3e8f783-fcef-4a99-9ae2-1c51e6f5ca81",
   "metadata": {},
   "source": [
    "These results are fairly similar to those provided in `twangMediation` (but the results here are not stratified and we are using a different nuisance model).\n",
    "\n",
    "## Pure Indirect Effect\n",
    "The pure (natural) indirect effect can be defined as\n",
    "$$ E[Y(a, M(a))] -  E[Y(a, M(a'))]$$\n",
    "which again involves a cross-world assumption. Note the second part here is the same as the left side of the pure direct effect. This shows how to decompose everything. The following code implements the estimator for the pure indirect effect"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6c3bd015-c3f2-43a2-8b10-0f9ca7329514",
   "metadata": {},
   "outputs": [],
   "source": [
    "def psi_pie(theta):\n",
    "    ndim_W = W_matrix.shape[1]\n",
    "    pie, mu1, mu1m = theta[:3]\n",
    "    alpha = theta[3: 3+ndim_W]\n",
    "    gamma = theta[3+ndim_W: ]\n",
    "\n",
    "    # Propensity score model for A\n",
    "    ee_ps = ee_regression(theta=alpha, X=W_matrix, y=a, model='logistic')\n",
    "    ps = inverse_logit(np.dot(W_matrix, alpha))\n",
    "    iptw = a/ps + (1-a)/(1-ps)\n",
    "\n",
    "    # Mediation score model for M\n",
    "    ee_ms = ee_regression(theta=gamma, X=X_matrix, y=m, model='logistic')\n",
    "    ms = inverse_logit(np.dot(X_matrix, gamma))\n",
    "    m0s = inverse_logit(np.dot(X0_matrix, gamma))\n",
    "    ipmw = m*(m0s/ms) + (1-m)*(1-m0s)/(1-ms)\n",
    "\n",
    "    # Hajek estimator for the causal risks\n",
    "    ee_r1 = iptw * a * (y - mu1)\n",
    "    ee_r1m = iptw * a * ipmw * (y - mu1m)\n",
    "    ee_pie = np.ones(y.shape[0]) * (mu1 - mu1m - pie)\n",
    "\n",
    "    return np.vstack([ee_pie, ee_r1, ee_r1m, ee_ps, ee_ms])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "41997b23-2084-4022-9919-1e7530ed92d7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PIE:    0.029079428972494544\n",
      "95% CI: [0.02309128 0.03506758]\n"
     ]
    }
   ],
   "source": [
    "init_vals = ([0., 0.5, 0.5] \n",
    "             + [0., ]*W_matrix.shape[1]\n",
    "             + [0., ]*X_matrix.shape[1])\n",
    "estr = MEstimator(psi_pie, init=init_vals)\n",
    "estr.estimate()\n",
    "ci = estr.confidence_intervals()\n",
    "\n",
    "print(\"PIE:   \", estr.theta[0])\n",
    "print(\"95% CI:\", ci[0, :])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4de40b40-a166-419c-8a2a-5deba8928216",
   "metadata": {},
   "source": [
    "That provides the pure indirect effect estimate. This is again similar to `twangMediation` and if we add the pure effects together, we get back to the average causal effect.\n",
    "\n",
    "## Pure Effects Altogether\n",
    "\n",
    "As a final step, we now package the pure direct and indirect effect estimators together into a single function. As with the controlled direct effect, interest is in a *pair* of parameters, so confidence bands are more appropriate to report here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "608af742-d255-4ea9-86c4-4e430c662515",
   "metadata": {},
   "outputs": [],
   "source": [
    "def psi_full(theta):\n",
    "    ndim_W = W_matrix.shape[1]\n",
    "    pde, pie, mu1, mu0, mu1m = theta[:5]\n",
    "    alpha = theta[5: 5+ndim_W]\n",
    "    gamma = theta[5+ndim_W: ]\n",
    "\n",
    "    # Propensity score model for A\n",
    "    ee_ps = ee_regression(theta=alpha, X=W_matrix, y=a, model='logistic')\n",
    "    ps = inverse_logit(np.dot(W_matrix, alpha))\n",
    "    iptw = a/ps + (1-a)/(1-ps)\n",
    "\n",
    "    # Mediation score model for M\n",
    "    ee_ms = ee_regression(theta=gamma, X=X_matrix, y=m, model='logistic')\n",
    "    ms = inverse_logit(np.dot(X_matrix, gamma))\n",
    "    m0s = inverse_logit(np.dot(X0_matrix, gamma))\n",
    "    ipmw = m*(m0s/ms) + (1-m)*(1-m0s)/(1-ms)\n",
    "\n",
    "    # Hajek estimator for the causal risks\n",
    "    ee_r1 = iptw * a * (y - mu1)\n",
    "    ee_r0 = iptw * (1-a) * (y - mu0)\n",
    "    ee_r1m = iptw * a * ipmw * (y - mu1m)\n",
    "    ee_pde = np.ones(y.shape[0]) * (mu1m - mu0 - pde)\n",
    "    ee_pie = np.ones(y.shape[0]) * (mu1 - mu1m - pie)\n",
    "\n",
    "    return np.vstack([ee_pde, ee_pie, ee_r1, ee_r0, ee_r1m, ee_ps, ee_ms])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "804fefc9-6bb0-4ff4-a963-7255ace03157",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PDE:    0.09552324822589882\n",
      "95% CB: [0.02309128 0.03506758]\n",
      "PIE:    0.029079428972494527\n",
      "95% CB: [0.30311362 0.33564762]\n"
     ]
    }
   ],
   "source": [
    "init_vals = ([0., 0., 0.5, 0.5, 0.5] \n",
    "             + [0., ]*W_matrix.shape[1]\n",
    "             + [0., ]*X_matrix.shape[1])\n",
    "estr = MEstimator(psi_full, init=init_vals)\n",
    "estr.estimate()\n",
    "cb = estr.confidence_bands(subset=[0, 1], seed=921453)\n",
    "\n",
    "print(\"PDE:   \", estr.theta[0])\n",
    "print(\"95% CB:\", ci[0, :])\n",
    "print(\"PIE:   \", estr.theta[1])\n",
    "print(\"95% CB:\", ci[1, :])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0a06865-ccf2-4a66-9b5f-3beeb9b2849d",
   "metadata": {},
   "source": [
    "This provides our results collected from before into a single estimating function. As a result, we are able to construct sup-t confidence bands for appropriate inference on the decomposition of the total effect.\n",
    "\n",
    "We can also check that our decomposition sums back to the average causal effect."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "987942a9-0c17-4ab8-9e86-0dfdaa0c9817",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "np.float64(0.12460267719839335)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "estr.theta[0] + estr.theta[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e41d8db3-4d97-4c77-8769-9b99160110ec",
   "metadata": {},
   "source": [
    "which matches the average causal effect for many of the decimals (we expect some minor differences due to floating point precision and the root-finding procedure tolerance not being infinite). \n",
    "\n",
    "This example provides another illustration of how to use `delicatessen` to construct custom estimating equations and use it for statistically valid inference. \n",
    "\n",
    "## References\n",
    "\n",
    "Coffman, D. L., & Zhong, W. (2012). Assessing mediation using marginal structural models in the presence of confounding and moderation. Psychological methods, 17(4), 642.\n",
    "\n",
    "Coffman, D. L., Schuler, M. S., McCaffrey, D. F., Castellano, K. E., Zhou, H., Vegetabile, B., & Griffin, B. A. (2021). A tutorial for conducting causal mediation analysis with the twangMediation package.\n",
    "\n",
    "Tchetgen Tchetgen, E. J., & Shpitser, I. (2012). Semiparametric theory for causal mediation analysis: efficiency bounds, multiple robustness, and sensitivity analysis. Annals of statistics, 40(3), 1816.\n",
    "\n",
    "Zivich, P. N., & Breskin, A. (2021). Machine learning for causal inference: on the use of cross-fit estimators. Epidemiology, 32(3), 393-401."
   ]
  }
 ],
 "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.13.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}