R/visualisation_recipe.estimate_means.R

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

#' @rdname visualisation_recipe.estimate_predicted
#'
#' @examples
#' # ==============================================
#' # estimate_means
#' # ==============================================
#' if (require("ggplot2")) {
#'   # Simple Model ---------------
#'   x <- estimate_means(lm(Sepal.Width ~ Species, data = iris))
#'   layers <- visualisation_recipe(x)
#'   layers
#'   plot(layers)
#' }
#' \donttest{
#' if (require("ggplot2")) {
#'   # Customize aesthetics
#'   layers <- visualisation_recipe(x,
#'     jitter = list(width = 0.03, color = "red"),
#'     line = list(linetype = "dashed")
#'   )
#'   plot(layers)
#'
#'   # Customize raw data
#'   plot(visualisation_recipe(x, show_data = c("violin", "boxplot", "jitter")))
#'
#'   # Two levels ---------------
#'   data <- mtcars
#'   data$cyl <- as.factor(data$cyl)
#'   data$new_factor <- as.factor(rep(c("A", "B"), length.out = nrow(mtcars)))
#'
#'   model <- lm(mpg ~ new_factor * cyl * wt, data = data)
#'   x <- estimate_means(model, at = c("new_factor", "cyl"))
#'   plot(visualisation_recipe(x))
#'
#'   # Modulations --------------
#'   x <- estimate_means(model, at = c("new_factor", "wt"))
#'   plot(visualisation_recipe(x))
#'
#'   x <- estimate_means(model, at = c("new_factor", "cyl", "wt"))
#'   plot(visualisation_recipe(x))
#'
#'   #'   # GLMs ---------------------
#'   data <- data.frame(vs = mtcars$vs, cyl = as.factor(mtcars$cyl))
#'   x <- estimate_means(glm(vs ~ cyl, data = data, family = "binomial"))
#'   plot(visualisation_recipe(x))
#' }
#' }
#' @export
visualisation_recipe.estimate_means <- function(x,
                                                show_data = "jitter",
                                                point = NULL,
                                                jitter = point,
                                                boxplot = NULL,
                                                violin = NULL,
                                                line = NULL,
                                                pointrange = NULL,
                                                labs = NULL,
                                                ...) {
  info <- attributes(x)
  layers <- list()


  # Main aesthetics -----------------
  data <- as.data.frame(x)
  y <- info$response
  color <- NULL
  alpha <- NULL

  levels <- info$at[info$at %in% names(data[!sapply(data, is.numeric)])]
  modulate <- info$at[info$at %in% names(data[sapply(data, is.numeric)])]
  x1 <- levels[1]
  if (length(levels) > 1L) {
    color <- levels[2]
    if (length(levels) > 2L) {
      # TODO: add facetting (needs updating see::geom_from_list to work with facets)
      insight::format_warning("Cannot deal with more than 2 levels variables for now. Other ones will be omitted.")
    }
  }
  if (!is.null(modulate) && length(modulate) > 0L) {
    alpha <- modulate[1]
    if (length(modulate) > 1L) {
      insight::format_warning("Cannot deal with more than 2 modulate variables for now. Other ones will be omitted.")
    }
  }


  # Layers -----------------------
  l <- 1

  # Show data (points, boxplot, violin, etc.)
  if (!is.null(show_data) && all(show_data != "none") && all(show_data != FALSE)) {
    # Default changes for binomial models
    shape <- 16
    stroke <- 0
    if (insight::model_info(info$model)$is_binomial) {
      shape <- "|"
      stroke <- 1
    }

    rawdata <- .visualisation_recipe_getrawdata(x)
    for (i in show_data) {
      if (i %in% c("point", "points", "jitter")) {
        layers[[paste0("l", l)]] <- .visualisation_predicted_points(
          rawdata,
          x1,
          y,
          color,
          shape = shape,
          stroke = stroke,
          type = i,
          width = 0.1,
          height = 0,
          point = jitter
        )
      } else if (i == "boxplot") {
        layers[[paste0("l", l)]] <-
          .visualisation_means_boxplot(rawdata, x1, y, color, type = "boxplot", boxplot = boxplot)
      } else if (i == "violin") {
        layers[[paste0("l", l)]] <-
          .visualisation_means_boxplot(rawdata, x1, y, color, type = "violin", boxplot = violin)
      } else {
        insight::format_error("`show_data` can only be some of 'points', 'boxplot', 'violin'. Check spelling.")
      }
      l <- l + 1
    }
  }

  # Line
  layers[[paste0("l", l)]] <- .visualisation_means_line(
    data,
    x1,
    y = info$coef_name[1],
    color = color,
    alpha = alpha,
    line = line
  )
  l <- l + 1

  # Pointrange
  layers[[paste0("l", l)]] <- .visualisation_means_pointrange(
    data,
    x1,
    y = info$coef_name[1],
    color = color,
    alpha = alpha,
    pointrange = pointrange
  )
  l <- l + 1

  # Labs
  layers[[paste0("l", l)]] <- .visualisation_means_labs(info, x1, y, labs = labs)

  # Out
  class(layers) <- unique(c("visualisation_recipe", "see_visualisation_recipe", class(layers)))
  attr(layers, "data") <- data
  layers
}


