R/get_marginaltrends.R

In modelbased: Estimation of Model-Based Predictions, Contrasts and Means

#' @rdname get_emmeans
#'
#' @examplesIf insight::check_if_installed("marginaleffects", quietly = TRUE)
#' model <- lm(Sepal.Width ~ Species * Petal.Length, data = iris)
#'
#' get_marginaltrends(model, trend = "Petal.Length", by = "Species")
#' get_marginaltrends(model, trend = "Petal.Length", by = "Petal.Length")
#' get_marginaltrends(model, trend = "Petal.Length", by = c("Species", "Petal.Length"))
#' @export
get_marginaltrends <- function(
  model,
  trend = NULL,
  by = NULL,
  predict = NULL,
  ci = 0.95,
  estimate = NULL,
  transform = NULL,
  p_adjust = "none",
  keep_iterations = FALSE,
  verbose = TRUE,
  ...
) {
  # check if available
  insight::check_if_installed("marginaleffects", minimum_version = "0.29.0")
  dots <- list(...)

  # set defaults
  estimate <- .validate_estimate_arg(estimate)

  # model details
  model_info <- insight::model_info(model, response = 1, verbose = FALSE)

  # Guess arguments
  myargs <- .guess_marginaltrends_arguments(model, trend, by, verbose, ...)

  # First step: create a data grid --------------------------------------------
  # ---------------------------------------------------------------------------

  # data grid only when we have by predictors
  if (is.null(by) && is.null(myargs$range)) {
    datagrid <- datagrid_info <- NULL
  } else {
    # setup arguments to create the data grid
    dg_factors <- switch(estimate, specific = "reference", "all")
    dg_args <- list(
      model,
      by = c(by, myargs$range),
      factors = dg_factors,
      include_random = TRUE,
      verbose = FALSE
    )
    # add user-arguments from "...", but remove those arguments that are already set
    dots[c("by", "factors", "include_random", "verbose")] <- NULL
    dg_args <- insight::compact_list(c(dg_args, dots))

    # Get corresponding datagrid (and deal with particular ats)
    datagrid <- do.call(insight::get_datagrid, dg_args)
    datagrid_info <- attributes(datagrid)

    # get "by" argument from data grid
    myargs$by <- datagrid_info$at_specs$varname

    # remove `range` (or: `trend`) variable from `by`. If user specified a range
    # for `trend`, we need the data (the defined range) in the data grid, but we
    # do not want to stratify by the same variable in `by` then (this won't
    # work). instead, in `.guess_marginaltrends_arguments()`, we copy the
    # `trend` value to `range`, which we pass to the data grid, and `trend` only
    # contains the variable name after that. This variable needs to be removed,
    # to not to be passed to `by`
    if (!is.null(myargs$range)) {
      myargs$by <- setdiff(myargs$by, myargs$trend)
      if (!length(myargs$by)) {
        myargs$by <- NULL
      }
    }
  }

  # Second step: prepare arguments for marginaleffects ------------------------
  # ---------------------------------------------------------------------------

  # sanity check
  if (!is.null(datagrid)) {
    datagrid <- as.data.frame(datagrid)
  }

  # remove user-arguments from "..." that will be used when calling marginaleffects
  dots[c("by", "conf_level", "digits")] <- NULL

  # handle weights - argument is named "wts" in marginal effects
  if (!is.null(dots$weights)) {
    dots$wts <- dots$weights
    dots$weights <- NULL
  }

  # handle custom hypothesis - officially, "estimate_slopes()" has no
  # comparison argument and contrasts of slopes is supported in
  # "estimate_contrasts()". But sometimes, it is useful to pass a custom
  # hypothesis to "avg_slopes()", so we just copy the "comparison" argument
  # to the "hypothesis" argument here.
  if (!is.null(dots$comparison)) {
    dots$hypothesis <- dots$comparison
    dots$comparison <- NULL
  }

  # model df - can be passed via `...`
  if (is.null(dots$df) && !model_info$is_bayesian) {
    dots$df <- insight::get_df(model, type = "wald", verbose = FALSE)
  }

  # setup arguments again
  fun_args <- insight::compact_list(c(
    list(model, variables = myargs$trend, by = myargs$by, conf_level = ci),
    dots
  ))

  # all other "marginalizations"
  # we don't want a datagrid for "average" option
  if (is.null(dots$newdata) && estimate != "average") {
    # we allow individual "newdata" options, so do not
    # # overwrite if explicitly set
    fun_args$newdata <- datagrid
  }

  # fmt: skip
  # handle distributional parameters
  if (!is.null(predict) && inherits(model, "brmsfit") && predict %in% .brms_aux_elements(model)) {
    fun_args$dpar <- predict
  } else {
    fun_args$type <- predict
  }

  # Third step: compute marginal slopes ---------------------------------------
  # ---------------------------------------------------------------------------

  # Compute stuff
  estimated <- suppressWarnings(do.call(marginaleffects::avg_slopes, fun_args))
  vcov_slopes <- .safe(stats::vcov(estimated))

  # Fourth step: back-transform response --------------------------------------
  # ---------------------------------------------------------------------------

  # transform reponse?
  if (isTRUE(transform)) {
    trans_fun <- insight::get_transformation(model, verbose = FALSE)$inverse
  } else {
    trans_fun <- transform
  }
  # if we have back-transformation, do that, but remove standard errors
  # these are no longer correct
  if (!is.null(trans_fun)) {
    estimated$estimate <- trans_fun(estimated$estimate)
    estimated$conf.low <- trans_fun(estimated$conf.low)
    estimated$conf.high <- trans_fun(estimated$conf.high)
    estimated$std.error <- NULL
  }

  # Fifth step: post-processing marginal means----------------------------------
  # ---------------------------------------------------------------------------

  # filter "by" rows when we have "average" marginalization, because we don't
  # pass data grid in such situations - but we still created the data grid based
  # on the `by` variables, for internal use, for example filtering at this point
  estimated <- .filter_datagrid_average(estimated, estimate, datagrid, datagrid_info)

  # Last step: Save information in attributes  --------------------------------
  # ---------------------------------------------------------------------------

  estimated <- .add_attributes(
    estimated,
    model_info = model_info,
    info = c(
      datagrid_info,
      list(
        trend = myargs$trend,
        datagrid = datagrid,
        coef_name = "Slope",
        p_adjust = p_adjust,
        ci = ci,
        transform = !is.null(transform),
        keep_iterations = keep_iterations,
        vcov = vcov_slopes
      )
    )
  )
  class(estimated) <- unique(c("marginaleffects_slopes", class(estimated)))

  # adjust p-values
  if (!model_info$is_bayesian) {
    estimated <- .p_adjust(model, estimated, p_adjust, verbose, ...)
  }

  estimated
}


