R/class_estimate__plot.R

In rcalinjageman/esci2: Effect Sizes and Confidence Intervals

plot_estimate <- function(
  estimate,
  rope = list(
    reference = NULL, 
    upper = NULL, 
    lower = NULL, 
    rope_units = "raw"
  ),
  plot_attributes = NULL,
  data_layout = c("random", "swarm", "none"), 
  data_spread = 0.25,
  error_layout = c("halfeye", "eye", "gradient", "none"),
  error_scale = 0.3,
  error_nudge = 0.35,
  ylim = c(NA, NA),
  breaks = 5,
  difference_axis_units = c("sd", "raw"),
  difference_axis_breaks = 5,
  y.axis.text = 10,
  y.axis.title = 12,
  x.axis.text = 10,
  x.axis.title = 12,
  ylab = "default",
  xlab = "default",
  ggtheme = NULL
) {
  
  
  # Input checks ---------------------------------------------------------------
  esci_assert_type(estimate, "is.estimate")
  data_layout <- match.arg(data_layout)
  error_layout <- match.arg(error_layout)
  error_normalize <- match.arg(error_normalize)
  difference_axis_units <- match.arg(difference_axis_units)
  
  if(is.null(ggtheme)) { ggtheme <- ggplot2::theme_classic()}
  
  if (is.null(ylab)) {ylab = "default"}
  if (is.null(xlab)) {xlab = "default"}
  
  plot_attributes <- esci_plot_attributes(plot_attributes)
  
  
  # Type of graphs -----------------------------------------------------------
  is_difference <- (estimate$properties$effect_size_category == "Difference")
  is_complex_difference <- (is_difference & nrow(estimate$overview) > 2)
  is_means <- (estimate$properties$effect_size_name == "M")
  plot_raw <- (!is.null(estimate$raw_data) & data_layout != "none")
  
  
  # Prep es_data
  if (!is.null(estimate$effect_sizes) & is_difference) {
    # Copy of effect sizes with only needed columns
    columns <- c("type", "effect", "effect_size", "df", "se")
    es_data <- estimate$effect_sizes[ , columns]
    # Swap rows so it is Reference, Comparison, Difference
    es_data <- es_data[c(2, 1, 3), ]
    # Store group
    es_data$group <- es_data$type
    # Set line_end to NA
    es_data$line_end <- NA
    # Set nudge to 0 or error_nudge
    es_data$nudge <- 0
    if (plot_raw & !is_complex_difference) {
      es_data[es_data$type != "Difference", ]$nudge <- error_nudge
    }
    
    # Store comparison and reference value
    #  and difference CI limits and pooled SD
    comparison_value <- 
      es_data[es_data$type == "Comparison", "effect_size"]
    reference_value <- 
      es_data[es_data$type == "Reference", "effect_size"]
    errorLower <- 
      es_data[es_data$type == "Difference", "lower"]
    errorUpper <- 
      es_data[es_data$type == "Difference", "upper"]
    pooled_sd <- 
      es_data[es_data$type == "Difference", "variability_component"]
    
    # Adjust difference effect sizes for sake of plotting
    es_data[es_data$type == "Difference", "effect_size"] <- 
      comparison_value
    
    # Handles calculations for plotting a difference in effect size
    weights <- if (is.null(estimate$properties$contrast))
      c(1)
    else
      estimate$properties$contrast
    
    comparison_levels <- names(weights)[which(weights > 0)]
    reference_levels <- names(weights)[which(weights < 0)]
    comparison_count <- length(comparison_levels)
    reference_count <- length(reference_levels)

    
    if (is_complex_difference) {
      # We will be plotting the 
      group_data <- data.frame(
        type = "Unused",
        effect = estimate$overview$group,
        effect_size = estimate$overview$m,
        df = estimate$overview$df,
        se = estimate$overview$se,
        group = estimate$overview$group,
        line_end = NA,
        nudge = error_nudge
      )
      group_data[group_data$effect %in% comparison_levels, ]$type <- "Comparison"
      group_data[group_data$effect %in% reference_levels, ]$type <- "Reference"
      
      es_data <- rbind(group_data, es_data)
    }
    
    # Es data has now been assembled, get info on group positioning
    
    es_data$group <- factor(es_data$group, levels = es_data$group)
    es_data$group_value <- as.integer(es_data$group)    
    
    comp_position <- es_data[es_data$x == "Comparison", "x_value"]
    ref_position <- es_data[es_data$x == "Reference", "x_value"]
    
    es_data[es_data$type == "Reference", "line_end"] <- ref_position
    es_data[es_data$type == "Comparison", "line_end"] <- comp_position    
    
    # Kludge to handle single-sample comparison
    es_data[is.na(es_data$df), "se"] <- .Machine$double.xmin
    es_data[is.na(es_data$df), "df"] <- 1
  }
  
  
  if (!is.null(estimate$raw_data)) {
    raw_data <- estimate$raw_data
    raw_data$group <- es_data$effect[
      match(unlist(raw_data$grouping_variable), es_data$effect)
    ]
    # raw_data$x <- es_data$type[
    #   match(unlist(raw_data$grouping_variable), es_data$effect)
    # ]
  }
  raw_data$x <- factor(raw_data$x, levels(es_data$x))

  
  # Data prep: Set graph attributes
  for (myattrib in names(plot_attributes)) {
    es_data[myattrib] <- esci_plot_match_attributes(
      es_data$type, plot_attributes[[myattrib]]
    )
    if (plot_raw) {
      raw_data[myattrib] <- esci_plot_match_attributes(
        raw_data$type, plot_attributes[[myattrib]]
      )
    }
  }
  
  
  
  # The actual plot -----------------------------------------------
  
  myplot <- ggplot2::ggplot(es_data, aes(x = x, y = effect_size))
  myplot <- myplot + geom_segment(
    data = es_data, 
    aes(x = x_value + nudge, xend = line_end + nudge, y = effect_size, yend = effect_size)
  )
  
  myplot <- myplot + ggbeeswarm::geom_beeswarm(data = raw_data, aes(x = x, y = outcome_variable))
  myplot <- myplot + scale_x_discrete(drop = FALSE)
  myplot <- myplot + ggdist::stat_dist_halfeye(
    data = es_data,
    orientation = 'vertical',
    aes(
      x = x,
      y = effect_size,
      dist = distributional::dist_student_t(
        df = df,
        mu = effect_size,
        sigma = se
      )
    ),
    position = position_nudge(x = es_data$nudge),
    .width = c(estimate$properties$conf_level),
    scale = 0.5
  )
  myplot
  
}