R/plot_apportionment.R

In OceanTracer: Marine Source Unmixing Toolkit

#' Plot Apportionment Results
#'
#' Displays the relative contribution of sediment sources to each mixture as
#' density or violin plots. Designed to take the output of unmixing functions directly.
#'
#' @param data A data frame, typically the output from \code{unmix()}, containing
#'   the relative source contributions. The first two columns must be the mixture
#'   identifier (e.g. \code{id}) and the goodness-of-fit (\code{gof}), followed
#'   by one column per source. Only the source columns are plotted.
#' @param violin Logical. If \code{TRUE}, violin + boxplot charts are used
#'   instead of density plots. Default is \code{FALSE}.
#' @param bounds A numeric vector of length 2 specifying the lower and upper
#'   display bounds (the camera window). Default is \code{c(0.0, 1.0)}.
#' @param scaled Logical. If \code{TRUE}, density curves are area-scaled to a
#'   common maximum. Only used when \code{violin = FALSE}. Default is \code{TRUE}.
#' @param y_high Maximum value for the y-axis. Only used when
#'   \code{violin = FALSE} and \code{scaled = FALSE}. Default is \code{1}.
#' @param colors An optional named character vector of colors, one per source. Default is \code{NULL}.
#' @param vlines An optional named numeric vector specifying true/expected values 
#'   for specific sources. Draws a dashed vertical line to allow visual validation 
#'   of the model's accuracy (e.g., \code{c("SourceA" = 0.45)}). Default is \code{NULL}.
#' @param title Optional character string to add a main title to the plot. Default is \code{NULL}.
#' @param bw The smoothing bandwidth to be used. Can be a numeric value (e.g., \code{0.06}) 
#'   or a character string like \code{"nrd"} (the default automatic method).
#'
#' @return A \code{\link[ggplot2]{ggplot}} object visualizing the source contribution distributions.
#'
#' @examples
#' # Generate mock unmixing output
#' set.seed(123)
#' mock_unmix <- data.frame(
#'   id = 1:100,
#'   gof = runif(100, 85, 95),
#'   SourceA = rnorm(100, mean = 0.45, sd = 0.05),
#'   SourceB = rnorm(100, mean = 0.35, sd = 0.04),
#'   SourceC = rnorm(100, mean = 0.20, sd = 0.03)
#' )
#'
#' # Density Plot
#' plot_apportionment(mock_unmix, title = "Source Contributions (Density)")
#'
#' # Violin Plot with true values
#' true_vals <- c("SourceA" = 0.45, "SourceB" = 0.35, "SourceC" = 0.20)
#' plot_apportionment(
#'   mock_unmix, violin = TRUE, 
#'   vlines = true_vals, 
#'   title = "Source Contributions (Violin)"
#' )
#' 
#' @export
plot_apportionment <- function(data, violin = FALSE, bounds = c(0.0, 1.0), 
                               scaled = TRUE, y_high = 1, colors = NULL, 
                               vlines = NULL, title = NULL, bw = "nrd") {
  
  # Reshape data using base R stats::reshape (assuming Col 1 is ID and Col 2 is GOF)
  source_cols <- colnames(data)[-c(1, 2)]
  data_plots  <- stats::reshape(
    data,
    varying = source_cols,
    v.names = "value",
    timevar = "variable",
    times = source_cols,
    direction = "long"
  )
  data_plots$variable <- factor(data_plots$variable, levels = source_cols)
  
  if (!violin) {
    if (scaled) {
      p <- ggplot2::ggplot(data_plots, ggplot2::aes(x = .data[["value"]])) +
        ggplot2::geom_density(
          ggplot2::aes(
            group = .data[["variable"]],
            colour = .data[["variable"]],
            fill = .data[["variable"]],
            y = ggplot2::after_stat(scaled)
          ),
          n = 2000L, bw = bw, alpha = 0.35
        )
    } else {
      p <- ggplot2::ggplot(data_plots, ggplot2::aes(x = .data[["value"]])) +
        ggplot2::geom_density(
          ggplot2::aes(
            group = .data[["variable"]],
            colour = .data[["variable"]],
            fill = .data[["variable"]]
          ),
          n = 2000L, bw = bw, alpha = 0.35
        ) +
        ggplot2::ylim(0, y_high)
    }
    
    p <- p + ggplot2::coord_cartesian(xlim = bounds, expand = TRUE) +
      ggplot2::labs(
        x = "Proportional Contribution",
        y = if (scaled) "Scaled Posterior Density" else "Density"
      )
    
    # Add vertical reference lines if provided
    if (!is.null(vlines)) {
      for (varname in names(vlines)) {
        val <- vlines[[varname]]
        col <- if (!is.null(colors) && varname %in% names(colors)) colors[[varname]] else "black"
        
        if (!is.na(val)) {
          p <- p + ggplot2::geom_vline(xintercept = val, color = col, linetype = "dashed", linewidth = 0.8)
        }
      }
    }
    
  } else {
    p <- ggplot2::ggplot(
      data_plots,
      ggplot2::aes(x = .data[["variable"]], y = .data[["value"]], colour = .data[["variable"]])
    ) +
      ggplot2::geom_violin(bw = bw, trim = FALSE, alpha = 0.65) +
      ggplot2::geom_boxplot(
        width = 0.2, alpha = 0.65, colour = "black",
        ggplot2::aes(fill = .data[["variable"]])
      ) +
      ggplot2::stat_summary(
        fun = mean, geom = "point", size = 2.5,
        color = "red", shape = 8
      ) +
      ggplot2::labs(x = "Source", y = "Apportionment") +
      ggplot2::coord_cartesian(ylim = bounds, expand = TRUE)
  }
  
  # Apply custom colors if specified
  if (!is.null(colors)) {
    p <- p +
      ggplot2::scale_fill_manual(values  = colors) +
      ggplot2::scale_color_manual(values = colors)
  }
  
  # Add Title
  if (!is.null(title)) {
    p <- p + ggplot2::ggtitle(title) + ggplot2::theme(plot.title = ggplot2::element_text(hjust = 0.5))
  }
  
  return(p)
}

# --- Backwards Compatibility Aliases ---
#' @rdname plot_apportionment
#' @export
plotResults <- plot_apportionment

#' @rdname plot_apportionment
#' @export
plot_results <- plot_apportionment