R/visualization.R

In PowRPriori: Power Analysis via Data Simulation for (Generalized) Linear Mixed Effects Models

#' Visualize Simulation Data or Power Simulation Results
#'
#' @description
#' Generic plotting function with methods for different objects.
#' - When used on an lme4-style `formula`, it simulates and plots a single plausible dataset.
#' - When used on a `PowRPriori` object, it plots either a power curve from the object or a dataset from the simulation.
#'
#' The plotting of the dataset is designed to aid in evaluating whether the simulated data is plausible in the context
#' of the desired study design and model specifications. It can help determine whether the chosen parameters are sensible or might
#' need some adapting. The power curve, plotted from the resulting `PowRPriori` object of the `power_sim` function visualizes the iterations
#' of the simulation across the different sample sizes for which the power was calculated during simulation.
#'
#' @details
#' The parameters `x_var, group_var, color_var and facet_var` are `NULL` by default. If left `NULL`, they are automatically extracted from the `PowRPriori` object
#' or the `design` object.
#'
#'
#' @param object The object to base the plot on. Can be either a `PowRPriori` object or an `lme4`-style formula
#' @param ... Additional arguments (not used).
#' @param type The type of plot to create. For `PowRPriori` objects, can be
#'   `"power_curve"` or `"data"`. For `lme4`-style formulas, only `"data"` is available.
#' @param design A `PowRPriori_design` object.
#' @param fixed_effects,random_effects Lists of effect parameters.
#' @param family The model family. Defaults to `"gaussian"`, other possible values are `"binomial"` or `"poisson"`.
#' @param center Controls if centering is applied to the predictors prior to plotting. Defaults to `"auto"`.
#' @param n The sample size to simulate.
#' @param x_var,group_var,color_var,facet_var Strings specifying variables for plot aesthetics.
#' @param n_data_points The maximum number of trajectories in spaghetti plots.
#'
#' @return A `ggplot` object.
#' @export
#'
#' @examples
#' # 1. Plot prior to simulation to check data plausibility
#' design <- define_design(
#'   sample_size = list(subject = 30),
#'   between = list(group = c("Control", "Treatment")),
#'   within = list(time = c("pre", "post"))
#' )
#'
#' fixed_effects <- list(
#'   `(Intercept)` = 10,
#'   groupTreatment = 2,
#'   timepost = 1,
#'   `groupTreatment:timepost` = 3
#' )
#'
#' random_effects <- list(
#'   subject = list(`(Intercept)` = 3),
#'   sd_resid = 3
#' )
#'
#' plot_sim_model(
#'   y ~ group * time + (1|subject),
#'   design = design,
#'   fixed_effects = fixed_effects,
#'   random_effects = random_effects
#' )
#' \donttest{
#' # 2. Plot from PowRPriori object after simulation
#'   power_results <- power_sim(
#'     formula = y ~ group * time + (1|subject),
#'     design = design,
#'     fixed_effects = fixed_effects,
#'     random_effects = random_effects,
#'     test_parameter = "groupTreatment:timepost",
#'     center = TRUE,
#'     n_start = 100,
#'     n_increment = 5,
#'     # The parameters defined in this example are to ensure low runtime.
#'     # Adapt these parameters!
#'     n_sims = 10,
#'     max_simulation_steps = 1,
#'     parallel_plan = "sequential"
#'   )
#'
#'   # Power curve
#'   plot_sim_model(power_results, type = "power_curve")
#'
#'   # Plot sample data with automated aesthetics extraction
#'   plot_sim_model(power_results, type = "data")
#'}
plot_sim_model <- function(object, type, design, fixed_effects, random_effects, family,
                           center, n, x_var, group_var, color_var, facet_var,
                           n_data_points, ...) {
  UseMethod("plot_sim_model")
}

