R/plot.R

In WMAP: Weighted Meta-Analysis with Pseudo-Populations

# ======================================================================================
# WMAP PACKAGE - PLOT METHODS
# ======================================================================================

# This file contains plot methods for WMAP objects:
# - plot.causal_estimates: Boxplot of bootstrap ESS distributions

# ======================================================================================
# VERSION HISTORY
# ======================================================================================

# Version 1.3 - Refactoring & Bug Fixes
#               * Fixed hardcoded B variable bug
#               * Updated variable naming (B -> n_bootstrap, ESS.mt -> ess_values)
#               * Added roxygen2 documentation to all functions
#               * Chnaged all = to <- for assignments
#               * Added additional examples

# ======================================================================================
# IMPORTS
# ======================================================================================

#' @importFrom ggplot2 aes theme element_text ggtitle xlab ylab ylim ggplot geom_boxplot
#' @importFrom forcats fct_inorder
NULL

# ======================================================================================
# PLOTS
# ======================================================================================

#' Plot method for objects of class 'causal_estimates'
#'
#' Creates a boxplot showing the distribution of effective sample size (ESS)
#' across bootstrap samples for the specified weighting method.
#'
#' @param x An object of class 'causal_estimates' returned by \code{causal.estimate()}.
#' @param ... Additional arguments:
#' \describe{
#'   \item{y_limit}{The y-axis range. Default is \code{c(0,50)}.}
#'   \item{color}{The boxplot color. Default is "red".}
#' }
#'
#' @return A boxplot of percent sample ESS for a specific weighting method (\code{FLEXOR}, \code{IC}, or \code{IGO}
#'
#' @examples
#' \dontrun{
#' data(demo)
#' result <- causal.estimate(S, Z, X, Y, B = 5, method = "FLEXOR", naturalGroupProp)
#' plot(result)
#' plot(result, y_limit = c(0, 100), color = "steelblue")
#' }
#'
#' @export
plot.causal_estimates <- function(x, ...) {

  # extract additional arguments, and define defaults
  args <- list(...)
  y_limit <- if ("y_limit" %in% names(args)) args$y_limit else c(0, 100)
  color <- if ("color" %in% names(args)) args$color else "red"

  # extract bootstrap ESS values
  ess_values <- x$collatedESS

  # calculate number of bootstrap samples from data
  n_bootstrap <- length(ess_values)

  # create df for plotting
  df <- data.frame(
    ESS = as.vector(ess_values),
    method = rep(x$method, times = n_bootstrap)
  )

  plot_obj <- ggplot(df, aes(x = fct_inorder(method), y = ESS)) +
    geom_boxplot(fill = color) +
    xlab("Method") +
    ylab("Percent ESS") +
    ggtitle("Effective Sample Sizes") +
    ylim(y_limit[1], y_limit[2]) +
    theme(
      axis.text = element_text(size = 12),
      axis.title = element_text(size = 14, face = "bold"),
      plot.title = element_text(hjust = 0.5, size = 16, face = "bold")
    )

  plot_obj
}