R/plot_pca.R

In OceanTracer: Marine Source Unmixing Toolkit

#' Principal Component Analysis biplot
#'
#' Performs a PCA on the tracer data and displays a biplot with source group
#' ellipses. Averaged datasets (with \code{Mean_}/\code{SD_} or 
#' \code{mean_}/\code{sd_} columns) are automatically expanded to individual
#' observations by sampling from a normal distribution before the PCA is
#' computed, so all formats are accepted.
#'
#' @param data A data frame containing source and mixture data in the standard
#'   fingerprinting format (raw or averaged).
#' @param components A numeric vector of length 2 indicating which principal
#'   components to display. Default \code{c(1, 2)}.
#' @param colors An optional character vector of colors, one per group. If
#'   \code{NULL}, ggplot2 defaults are used.
#' @param seed An integer value used to initialize the random number generator.
#'
#' @return A ggplot object (also printed to the active graphics device).
#'
#' @importFrom stats prcomp rnorm
#' @importFrom ggplot2 ggplot aes geom_point geom_hline geom_vline stat_ellipse 
#'   ggtitle theme element_text scale_colour_manual scale_fill_manual guides
#' @export
plot_pca <- function(data, components = c(1L, 2L), colors = NULL, seed = NULL) {
  
  if (!is.null(seed)) {
    set.seed(seed)
  }
  
  # ── Internal Helpers ────────────────────────────────────────────────────────
  is_averaged_internal <- function(d) {
    if (sum(grepl("^n$", tolower(colnames(d)))) != 1L) return(FALSE)
    if ((ncol(d) - 3L) %% 2L != 0L) return(FALSE)
    ntracer <- (ncol(d) - 3L) / 2L
    nms     <- colnames(d)[-c(1L, 2L)][1L:(2L * ntracer)]
    v1      <- nms[grepl("^[Mm]ean_", nms)]
    v2      <- nms[grepl("^[Ss][Dd]_", nms)]
    
    length(v1) == ntracer && length(v2) == ntracer &&
      all(gsub("^[Mm]ean_", "", v1) == gsub("^[Ss][Dd]_", "", v2))
  }
  
  # ── Data Expansion / Formatting ─────────────────────────────────────────────
  if (is_averaged_internal(data)) {
    src      <- inputSource(data)   # Expected layout: ID, Mean_*, SD_*, n
    n_col    <- ncol(src)
    tracer_n <- as.integer((n_col - 2L) / 2L)
    
    rows <- list()
    for (i in seq_len(nrow(src))) {
      n_obs <- as.integer(src[i, n_col])
      means <- as.numeric(src[i, 2L:(1L + tracer_n)])
      sds   <- as.numeric(src[i, (2L + tracer_n):(1L + 2L * tracer_n)])
      mat   <- mapply(function(m, s) stats::rnorm(n_obs, m, s), means, sds)
      
      if (n_obs == 1L) mat <- matrix(mat, nrow = 1L)
      df <- as.data.frame(mat)
      colnames(df) <- gsub("^[Mm]ean_", "", colnames(src)[2L:(1L + tracer_n)])
      df$samples   <- as.character(src[i, 1L])
      rows[[i]]    <- df
    }
    
    # Append mixture row(s)
    mix    <- inputSample(data)
    mix_df <- as.data.frame(mix[, -1L, drop = FALSE])
    colnames(mix_df) <- gsub("^[Mm]ean_", "", colnames(mix_df))
    mix_df$samples   <- as.character(mix[, 1L])
    rows[[length(rows) + 1L]] <- mix_df
    
    raw      <- do.call(rbind, rows)
    data_pca <- raw[, !colnames(raw) %in% "samples", drop = FALSE]
    groups   <- factor(raw$samples, levels = unique(raw$samples))
    
  } else {
    # Raw format processing
    groups   <- factor(data[, 2L], levels = unique(data[, 2L]))
    data_pca <- data[, 3L:ncol(data), drop = FALSE]
  }
  
  # ── PCA Computation ────────────────────────────────────────────────────────
  pca_fit <- stats::prcomp(data_pca, center = TRUE, scale. = TRUE)
  
  pca_scores        <- as.data.frame(pca_fit$x[, components, drop = FALSE])
  pca_scores$groups <- groups
  
  # Handle NA values within score matrix
  if (any(is.na(pca_scores))) {
    message("NA values in PCA scores were replaced with group minimums.")
    pca_scores <- dplyr::group_by(pca_scores, groups)
    pca_scores <- dplyr::mutate(
      pca_scores,
      dplyr::across(
        dplyr::where(is.numeric),
        ~ ifelse(is.na(.), min(., na.rm = TRUE), .)
      )
    )
    pca_scores <- as.data.frame(pca_scores)
  }
  
  # ── Base Biplot Construction ────────────────────────────────────────────────
  p <- ggbiplot::ggbiplot(
    pca_fit,
    choices        = components,
    obs.scale      = 1,
    var.scale      = 1,
    ellipse        = FALSE,
    groups         = pca_scores$groups,
    circle         = FALSE,
    varname.size   = 4,
    varname.adjust = 1.2,
    var.axes       = TRUE
  )
  
  # Filter scaled scores to groups with at least 3 points for stat_ellipse
  scaled_df      <- p$data
  ellipse_groups <- names(which(table(scaled_df$groups) >= 3L))
  ellipse_df     <- scaled_df[scaled_df$groups %in% ellipse_groups, , drop = FALSE]
  
  # ── Add Ellipses & Custom Styling ───────────────────────────────────────────
  p <- p +
    ggplot2::geom_hline(yintercept = 0, colour = "grey50", linetype = "longdash", linewidth = 0.5) +
    ggplot2::geom_vline(xintercept = 0, colour = "grey50", linetype = "longdash", linewidth = 0.5) +
    ggplot2::geom_point(ggplot2::aes(shape = .data$groups, colour = .data$groups), size = 3) +
    ggplot2::stat_ellipse(
      data    = ellipse_df,
      mapping = ggplot2::aes(x = .data$xvar, y = .data$yvar, fill = .data$groups),
      type    = "t",
      geom    = "polygon",
      alpha   = 0.15,
      level   = 0.85,
      linewidth = 0.5
    ) +
    ggplot2::stat_ellipse(
      data    = ellipse_df,
      mapping = ggplot2::aes(x = .data$xvar, y = .data$yvar, colour = .data$groups),
      type    = "t",
      level   = 0.85,
      linewidth = 0.8
    ) +
    ggplot2::ggtitle("Principal Component Analysis") +
    ggplot2::theme_bw() +
    ggplot2::theme(
      legend.direction = "horizontal",
      legend.position  = "top",
      legend.title     = ggplot2::element_blank(),
      text             = ggplot2::element_text(size = 12),
      plot.title       = ggplot2::element_text(hjust = 0.5, face = "bold")
    )
  
  if (!is.null(colors)) {
    num_groups <- length(levels(groups))
    if (length(colors) < num_groups) {
      stop("Not enough colors: need ", num_groups, " but only ", length(colors), " supplied.")
    }
    p <- p +
      ggplot2::scale_colour_manual(values = colors) +
      ggplot2::scale_fill_manual(values = colors)
  }
  
  return(p)
}

# ── Backwards Compatibility Aliases ───────────────────────────────────────────
#' @rdname plot_pca
#' @export
PCAPlot <- plot_pca

#' @rdname plot_pca
#' @export
PCA_plot <- plot_pca