R/treatment_effect.R

In RobinCar2: ROBust INference for Covariate Adjustment in Randomized Clinical Trials

#' Treatment Effect
#' @description Obtain treatment effect and variance from counter-factual prediction
#'
#' @param object Object from which to obtain treatment effect.
#' @param pair (`contrast`) Contrast choices.
#' @param eff_measure (`function`) Treatment effect measurement function.
#' @param eff_jacobian (`function`) Treatment effect jacobian function.
#' @param contrast_name (`string`) Name of the contrast.
#' @param ... Additional arguments for variance.
#' @return A list of `treatment_effect` object with following elements:
#' - `estimate`: estimate of the treatment effect.
#' - `pair`: `contrast` object indicating the pairwise treatment effect.
#' - `contrast`: name of the contrast function.
#' - `euqal_val`: the value for no treatment effect given the contrast.
#' - `variance`: the variance of the treatment effect.
#' - `jacobian`: the Jacobian matrix.
#' - `contrast_mat`: contrast summary matrix.
#'
#' @export
treatment_effect <- function(
    object,
    pair = pairwise(names(object$estimate)),
    eff_measure,
    eff_jacobian = eff_jacob(eff_measure),
    contrast_name,
    ...) {
  UseMethod("treatment_effect")
}

#' @export
treatment_effect.prediction_cf <- function(
    object,
    pair = pairwise(names(object$estimate)),
    eff_measure,
    eff_jacobian = eff_jacob(eff_measure),
    contrast_name = deparse(substitute(eff_measure)),
    ...) {
  assert_function(eff_measure)
  assert_class(pair, "contrast")
  assert_string(contrast_name)
  # make sure levels match
  pair <- update_levels(pair, names(object$estimate))
  trt_effect <- unname(eff_measure(object$estimate[pair[[1]]], object$estimate[pair[[2]]]))
  trt_jac <- eff_jacobian(object$estimate[pair[[1]]], object$estimate[pair[[2]]])
  trt_jac_mat <- jac_mat(trt_jac, pair)
  equal_val <- eff_measure(object$estimate[1], object$estimate[1])

  contrast_var <- trt_jac_mat %*% object$variance %*% t(trt_jac_mat)

  if (is.null(contrast_var)) {
    trt_sd <- rep(NA, length(object))
  } else {
    trt_sd <- sqrt(diag(contrast_var))
  }

  z_value <- as.numeric(trt_effect - equal_val) / trt_sd
  p <- 2 * pnorm(abs(z_value), lower.tail = FALSE)
  contrast_mat <- matrix(
    c(
      trt_effect,
      trt_sd,
      z_value,
      p
    ),
    nrow = length(trt_effect)
  )
  colnames(contrast_mat) <- c("Estimate", "Std.Err", "Z Value", "Pr(>|z|)")
  row.names(contrast_mat) <- sprintf("%s v.s. %s", attr(pair, "levels")[pair[[1]]], attr(pair, "levels")[pair[[2]]])
  structure(
    list(
      estimate = trt_effect,
      pair = pair,
      contrast = contrast_name,
      equal_val = equal_val,
      variance = contrast_var,
      jacobian = trt_jac_mat,
      contrast_mat = contrast_mat
    ),
    class = "treatment_effect"
  )
}


#' @export
#' @inheritParams predict_counterfactual
treatment_effect.lm <- function(
    object,
    pair,
    eff_measure,
    eff_jacobian = eff_jacob(eff_measure),
    contrast_name = deparse(substitute(eff_measure)),
    vcov = "vcovG",
    vcov_args = list(),
    treatment,
    data = find_data(object),
    ...) {
  pc <- predict_counterfactual(object, data = data, treatment, vcov = vcov, vcov_args = vcov_args)
  if (missing(pair)) {
    pair <- pairwise(names(pc$estimate))
  }
  treatment_effect(pc, pair = pair, eff_measure = eff_measure, eff_jacobian = eff_jacobian, ...)
}

#' @export
treatment_effect.glm <- function(
    object,
    pair,
    eff_measure,
    eff_jacobian = eff_jacob(eff_measure),
    contrast_name = deparse(substitute(eff_measure)),
    vcov = "vcovG",
    vcov_args = list(),
    treatment,
    data = find_data(object),
    ...) {
  pc <- predict_counterfactual(object, treatment, data, vcov = vcov, vcov_args = vcov_args)
  if (missing(pair)) {
    pair <- pairwise(names(pc$estimate))
  }
  treatment_effect(pc, pair = pair, eff_measure = eff_measure, eff_jacobian = eff_jacobian, ...)
}

