R/plot.check_predictions.R

In see: Model Visualisation Toolbox for 'easystats' and 'ggplot2'

#' @export
data_plot.performance_pp_check <- function(x, type = "density", ...) {
  # 1. Handle data from bayesplot ----------------------------------------------
  # If the data comes from "bayesplot::pp_check()", it is already in the correct
  # shape. We just assign classes and attributes and return early.
  if (isTRUE(attributes(x)$is_stan) && type != "density") {
    class(x) <- c("data_plot", "see_performance_pp_check", "data.frame")
    attr(x, "info") <- list(
      xlab = attr(x, "response_name"),
      ylab = ifelse(identical(type, "density"), "Density", "Counts"),
      title = "Posterior Predictive Check",
      check_range = attr(x, "check_range"),
      bandwidth = attr(x, "bandwidth"),
      model_info = attr(x, "model_info")
    )
    return(x)
  }

  columns <- colnames(x)

  # 2. Prepare ordinal models --------------------------------------------------
  # For ordinal models, the simulated response might be numeric, while the
  # original response is an ordered factor. We fix the simulations to match here.
  if (is.ordered(x$y) && is.numeric(x$sim_1)) {
    sims <- startsWith(colnames(x), "sim_")
    x[sims] <- lapply(x[sims], function(i) {
      factor(i, labels = levels(x$y), ordered = TRUE)
    })
  }

  # 3. Reshape data ------------------------------------------------------------
  # Convert data from wide to long format for ggplot2 compatibility
  dataplot <- stats::reshape(
    x,
    times = columns,
    timevar = "key",
    v.names = "values",
    varying = list(columns),
    direction = "long"
  )

  # Ensure factor values are treated as characters to avoid grouping issues
  if (is.factor(dataplot[["values"]])) {
    dataplot[["values"]] <- as.character(dataplot[["values"]])
  }

  # Drop the last column (usually the 'id' column generated by reshape)
  dataplot <- dataplot[, 1:(ncol(dataplot) - 1), drop = FALSE]

  # Relabel the keys for clearer plotting legends
  dataplot$key[dataplot$key != "y"] <- "Model-predicted data"
  dataplot$key[dataplot$key == "y"] <- "Observed data"

  # Create a grouping variable tying the long rows back to original columns
  dataplot$grp <- rep(seq_len(ncol(x)), each = nrow(x))

  # 4. Attach metadata ---------------------------------------------------------
  attr(dataplot, "info") <- list(
    xlab = attr(x, "response_name"),
    ylab = ifelse(identical(type, "density"), "Density", "Counts"),
    title = "Posterior Predictive Check",
    check_range = attr(x, "check_range"),
    bandwidth = attr(x, "bandwidth"),
    model_info = attr(x, "model_info")
  )

  class(dataplot) <- unique(c(
    "data_plot",
    "see_performance_pp_check",
    class(dataplot)
  ))

  dataplot
}


# Plot --------------------------------------------------------------------

#' Plot method for posterior predictive checks
#'
#' The `plot()` method for the `performance::check_predictions()` function.
#'
#' @param alpha_line Numeric value specifying alpha of lines indicating `yrep`.
#' @param theme A ggplot2-theme function, e.g. `theme = theme_lucid()` or
#' `theme = ggplot2::theme_dark()`.
#' @param type Plot type for the posterior predictive checks plot. Can be `"density"`
#' (default), `"discrete_dots"`, `"discrete_interval"` or `"discrete_both"` (the
#' `discrete_*` options are appropriate for models with discrete - binary, integer
#' or ordinal etc. - outcomes).
#' @param x_limits Numeric vector of length 2 specifying the limits of the x-axis.
#' If not `NULL`, will zoom in the x-axis to the specified limits.
#' @inheritParams data_plot
#' @inheritParams plot.see_check_normality
#' @inheritParams plot.see_parameters_distribution
#'
#' @return A ggplot2-object.
#'
#' @seealso See also the vignette about [`check_model()`](https://easystats.github.io/performance/articles/check_model.html).
#'
#' @examples
#' library(performance)
#'
#' model <- lm(Sepal.Length ~ Species * Petal.Width + Petal.Length, data = iris)
#' check_predictions(model)
#'
#' # dot-plot style for count-models
#' d <- iris
#' d$poisson_var <- rpois(150, 1)
#' model <- glm(
#'   poisson_var ~ Species + Petal.Length + Petal.Width,
#'   data = d,
#'   family = poisson()
#' )
#' out <- check_predictions(model)
#' plot(out, type = "discrete_dots")
#' @export
print.see_performance_pp_check <- function(
  x,
  linewidth = 0.5,
  size_point = 2,
  size_bar = 0.7,
  size_axis_title = base_size,
  size_title = 12,
  base_size = 10,
  alpha_line = 0.15,
  theme = NULL,
  colors = unname(social_colors(c("green", "blue"))),
  type = "density",
  x_limits = NULL,
  ...
) {
  orig_x <- x
  check_range <- isTRUE(attributes(x)$check_range)

  p1 <- .prepare_pp_check(
    x,
    linewidth = linewidth,
    size_point = size_point,
    size_bar = size_bar,
    size_axis_title = size_axis_title,
    size_title = size_title,
    base_size = base_size,
    alpha_line = alpha_line,
    theme = theme,
    colors = colors,
    type = type,
    x_limits = x_limits,
    ...
  )

  # Build and append range plot if requested
  if (isTRUE(check_range)) {
    p2 <- .plot_pp_check_range(orig_x, size_bar, colors = colors)
    graphics::plot(plots(p1, p2, n_columns = 1))
  } else {
    suppressWarnings(graphics::plot(p1))
  }

  invisible(orig_x)
}