# Layer - Violin / boxplot ------------------------------------------------------------

.visualisation_means_boxplot <- function(raw_data, x1, y, color, type = "boxplot", boxplot = NULL) {
  out <- list(
    data = as.data.frame(raw_data),
    geom = type,
    aes = list(x = x1, y = y, fill = color)
  )

  if (type == "boxplot") {
    out$outlier.shape <- NA
  }
  if (!is.null(boxplot)) out <- utils::modifyList(out, boxplot) # Update with additional args
  out
}

# Layer - Line -------------------------------------------------------------

.visualisation_means_line <- function(data, x1, y, color, alpha, line = NULL) {
  if (is.null(color) && is.null(alpha)) {
    group <- 1
  } else if (!is.null(color) && is.null(alpha)) {
    group <- color
  } else if (is.null(color) && !is.null(alpha)) {
    group <- alpha
  } else {
    group <- paste0("interaction(", alpha, ", ", color, ")")
  }

  out <- list(
    geom = "line",
    data = data,
    aes = list(
      y = y,
      x = x1,
      color = color,
      group = group,
      alpha = alpha
    )
  )

  if (!is.null(color) || !is.null(alpha)) {
    out$position <- "dodge"
    out$width <- 0.1
  }
  if (!is.null(line)) out <- utils::modifyList(out, line) # Update with additional args
  out
}

# Layer - Pointrange -------------------------------------------------------

.visualisation_means_pointrange <- function(data, x1, y, color, alpha, pointrange = NULL) {
  out <- list(
    geom = "pointrange",
    data = data,
    aes = list(
      y = y,
      x = x1,
      ymin = "CI_low",
      ymax = "CI_high",
      color = color,
      alpha = alpha
    )
  )

  if (!is.null(color) || !is.null(alpha)) {
    out$position <- "dodge"
    out$width <- 0.1
  }
  if (!is.null(pointrange)) out <- utils::modifyList(out, pointrange) # Update with additional args
  out
}



# Layer - Labels --------------------------------------------------------------

.visualisation_means_labs <- function(info, x1, y, labs = NULL) {
  if (all(info$coef_name == "Probability")) {
    title <- "Estimated Mean Probabilities"
  } else {
    title <- "Estimated Means"
  }

  out <- list(
    geom = "labs",
    x = x1,
    y = y,
    title = paste0(title, " (", format(insight::find_formula(info$model)), ")")
  )
  if (!is.null(labs)) out <- utils::modifyList(out, labs) # Update with additional args
  out
}