#' @param n The total sample size to simulate for the plot (overwrites the lowest design level).
#'
#' @rdname plot_sim_model
#' @export
plot_sim_model.formula <- function(object, type = "data", design, fixed_effects, random_effects, family = "gaussian",
                                   center = "auto", n = NULL, x_var = NULL, group_var = NULL,
                                   color_var = NULL, facet_var = NULL,
                                   n_data_points = 10, ...) {
  if(type == "data"){

    temp_design <- design

    if (!missing(n) && !is.null(n)) {
      target_lvl <- names(temp_design$sample_size)[length(temp_design$sample_size)]
      temp_design$sample_size[[target_lvl]] <- n
    }

    sim_data <- .create_design_matrix(design = temp_design, formula = object)

    if (identical(center, "auto")) {
      is_centered_attr <- attr(fixed_effects, "is_centered")
      if (!is.null(is_centered_attr)) {
        center <- is_centered_attr
      } else {
        center <- FALSE
      }
    }

    if (identical(center, "auto")) {
      is_centered_attr <- attr(fixed_effects, "is_centered")

      if (!is.null(is_centered_attr)) {
        center <- is_centered_attr
      } else{
        hasInteraction <- any(grepl(":", attr(stats::terms(nobars(object)), "term.labels")))

        if(hasInteraction){
          stop("Your model contains interaction effects and you manually specified your fixed effects without specifying if the predictors should be centered. Since centering changes the interpretation and power of main effects, PowRPriori cannot safely guess your intent. Please explicitly add 'center = TRUE' (recommended if your weights are effect-coded) or 'center = FALSE' (if your weights use dummy-coding) to your power_sim() call.", call. = FALSE)
        } else {
          center <- FALSE
        }
      }
    }

    sim_data <- .simulate_outcome(sim_data, formula = object, fixed_effects = fixed_effects,
                                  sds_random = random_effects, family = family)

    .plot_data(data = sim_data, design = temp_design, formula = object, x_var = x_var, family = family,
               group_var = group_var, color_var = color_var,
               facet_var = facet_var, n_data_points = n_data_points)
  } else {
    stop(
      "When calling `plot_sim_model` with an lme4-style formula supplied to the `object` parameter, `type` must be set to `data`",
      call. = FALSE
    )
  }
}


#' @param type The type of plot to create: `"power_curve"` (default) or `"data"`
#'   (to visualize the sample data from the simulation).
#'
#' @rdname plot_sim_model
#' @export
plot_sim_model.PowRPriori <- function(object, type = "power_curve", design = NULL, fixed_effects = NULL,
                                      random_effects = NULL, family = NULL, center = NULL, n = NULL,
                                      x_var = NULL, group_var = NULL,
                                      color_var = NULL, facet_var = NULL,
                                      n_data_points = 10, ...) {
  if (type == "power_curve") {
    power_data <- object$power_table
    n_var <- names(power_data)[1]

    plot_data_long <- power_data %>%
      tidyr::pivot_longer(cols = tidyselect::starts_with(c("power", "ci_lower", "ci_upper")),
                          names_pattern = "(power|ci_lower|ci_upper)_(.*)",
                          names_to = c(".value", "parameter"))

    dodge <- ggplot2::position_dodge(width = ifelse(length(unique(plot_data_long$parameter))>1, 0.8, 0))
    p <- ggplot2::ggplot(plot_data_long,
                         ggplot2::aes(x = .data[[n_var]],
                                      y = .data$power,
                                      color = .data$parameter,
                                      group = .data$parameter)) +
      ggplot2::geom_line(linewidth = 1, position=dodge) +
      ggplot2::geom_point(size = 1.5, position=dodge) +
      ggplot2::geom_errorbar(ggplot2::aes(ymin=.data$ci_lower, ymax=.data$ci_upper), width=.4, position=dodge) +
      ggplot2::geom_hline(yintercept = object$power_crit, linetype = "dashed", color = "red") +
      ggplot2::scale_y_continuous(labels = scales::percent, limits = c(0, 1), breaks = seq(0,1,0.1)) +
      ggplot2::scale_x_continuous(limits = c(min(plot_data_long[[n_var]]) - 1, max(plot_data_long[[n_var]]) + 1)) +
      ggplot2::labs(title = "Power-Curve", y = "Power", x = n_var, color = "Parameter") +
      ggplot2::theme_minimal() + ggplot2::theme(legend.position = "bottom")
  } else if (type == "data") {
    p <- .plot_data(data = object$sample_data, design = object$design, formula = object$formula, family = object$family,
                    x_var = x_var, group_var = group_var, color_var = color_var,
                    facet_var = facet_var, n_data_points = n_data_points)
  } else {
    stop("Parameter 'type' can only be either 'power_curve' or 'data'", call. = F)
  }
  return(p)
}