#' @rdname print.see_performance_pp_check
#' @export
plot.see_performance_pp_check <- function(
  x,
  linewidth = 0.5,
  size_point = 2,
  size_bar = 0.7,
  size_axis_title = base_size,
  size_title = 12,
  base_size = 10,
  alpha_line = 0.15,
  theme = NULL,
  colors = unname(social_colors(c("green", "blue"))),
  type = "density",
  x_limits = NULL,
  ...
) {
  check_range <- isTRUE(attributes(x)$check_range)

  p1 <- .prepare_pp_check(
    x,
    linewidth = linewidth,
    size_point = size_point,
    size_bar = size_bar,
    size_axis_title = size_axis_title,
    size_title = size_title,
    base_size = base_size,
    alpha_line = alpha_line,
    theme = theme,
    colors = colors,
    type = type,
    x_limits = x_limits,
    default_theme = ggplot2::theme_grey(),
    ...
  )

  # Combine with range plot if requested, otherwise return just p1
  if (isTRUE(check_range)) {
    p2 <- .plot_pp_check_range(x, size_bar, colors = colors)
    plots(p1, p2)
  } else {
    p1
  }
}


# internal function, used by both the `print()` and `plot()` -----------------
.prepare_pp_check <- function(
  x,
  linewidth = 0.5,
  size_point = 2,
  size_bar = 0.7,
  size_axis_title = base_size,
  size_title = 12,
  base_size = 10,
  alpha_line = 0.15,
  theme = NULL,
  colors = unname(social_colors(c("green", "blue"))),
  type = "density",
  x_limits = NULL,
  default_theme = NULL,
  ...
) {
  # Save original object for the range plot
  plot_type <- attributes(x)$type
  is_stan <- attributes(x)$is_stan

  theme <- .set_default_theme(
    x,
    theme,
    base_size,
    size_axis_title,
    size_title,
    default_theme
  )

  # Validate and set the plot type
  if (
    missing(type) &&
      !is.null(plot_type) &&
      plot_type %in%
        c("density", "discrete_dots", "discrete_interval", "discrete_both")
  ) {
    type <- plot_type
  } else {
    type <- insight::validate_argument(
      type,
      c("density", "discrete_dots", "discrete_interval", "discrete_both")
    )
  }

  # Ensure the data is prepared for plotting
  if (!inherits(x, "data_plot")) {
    x <- data_plot(x, type)
  }

  # Build main plot
  .plot_pp_check(
    x,
    linewidth = linewidth,
    size_point = size_point,
    alpha_line = alpha_line,
    theme = theme,
    colors = colors,
    base_size = base_size,
    size_title = size_title,
    size_axis_title = size_axis_title,
    type = type,
    x_limits = x_limits,
    is_stan = is_stan,
    ...
  )
}