#' @rdname treatment_effect
difference <- function(object, ...) {
  treatment_effect(object, eff_measure = h_diff, eff_jacobian = h_jac_diff, ...)
}
#' @rdname treatment_effect
risk_ratio <- function(object, ...) {
  treatment_effect(object, eff_measure = h_risk_ratio, eff_jacobian = h_jac_risk_ratio, ...)
}
#' @rdname treatment_effect
odds_ratio <- function(object, ...) {
  treatment_effect(object, eff_measure = h_odds_ratio, eff_jacobian = h_jac_odds_ratio, ...)
}

#' @rdname treatment_effect
log_risk_ratio <- function(object, ...) {
  treatment_effect(object, eff_measure = h_log_risk_ratio, eff_jacobian = h_jac_log_risk_ratio, ...)
}
#' @rdname treatment_effect
log_odds_ratio <- function(object, ...) {
  treatment_effect(object, eff_measure = h_log_odds_ratio, eff_jacobian = h_jac_log_odds_ratio, ...)
}

#' Contrast Functions and Jacobians
#' @rdname contrast
#' @param x (`numeric`) Vector of values.
#' @return Vector of contrasts, or matrix of jacobians.
#' @examples
#' h_diff(1:3, 4:6)
#' h_jac_risk_ratio(1:3, 4:6)
#' @export
h_diff <- function(x, y) {
  assert_numeric(x)
  assert_numeric(y, len = length(x))
  x - y
}

#' @rdname contrast
#' @export
h_jac_diff <- function(x, y) {
  assert_numeric(x)
  assert_numeric(y, len = length(x))
  n <- length(x)
  matrix(c(1, -1), nrow = n, ncol = 2, byrow = TRUE)
}

#' @rdname contrast
#' @export
h_risk_ratio <- function(x, y) {
  assert_numeric(x, lower = 0)
  assert_numeric(y, lower = 0, len = length(x))
  x / y
}

#' @rdname contrast
#' @export
h_jac_risk_ratio <- function(x, y) {
  assert_numeric(x, lower = 0)
  assert_numeric(y, lower = 0, len = length(x))
  cbind(1 / y, -x / y^2)
}

#' @rdname contrast
#' @export
h_odds_ratio <- function(x, y) {
  assert_numeric(x, lower = 0, upper = 1)
  assert_numeric(y, lower = 0, upper = 1, len = length(x))
  h_risk_ratio(x / (1 - x), y / (1 - y))
}

#' @rdname contrast
#' @export
h_jac_odds_ratio <- function(x, y) {
  assert_numeric(x, lower = 0, upper = 1)
  assert_numeric(y, lower = 0, upper = 1, len = length(x))
  cbind(1 / (1 - x)^2 / y * (1 - y), -x / (1 - x) / y^2)
}

#' @rdname contrast
#' @export
h_log_risk_ratio <- function(x, y) {
  log(h_risk_ratio(x, y))
}

#' @rdname contrast
#' @export
h_jac_log_risk_ratio <- function(x, y) {
  cbind(1 / x, -1 / y)
}

#' @rdname contrast
#' @export
h_log_odds_ratio <- function(x, y) {
  log(h_odds_ratio(x, y))
}

#' @rdname contrast
#' @export
h_jac_log_odds_ratio <- function(x, y) {
  assert_numeric(x, lower = 0, upper = 1)
  assert_numeric(y, lower = 0, upper = 1, len = length(x))
  cbind(1 / x + 1 / (1 - x), -1 / y - 1 / (1 - y))
}

#' @rdname contrast
#' @param f (`function`) Function with argument x and y to compute treatment effect.
#' @export
eff_jacob <- function(f) {
  assert_function(f, args = c("x", "y"))
  function(x, y) {
    cbind(
      numDeriv::grad(function(x) f(x = x, y = y), x),
      numDeriv::grad(function(y) f(x = x, y = y), y)
    )
  }
}

#' @export
print.treatment_effect <- function(x, ...) {
  stats::printCoefmat(
    x$contrast_mat
  )
}

#' Confidence interval function.
#' @rdname confint
#' @export
confint.treatment_effect <- function(object, parm, level = 0.95, transform, ...) {
  if (missing(transform)) {
    transform <- if (
      object$contrast %in%
        c(
          "log_odds_ratio",
          "log_risk_ratio",
          "h_log_odds_ratio",
          "h_log_risk_ratio"
        )
    ) {
      exp
    } else {
      identity
    }
  }
  h_confint(object$contrast_mat, parm = parm, level = level, transform = transform, ...)
}