R/plot.R

In mixedbiastest: Bias Diagnostic for Linear Mixed Models

# ===== plot.R =====
#' Plot Method for Bias Diagnostic Results
#'
#' Plots the permutation distributions and observed test statistics for each fixed effect.
#'
#' @param x An object of class `"mixedbiastest"`.
#' @param bins Integer, number of bins for the histograms (default 30).
#' @param ... Additional arguments (currently not used).
#' @return A \code{ggplot} object (returned invisibly) showing permutation distributions for all fixed effects.
#' @method plot mixedbiastest
#' @importFrom ggplot2 ggplot geom_histogram geom_vline facet_wrap labs theme_minimal geom_text aes
#' @importFrom stats rnorm
#' @importFrom rlang .data
#' @export
plot.mixedbiastest <- function(x, bins = 30, ...) {
  permutation_values_list <- x$permutation_values
  results <- x$results_table
  k_list_names <- results$Fixed_Effect

  df_list <- list()
  annotations <- data.frame(Fixed_Effect = character(), label = character(), stringsAsFactors = FALSE)

  for (fe_name in k_list_names) {
    permutation_values <- permutation_values_list[[fe_name]]
    observed_value <- as.numeric(results$Bias_Estimate[results$Fixed_Effect == fe_name])

    if (length(permutation_values) == 0L) next

    if (length(unique(permutation_values)) == 1) {
      annotations <- rbind(annotations, data.frame(
        Fixed_Effect = fe_name,
        label = "No variation in permutation values",
        stringsAsFactors = FALSE
      ))
    }

    df_list[[fe_name]] <- data.frame(
      PermutationValues = permutation_values,
      Fixed_Effect = fe_name,
      Observed_Value = observed_value
    )
  }

  # If nothing to plot, return quietly
  if (!length(df_list)) {
    message("No permutation values to plot.")
    return(invisible(NULL))
  }

  combined_df <- do.call(rbind, df_list)
  combined_df$Fixed_Effect <- factor(combined_df$Fixed_Effect, levels = k_list_names)
  if (nrow(annotations) > 0) {
    annotations$Fixed_Effect <- factor(annotations$Fixed_Effect, levels = k_list_names)
  }

  p <- ggplot(combined_df, aes(x = .data$PermutationValues)) +
    geom_histogram(fill = "lightblue", color = "black", na.rm = TRUE, bins = bins) +
    geom_vline(aes(xintercept = .data$Observed_Value), color = "red",
               linetype = "dashed", size = 1) +
    facet_wrap(~ Fixed_Effect, scales = "free") +
    labs(title = "Bias Diagnostic for Fixed Effects",
         x = expression(hat(nu) %*% hat(eta)),
         y = "Frequency") +
    theme_minimal()

  if (nrow(annotations) > 0) {
    p <- p + geom_text(
      data = annotations,
      aes(
        x = mean(combined_df$PermutationValues, na.rm = TRUE),
        y = 0,
        label = .data$label
      ),
      color = "red",
      size = 3,
      hjust = 0.5,
      vjust = -1,
      inherit.aes = FALSE
    )
  }

  suppressWarnings(print(p))
  invisible(p)
}