.plot_pp_check <- function(
  x,
  linewidth,
  size_point,
  alpha_line,
  theme = NULL,
  base_size = 10,
  size_axis_title = 10,
  size_title = 12,
  colors,
  type = "density",
  x_limits = NULL,
  is_stan = NULL,
  ...
) {
  info <- attr(x, "info")

  theme <- .set_default_theme(
    x,
    theme,
    base_size,
    size_axis_title,
    size_title,
    default_theme = ggplot2::theme_grey()
  )

  # 1. Dispatch for Stan objects -----------------------------------------------
  # Discrete plot type from "bayesplot::pp_check()" returns a different data
  # structure, so we pass it to a dedicated handler
  if (isTRUE(is_stan) && type != "density") {
    return(.plot_check_predictions_stan_dots(
      x,
      colors,
      info,
      linewidth,
      size_point,
      alpha_line,
      ...
    ))
  }

  # Default bandwidth for smoothing density curves
  bandwidth <- info$bandwidth
  if (is.null(bandwidth)) {
    bandwidth <- "nrd"
  }

  # Check if model properties suggest using discrete dots instead of lines
  minfo <- info$model_info
  suggest_dots <- (minfo$is_bernoulli ||
    minfo$is_count ||
    minfo$is_ordinal ||
    minfo$is_categorical)

  # 2. Dispatch Based on Type --------------------------------------------------
  if (
    !is.null(type) &&
      type %in% c("discrete_dots", "discrete_interval", "discrete_both") &&
      suggest_dots
  ) {
    # Discrete dot plots (typically for integer/count outcomes)
    out <- .plot_check_predictions_dots(
      x,
      colors,
      info,
      linewidth,
      size_point,
      alpha_line,
      type,
      ...
    )
  } else {
    # Warn user if they are using density plots for discrete data
    if (suggest_dots) {
      insight::format_alert(
        "The model has an integer or a discrete response variable.",
        "It is recommended to switch to a dot-plot style, e.g. `plot(check_model(model), type = \"discrete_dots\"`."
      )
    }
    # Standard continuous density plot
    out <- .plot_check_predictions_density(
      x,
      colors,
      info,
      linewidth,
      alpha_line,
      bandwidth,
      ...
    )
  }

  # 3. Apply Styling -----------------------------------------------------------
  dots <- list(...)
  if (isTRUE(dots[["check_model"]])) {
    out <- out + theme
  }

  # Adjust legend layout if required
  if (isTRUE(dots[["adjust_legend"]]) || isTRUE(info$check_range)) {
    out <- out +
      ggplot2::theme(
        legend.position = "bottom",
        legend.margin = ggplot2::margin(0, 0, 0, 0),
        legend.box.margin = ggplot2::margin(-5, -5, -5, -5)
      )
  }

  # Apply custom zoom limits
  if (!is.null(x_limits)) {
    out <- out + ggplot2::coord_cartesian(xlim = x_limits)
  }

  out
}


.plot_check_predictions_density <- function(
  x,
  colors,
  info,
  linewidth,
  alpha_line,
  bandwidth,
  ...
) {
  # Generates standard overlaid density curves
  ggplot2::ggplot(x) +
    ggplot2::stat_density(
      mapping = ggplot2::aes(
        x = .data$values,
        group = .data$grp,
        color = .data$key,
        linewidth = .data$key,
        alpha = .data$key
      ),
      geom = "line",
      position = "identity",
      bw = bandwidth
    ) +
    ggplot2::scale_y_continuous() +
    ggplot2::scale_color_manual(
      values = c(
        "Observed data" = colors[1],
        "Model-predicted data" = colors[2]
      )
    ) +
    ggplot2::scale_linewidth_manual(
      values = c(
        "Observed data" = 1.7 * linewidth,
        "Model-predicted data" = linewidth
      ),
      guide = "none"
    ) +
    ggplot2::scale_alpha_manual(
      values = c(
        "Observed data" = 1,
        "Model-predicted data" = alpha_line
      ),
      guide = "none"
    ) +
    ggplot2::labs(
      x = info$xlab,
      y = info$ylab,
      color = "",
      size = "",
      alpha = "",
      title = "Posterior Predictive Check",
      subtitle = "Model-predicted lines should resemble observed data line"
    ) +
    ggplot2::guides(
      color = ggplot2::guide_legend(reverse = TRUE),
      size = ggplot2::guide_legend(reverse = TRUE)
    )
}