# =========================================================================
# HELPERS (guess arguments) -----------------------------------------------
# =========================================================================

#' @keywords internal
.guess_marginaltrends_arguments <- function(model,
                                            trend = NULL,
                                            by = NULL,
                                            verbose = TRUE,
                                            ...) {
  # Gather info
  model_data <- insight::get_data(model, verbose = FALSE)
  predictors <- intersect(
    colnames(model_data),
    insight::find_predictors(model, effects = "fixed", flatten = TRUE, ...)
  )

  # Guess arguments
  if (is.null(trend)) {
    trend <- predictors[sapply(model_data[predictors], is.numeric)][1]
    if (!length(trend) || is.na(trend)) {
      insight::format_error("Model contains no numeric predictor. Please specify `trend`.")
    }
    if (verbose) {
      insight::format_alert(paste0("No numeric variable was specified for slope estimation. Selecting `trend = \"", trend, "\"`.")) # nolint
    }
  }

  # check that we have only one predictor
  if (length(trend) > 1) {
    if (verbose) {
      insight::format_alert(paste0(
        "More than one numeric variable was selected for slope estimation. Keeping only `", trend[1], "`. ", # nolint
        "If you want to estimate the slope of `", trend[1], "` at different values of `", trend[2], "`, use `by=\"", trend[2], "\"` instead." # nolint
      ))
    }
    trend <- trend[1]
  }

  # check if user provided values in `trend`, e.g. `trend=1:10`. We then pass
  # this argument to also create a data grid, but we also need to "clean" trend
  if (grepl("=", trend, fixed = TRUE)) {
    range <- trend
    trend <- gsub("=.*", "\\1", trend)
  } else {
    range <- NULL
  }

  # make sure range in `trend` is not also in `by`
  if (!is.null(by) && !is.null(range) && startsWith(by, trend)) {
    insight::format_error(
      paste0(
        "To calculate average marginal effects over a range of `", trend, "` ",
        "values, use `trend=\"", trend, "=seq(1, 3, 0.1)\"` (or similar) and omit `",
        trend, "` from the `by` argument."
      ),
      paste0(
        "To get marginal effects at specific `", trend, "` values, use `trend=\"",
        trend, "\"` along with `by=\"", trend, "=c(1, 3, 5)\"`."
      )
    )
  }

  list(trend = trend, range = range, by = NULL)
}


#' @keywords internal
.marginaleffects_slopes <- function() {
  c("dY/dX", "eY/eX", "eY/dX", "dY/eX")
}