`marginaleffects` unconditional vcov

get_vcov_unconditional.R

library(marginaleffects)

# sanity checks: let's be super strict
sanity_vcov_unconditional <- function(p) {
  insight::check_if_installed("sandwich")
  cal <- attr(p, "call")

  callfun <- deparse(cal[[1]])
  if (!identical(callfun, "predictions")) {
    msg <- "`vcov='unconditional'` is only supported for `avg_predictions()` calls."
    stop(msg, call. = FALSE)
  }

  if (!is.null(cal[["newdata"]])) {
    msg <- "`vcov='unconditional'` does not support `newdata`."
    stop(msg, call. = FALSE)
  }

  variables <- cal[["variables"]]
  checkmate::assert_string(variables, null.ok = FALSE)

  by <- cal[["by"]]
  flag <- !isTRUE(by) && !identical(by, variables) && !is.null(by)
  if (flag) {
    msg <- "`vcov='unconditional'`. `by` must be `TRUE` or identical to `variables`."
    stop(msg, call. = FALSE)
  }
}

get_vcov_unconditional <- function(p) {
  sanity_vcov_unconditional(p)

  cal <- cal_unit <- attr(p, "call")
  newdata <- attr(p, "newdata")
  model <- attr(p, "model")

  variables <- cal[["variables"]]

  # unit-level yhat and dyhat/dbeta
  cal_unit[["by"]] <- FALSE
  p_unit <- eval(cal_unit)
  yhat <- p_unit$estimate
  J <- attr(p_unit, "jacobian")

  # unconditional variance
  beta_dot <- tcrossprod(sandwich::bread(model), sandwich::estfun(model))
  S <- do.call("cbind", lapply(p[[variables]], function(v) {
    in_v <- newdata[[variables]] == v
    (yhat[in_v] - mean(yhat[in_v])) + drop((colMeans(J[in_v, ]) %*% beta_dot))
  }))
  V <- crossprod(S / ncol(beta_dot))

  return(V)
}