.plot_check_predictions_dots <- function(
  x,
  colors,
  info,
  linewidth,
  size_point,
  alpha_line,
  type = "discrete_dots",
  ...
) {
  # 1. Data Aggregation and Reshaping ------------------------------------------
  # Force factor so "table()" generates frequencies for all levels (including 0 counts).
  # We need tables of same size to correctly bind data frames.
  x$values <- as.factor(x$values)
  x <- stats::aggregate(x["values"], list(grp = x$grp), table)
  x <- cbind(
    data.frame(key = "Model-predicted data", stringsAsFactors = FALSE),
    x
  )
  x <- cbind(x[1:2], as.data.frame(x[[3]]))
  x$key[nrow(x)] <- "Observed data" # The last row is always the original observed data

  # Convert aggregated wide data back to long format for ggplot
  x <- datawizard::data_to_long(
    x,
    select = -1:-2,
    names_to = "x",
    values_to = "count"
  )

  if (insight::n_unique(x$x) > 8) {
    x$x <- datawizard::to_numeric(x$x, dummy_factors = TRUE)
  }

  p1 <- p2 <- NULL

  # 2. Build Error Bars Layer (if requested) -----------------------------------
  if (!is.null(type) && type %in% c("discrete_interval", "discrete_both")) {
    # Helper to calculate median and confidence intervals
    centrality_dispersion <- function(i) {
      c(
        count = stats::median(i, na.rm = TRUE),
        unlist(bayestestR::ci(i)[c("CI_low", "CI_high")])
      )
    }

    x_errorbars <- stats::aggregate(
      x["count"],
      list(x$x),
      centrality_dispersion
    )

    x_errorbars <- cbind(x_errorbars[1], as.data.frame(x_errorbars[[2]]))
    colnames(x_errorbars) <- c("x", "count", "CI_low", "CI_high")
    x_errorbars <- cbind(
      data.frame(key = "Model-predicted data", stringsAsFactors = FALSE),
      x_errorbars
    )

    x_tmp <- x[x$key == "Observed data", ]
    x_tmp$CI_low <- NA
    x_tmp$CI_high <- NA
    x_tmp$grp <- NULL

    x_errorbars <- rbind(x_errorbars, x_tmp)

    p1 <- ggplot2::ggplot() +
      ggplot2::geom_pointrange(
        data = x_errorbars[x_errorbars$key == "Model-predicted data", ],
        mapping = ggplot2::aes(
          x = .data$x,
          y = .data$count,
          ymin = .data$CI_low,
          ymax = .data$CI_high,
          color = .data$key
        ),
        position = ggplot2::position_nudge(x = 0.2),
        size = 0.4 * size_point,
        linewidth = linewidth,
        stroke = 0,
        shape = 16
      ) +
      ggplot2::geom_point(
        data = x_errorbars[x_errorbars$key == "Observed data", ],
        mapping = ggplot2::aes(
          x = .data$x,
          y = .data$count,
          color = .data$key
        ),
        size = 1.5 * size_point,
        stroke = 0,
        shape = 16
      )
  }

  # 3. Build Raw Dots Layer (if requested) -------------------------------------
  if (!is.null(type) && type %in% c("discrete_dots", "discrete_both")) {
    if (is.null(p1)) {
      p2 <- ggplot2::ggplot()
    } else {
      p2 <- p1
    }

    p2 <- p2 +
      ggplot2::geom_point(
        data = x[x$key == "Model-predicted data", ],
        mapping = ggplot2::aes(
          x = .data$x,
          y = .data$count,
          group = .data$grp,
          color = .data$key
        ),
        alpha = alpha_line,
        position = ggplot2::position_jitter(width = 0.1, height = 0.02),
        size = 0.8 * size_point,
        stroke = 0,
        shape = 16
      ) +
      # Hidden layer specifically to force the legend behavior
      ggplot2::geom_point(
        data = x[x$key == "Observed data", ],
        mapping = ggplot2::aes(
          x = .data$x,
          y = .data$count,
          group = .data$grp,
          color = .data$key
        ),
        size = 0.8 * size_point
      ) +
      ggplot2::geom_point(
        data = x[x$key == "Observed data", ],
        mapping = ggplot2::aes(x = .data$x, y = .data$count),
        size = size_point,
        shape = 21,
        colour = "white",
        fill = colors[1]
      )
  }

  if (is.null(p2)) {
    p <- p1
  } else {
    p <- p2
  }

  if (type == "discrete_interval") {
    subtitle <- "Model-predicted intervals should include observed data points"
  } else {
    subtitle <- "Model-predicted points should be close to observed data points"
  }

  p +
    ggplot2::scale_y_continuous() +
    ggplot2::scale_color_manual(
      values = c(
        "Observed data" = colors[1],
        "Model-predicted data" = colors[2]
      )
    ) +
    ggplot2::labs(
      x = info$xlab,
      y = info$ylab,
      color = "",
      size = "",
      alpha = "",
      title = "Posterior Predictive Check",
      subtitle = subtitle
    ) +
    ggplot2::guides(
      color = ggplot2::guide_legend(reverse = TRUE),
      size = ggplot2::guide_legend(reverse = TRUE)
    )
}


