R/est_ipw.R

In nhejazi/medshift: Causal mediation analysis for stochastic interventions

#' Inverse probability weighted (IPW) estimator
#'
#' @param data A \code{data.table} containing the observed data, with columns
#'  in the order specified by the NPSEM (Y, Z, A, W), with column names set
#'  appropriately based on the original input data. Such a structure is merely
#'  a convenience utility to passing data around to the various core estimation
#'  routines and is automatically generated by \code{\link{medshift}}.
#' @param delta A \code{numeric} value indicating the degree of shift in the
#'  intervention to be used in defining the causal quantity of interest. In the
#'  case of binary interventions, this takes the form of an incremental
#'  propensity score shift, acting as a multiplier of the odds with which a
#'  given observational unit receives the intervention (EH Kennedy, 2018, JASA;
#'  <doi:10.1080/01621459.2017.1422737>).
#' @param g_learners A \code{\link[sl3]{Stack}} (or other learner class that
#'   inherits from \code{\link[sl3]{Lrnr_base}}), containing a single or set of
#'   instantiated learners from \pkg{sl3}, to be used in fitting the propensity
#'   score, i.e., g = P(A | W).
#' @param e_learners A \code{\link[sl3]{Stack}} (or other learner class that
#'   inherits from \code{\link[sl3]{Lrnr_base}}), containing a single or set of
#'   instantiated learners from \pkg{sl3}, to be used in fitting a propensity
#'   score that conditions on the mediators, i.e., e = P(A | Z, W).
#' @param w_names A \code{character} vector of the names of the columns that
#'  correspond to baseline covariates (W). The input for this argument is
#'  automatically generated by a call to the wrapper function \code{medshift}.
#' @param z_names A \code{character} vector of the names of the columns that
#'  correspond to mediators (Z). The input for this argument is automatically
#'  generated by \code{\link{medshift}}.
#' @param ... Other arguments currently ignored.
est_ipw <- function(data,
                    delta,
                    g_learners,
                    e_learners,
                    w_names,
                    z_names,
                    ...) {
  # fit regression for incremental propensity score intervention
  g_out <- fit_g_mech(
    data = data, delta = delta,
    learners = g_learners, w_names = w_names
  )

  # fit clever regression for treatment, conditional on mediators
  e_out <- fit_e_mech(
    data = data, learners = e_learners,
    z_names = z_names, w_names = w_names
  )

  # get indices of treated and control units in validation data
  idx_A1 <- which(data$A == 1)
  idx_A0 <- which(data$A == 0)

  # compute IPW  estimator components from estimates of nuisance parameters
  ipw_out <- compute_ipw(
    g_output = g_out, e_output = e_out,
    idx_treat = idx_A1, idx_cntrl = idx_A0
  )
  g_shifted <- ipw_out$g_shifted
  e_pred <- ipw_out$e_pred
  mean_aipw <- ipw_out$mean_aipw

  # compute estimator
  estim_ipw <- mean(((g_shifted / e_pred) / mean_aipw) * data$Y)

  # output
  estim_ipw_out <- list(theta = estim_ipw, type = "re-weighted (IPW)")
  return(estim_ipw_out)
}

###############################################################################

#' Get inverse probability weighted (IPW) estimate from nuisance parameters
#'
#' @param g_output Object containing results from fitting the propensity score
#'  regression, as produced by a call to \code{\link{fit_g_mech}}.
#' @param e_output Object containing results from fitting the propensity score
#'  regression while conditioning on mediators, as produced by a call to
#'  \code{\link{fit_e_mech}}.
#' @param idx_treat A \code{numeric} vector providing the indices corresponding
#'  to units that received treatment (A = 1), for a binary intervention.
#' @param idx_cntrl A \code{numeric} vector providing the indices corresponding
#'  to units that did not receive treatment (A = 0), for a binary intervention.
#'
#' @keywords internal
compute_ipw <- function(g_output, e_output, idx_treat, idx_cntrl) {
  # compute components for A = 0 based on symmetry with A = 1 case
  g_shifted_A1 <- g_output$g_est$g_pred_shifted_A1
  g_shifted_A0 <- g_output$g_est$g_pred_shifted_A0
  e_pred_A1 <- e_output$e_est$e_pred_A1
  e_pred_A0 <- e_output$e_est$e_pred_A0

  # subset computed components based on observed treatment status for g
  g_shifted_obs <- rep(NA, length(idx_treat) + length(idx_cntrl))
  g_shifted_A1_obs <- g_shifted_A1[idx_treat]
  g_shifted_A0_obs <- g_shifted_A0[idx_cntrl]
  g_shifted_obs[idx_treat] <- g_shifted_A1_obs
  g_shifted_obs[idx_cntrl] <- g_shifted_A0_obs

  # subset computed components based on observed treatment status for e
  e_pred_obs <- rep(NA, length(idx_treat) + length(idx_cntrl))
  e_pred_A1_obs <- e_pred_A1[idx_treat]
  e_pred_A0_obs <- e_pred_A0[idx_cntrl]
  e_pred_obs[idx_treat] <- e_pred_A1_obs
  e_pred_obs[idx_cntrl] <- e_pred_A0_obs

  # stabilize weights in A-IPW by dividing by sample average since E[g/e] = 1
  mean_aipw <- mean(g_shifted_obs / e_pred_obs)

  # output as simple list
  return(list(
    g_shifted = g_shifted_obs,
    e_pred = e_pred_obs,
    mean_aipw = mean_aipw
  ))
}