#' Internal Data Plotting Engine
#'
#' @description
#' An internal helper function containing the logic to "intelligently" create
#' plots from simulated data. It automatically chooses between spaghetti plots and jitter/point-range plots depending on
#' the specified design and model family.It derives sensible defaults for plot aesthetics from the design,
#' if they are not supplied directly via the `plot_sim_model` function.
#'
#' @param data The data frame to plot.
#' @param design The `PowRPriori_design` object.
#' @param formula An lme4-style formula (e.g. `outcome ~ predictor1 * predictor2 + (1 | subject)`)
#' @param family The model family. Defaults to `"gaussian"`, other possible values are `"binomial"` or `"poisson"`.
#' @param x_var,group_var,color_var,facet_var Strings specifying variables for plot aesthetics.
#' @param n_data_points The maximum number of trajectories in spaghetti plots.
#'
#' @return A `ggplot` object.
.plot_data <- function(data, design, formula, family, x_var, group_var,
                       color_var, facet_var, n_data_points) {
  y_var <- all.vars(formula)[1]

  id_var <- ".PowR_id"

  between_vars <- .extract_plot_vars(design$between)
  within_vars <- .extract_plot_vars(design$within)

  if (is.null(group_var)) {
    group_var <- id_var
  }

  is_categorical <- function(v) {
    is.factor(data[[v]]) || is.character(data[[v]])
  }

  if (length(within_vars) > 0) {
    if (is.null(x_var)) {
      cat_within <- within_vars[sapply(within_vars, is_categorical)]
      if (length(cat_within) > 0) {
        x_var <- cat_within[1]
      } else {
        x_var <- within_vars[1]
      }
    }

    if (is.null(facet_var) && is.null(color_var) && length(between_vars) > 0) {
      cat_between <- between_vars[sapply(between_vars, is_categorical)]
      if (length(cat_between) > 0) {
        facet_var <- cat_between[1]
      }
    }

    vars_to_average <- setdiff(within_vars, c(group_var, x_var, color_var, facet_var))

    if (length(vars_to_average) > 0) {
      message(paste0("Note: Effects of '", paste(vars_to_average, collapse = ", "), "' averaged for plot."))
      data <- data %>%
        dplyr::group_by(dplyr::across(dplyr::all_of(unlist(purrr::compact(list(group_var, x_var, color_var, facet_var)))))) %>%
        dplyr::summarise(!!y_var := mean(.data[[y_var]]), .groups = "drop")
    }

    y_var <- all.vars(formula)[1]

    if(family == "binomial") {
      grouping_vars <- purrr::compact(list(x_var, color_var, facet_var, y_var))

      count_data <- data %>%
        dplyr::group_by(dplyr::across(dplyr::all_of(unlist(grouping_vars)))) %>%
        dplyr::summarise(n_obs = dplyr::n(), .groups = "drop")

      p <- ggplot2::ggplot(data, ggplot2::aes(x = .data[[x_var]], y = .data[[y_var]])) +
        ggplot2::geom_jitter(ggplot2::aes(color = if(!is.null(color_var)) .data[[color_var]] else NULL),
                             width = 0.25, height = 0.15, alpha = 0.25, size = 1.5) +
        ggplot2::geom_point(
          data = count_data,
          ggplot2::aes(size = .data$n_obs^2, fill = if(!is.null(color_var)) .data[[color_var]] else NULL),
          shape = 23,
          alpha = 1, fill = "deepskyblue3", color = "black", stroke = 0.8
        ) +
        ggplot2::labs(title = paste("Jitter-Plot for Simulated Data"), subtitle = paste("Family:", family),
                      x = x_var, y = y_var) +
        ggplot2::coord_cartesian(ylim = c(-0.1, 1.1)) +
        ggplot2::scale_y_continuous(breaks = c(0, 1)) +
        ggplot2::theme_minimal() + ggplot2::theme(legend.position = "bottom")

      if (!is.null(facet_var)) {
        p <- p + ggplot2::facet_wrap(ggplot2::vars(.data[[facet_var]]))
      }
    } else {
      if (!is.null(facet_var)) {
        plot_spaghetti_data <- data %>%
          dplyr::group_by(.data[[facet_var]]) %>%
          dplyr::filter(.data[[group_var]] %in% sample(unique(.data[[group_var]]),
                                                       min(n_data_points, dplyr::n_distinct(.data[[group_var]])))) %>%
          dplyr::ungroup()
      } else if (length(unique(data[[group_var]])) > n_data_points) {
        sampled_groups <- sample(unique(data[[group_var]]), n_data_points)
        plot_spaghetti_data <- data[data[[group_var]] %in% sampled_groups, ]
      } else {
        plot_spaghetti_data <- data
      }

      p <- ggplot2::ggplot(plot_spaghetti_data, ggplot2::aes(x = .data[[x_var]], y = .data[[y_var]]))
      if (!is.null(facet_var)) {
        p <- p + ggplot2::geom_line(ggplot2::aes(group = .data[[group_var]]), alpha = 0.4, color = "grey50") +
          ggplot2::facet_wrap(rlang::syms(facet_var)) +
          ggplot2::stat_summary(data = data, ggplot2::aes(group = 1), fun = mean, geom = "line", linewidth = 1.2, color = "black") +
          ggplot2::theme(legend.position = "none")
      } else if (!is.null(color_var)) {
        p <- p + ggplot2::geom_line(ggplot2::aes(group = .data[[group_var]], color = .data[[color_var]]), alpha = 0.6) +
          ggplot2::stat_summary(data = data, ggplot2::aes(group = .data[[color_var]], color = .data[[color_var]]),
                                fun = mean, geom = "line", linewidth = 1.2)
      } else {
        p <- p + ggplot2::geom_line(ggplot2::aes(group = .data[[group_var]]), alpha = 0.4) +
          ggplot2::stat_summary(data = data, ggplot2::aes(group = 1), fun = mean, geom = "line", linewidth = 1.2, color = "black")
      }
      p <- p + ggplot2::labs(title = paste("Spaghetti-Plot for Simulated Data"), x = x_var, y = y_var, color = color_var) +
        ggplot2::theme_minimal() + ggplot2::theme(legend.position = "bottom")
    }
  } else {
    # Cross-sectional plot (no within-variables)
    if (is.null(x_var)) {
      if (length(between_vars) > 0) {
        x_var <- between_vars[1]
      } else {
        x_var <- id_var
      }
    }

    p <- ggplot2::ggplot(data, ggplot2::aes(x = .data[[x_var]], y = .data[[y_var]])) +
      ggplot2::geom_jitter(width = 0.1, alpha = 0.2) +
      ggplot2::stat_summary(fun = mean, geom = "point", size = 4, shape = 18) +
      ggplot2::stat_summary(
        fun.min = function(y) mean(y) - 1.96 * sd(y) / sqrt(length(y)),
        fun.max = function(y) mean(y) + 1.96 * sd(y) / sqrt(length(y)),
        geom = "errorbar",
        width = 0.2
      ) +
      ggplot2::labs(title = paste("Simulated Distribution"), x = x_var, y = y_var) +
      ggplot2::theme_minimal() + ggplot2::theme(legend.position = "bottom")
  }

  return(p)
}


#' Internal Helper for Plotting Functions to Extract Variables from Design Object
#'
#' @param design_list A list from the design object (e.g., design$between)
#'
#' @returns A character vector of actual variable names
#'
#' @noRd
.extract_plot_vars <- function(design_list) {
  if (is.null(design_list)) return(character(0))
  unname(unlist(lapply(design_list, names)))
}