.plot_check_predictions_stan_dots <- function(
  x,
  colors,
  info,
  linewidth,
  size_point,
  alpha_line,
  ...
) {
  # Normalize group names to match our established color scales
  x$Group[x$Group == "y"] <- "Observed data"
  x$Group[x$Group == "Mean"] <- "Model-predicted data"

  # Sanity check: remove NA rows
  x <- x[!is.na(x$Count), ]

  p <- ggplot2::ggplot() +
    ggplot2::geom_pointrange(
      data = x[x$Group == "Model-predicted data", ],
      mapping = ggplot2::aes(
        x = .data$x,
        y = .data$Count,
        ymin = .data$CI_low,
        ymax = .data$CI_high,
        color = .data$Group
      ),
      position = ggplot2::position_nudge(x = 0.2),
      size = 0.4 * size_point,
      linewidth = linewidth,
      stroke = 0,
      shape = 16
    ) +
    ggplot2::geom_point(
      data = x[x$Group == "Observed data", ],
      mapping = ggplot2::aes(
        x = .data$x,
        y = .data$Count,
        color = .data$Group
      ),
      size = 1.5 * size_point,
      stroke = 0,
      shape = 16
    ) +
    ggplot2::scale_y_continuous() +
    ggplot2::scale_color_manual(
      values = c(
        "Observed data" = colors[1],
        "Model-predicted data" = colors[2]
      )
    ) +
    ggplot2::labs(
      x = info$xlab,
      y = info$ylab,
      color = "",
      size = "",
      alpha = "",
      title = "Posterior Predictive Check",
      subtitle = "Model-predicted intervals should include observed data points"
    ) +
    ggplot2::guides(
      color = ggplot2::guide_legend(reverse = TRUE),
      size = ggplot2::guide_legend(reverse = TRUE)
    )

  p
}


.plot_pp_check_range <- function(
  x,
  size_bar = 0.7,
  colors = unname(social_colors(c("green", "blue")))
) {
  # Calculate min and max for the original observed data
  original <- data.frame(
    x = c(min(x$y), max(x$y)),
    group = factor(c("Minimum", "Maximum"), levels = c("Minimum", "Maximum")),
    color = "Observed data",
    stringsAsFactors = FALSE
  )

  # Calculate min and max for the replicated/simulated data runs
  replicated <- rbind(
    data.frame(
      x = vapply(x[which(names(x) != "y")], min, numeric(1)),
      group = "Minimum",
      color = "Model-predicted data",
      stringsAsFactors = FALSE
    ),
    data.frame(
      x = vapply(x[which(names(x) != "y")], max, numeric(1)),
      group = "Maximum",
      color = "Model-predicted data",
      stringsAsFactors = FALSE
    )
  )
  replicated$group <- factor(replicated$group, levels = c("Minimum", "Maximum"))

  p <- ggplot2::ggplot(
    replicated,
    ggplot2::aes(x = .data$x, group = .data$group)
  ) +
    ggplot2::facet_wrap(~group, scales = "free_x")

  # Choose geometric shape based on the number of unique integer breaks
  if (insight::n_unique(replicated$x) <= 12) {
    p <- p + ggplot2::geom_bar(width = size_bar, fill = colors[2], color = NA)
  } else if (.is_integer(replicated$x)) {
    p <- p +
      ggplot2::geom_bar(width = size_bar, fill = colors[2], color = NA) +
      ggplot2::scale_x_continuous(
        n.breaks = round(insight::n_unique(replicated$x) / 4)
      )
  } else {
    p <- p +
      ggplot2::geom_histogram(binwidth = size_bar, fill = colors[2], color = NA)
  }

  p +
    ggplot2::geom_vline(
      data = original,
      mapping = ggplot2::aes(xintercept = .data$x),
      color = colors[1],
      linewidth = 1
    ) +
    ggplot2::labs(
      x = NULL,
      y = NULL,
      subtitle = "Model-predicted extrema should contain observed data extrema"
    )
}