R/get_draws.R

In marginaleffects: Predictions, Comparisons, Slopes, Marginal Means, and Hypothesis Tests

#' Extract Posterior Draws or Bootstrap Resamples from `marginaleffects` Objects
#'
#' @param x An object produced by a `marginaleffects` package function, such as `predictions()`, `avg_slopes()`, `hypotheses()`, etc.
#' @param shape string indicating the shape of the output format:
#' * "long": long format data frame
#' * "DxP": Matrix with draws as rows and parameters as columns
#' * "PxD": Matrix with draws as rows and parameters as columns
#' * "rvar": Random variable datatype (see `posterior` package documentation).
#' @return A data.frame with `drawid` and `draw` columns.
#' @details
#' If DxP and PxD and the names returned by `coef(x)` are unique, `marginaleffects` sets parameter names to those names. Otherwise, it sets them to `b1`, `b2`, etc.
#' @export
get_draws <- function(x, shape = "long") {
  checkmate::assert_choice(shape, choices = c("long", "DxP", "PxD", "rvar"))

  # tidy.comparisons() sometimes already saves draws in a nice long format
  draws <- attr(x, "posterior_draws")
  if (inherits(draws, "posterior_draws")) {
    return(draws)
  }

  if (is.null(attr(x, "posterior_draws"))) {
    warning('This object does not include a "posterior_draws" attribute. The `posterior_draws` function only supports bayesian models produced by the `marginaleffects` or `predictions` functions of the `marginaleffects` package.',
      call. = FALSE)
    return(x)
  }

  if (nrow(draws) != nrow(x)) {
    stop("The number of parameters in the object does not match the number of parameters for which posterior draws are available.", call. = FALSE)
  }

  if (shape %in% c("PxD", "DxP")) {
    parnames <- names(stats::coef(x))
    if (length(parnames) != nrow(draws) || anyDuplicated(parnames) > 0) {
      parnames <- paste0("b", seq_len(nrow(draws)))
    }
    row.names(draws) <- parnames
    colnames(draws) <- paste0("draw", seq_len(ncol(draws)))
  }

  if (shape == "PxD") {
    return(draws)
  }

  if (shape == "DxP") {
    return(t(draws))
  }

  if (shape == "rvar") {
    insight::check_if_installed("posterior")
    draws <- t(draws)
    if (!is.null(attr(x, "nchains"))) {
      x[["rvar"]] <- posterior::rvar(draws, nchains = attr(x, "nchains"))
    } else {
      x[["rvar"]] <- posterior::rvar(draws)
    }
    return(x)
  }

  if (shape == "long") {
    draws <- data.table(draws)
    setnames(draws, as.character(seq_len(ncol(draws))))
    for (v in colnames(x)) {
      draws[[v]] <- x[[v]]
    }
    out <- melt(
      draws,
      id.vars = colnames(x),
      variable.name = "drawid",
      value.name = "draw")
    cols <- unique(c("drawid", "draw", "rowid", colnames(out)))
    cols <- intersect(cols, colnames(out))
    setcolorder(out, cols)
    data.table::setDF(out)
    return(out)
  }
}


average_draws <- function(data, index, draws, byfun = NULL) {
  insight::check_if_installed("collapse", minimum_version = "2.0.18")

  w <- data[["marginaleffects_wts_internal"]]
  if (all(is.na(w))) {
    w <- NULL
  }

  if (is.null(index)) {
    index <- intersect(colnames(data), "type")
  }

  if (length(index) > 0) {
    if ("term" %in% colnames(data)) {
      Encoding(data[["term"]]) <- "UTF-8"
    }
    g <- collapse::GRP(data, by = index)

    if (is.null(byfun)) {
      draws <- collapse::fmean(
        draws,
        g = g,
        w = w,
        drop = FALSE)
    } else {
      draws <- collapse::BY(
        draws,
        g = g,
        FUN = byfun,
        drop = FALSE)
    }
    out <- data.table(
      g[["groups"]],
      average = collapse::dapply(draws, MARGIN = 1, FUN = collapse::fmedian))
  } else {
    if (is.null(byfun)) {
      draws <- collapse::fmean(
        draws,
        w = w,
        drop = FALSE)
    } else {
      draws <- collapse::BY(
        draws,
        g = g,
        FUN = byfun,
        drop = FALSE)
    }
    out <- data.table(average = collapse::dapply(draws, MARGIN = 1, FUN = collapse::fmedian))
  }

  setnames(out, old = "average", new = "estimate")
  attr(out, "posterior_draws") <- draws
  return(out)
}




#' alias to `get_draws()` keep forever for backward compatibility with JSS
#'
#' @inherit posterior_draws
#' @keywords internal
#' @export
posterior_draws <- get_draws