R/ggbiplot.R

In OceanTracer: Marine Source Unmixing Toolkit

#' Biplot for Principal Components using ggplot2
#'
#' Renders a custom publication-ready biplot for PCA models created via
#' \code{prcomp()}, \code{princomp()}, \code{FactoMineR::PCA()}, or \code{MASS::lda()}.
#'
#' @param pcobj Object. A PCA object returned by \code{\link[stats]{prcomp}},
#'   \code{\link[stats]{princomp}}, or \code{PCA()}.
#' @param choices Numeric vector of length 2. Indices of principal components to plot. Default \code{c(1, 2)}.
#' @param scale Numeric. Scale factor (0 or 1). Covariance biplot (scale = 1), form biplot (scale = 0). Default 1.
#' @param pc.biplot Logical. If \code{TRUE}, ensures compatibility with \code{biplot.princomp()} scaling. Default \code{TRUE}.
#' @param obs.scale Numeric. Scale factor applied to observations. Default \code{1 - scale}.
#' @param var.scale Numeric. Scale factor applied to variables. Default \code{scale}.
#' @param groups Factor. Optional grouping variable to color-code observations. Default \code{NULL}.
#' @param ellipse Logical. If \code{TRUE}, draws a normal concentration ellipse for each group. Default \code{FALSE}.
#' @param ellipse.prob Numeric. Probability level for concentration ellipse size. Default \code{0.68}.
#' @param labels Character vector. Optional custom labels for observations. Default \code{NULL}.
#' @param labels.size Numeric. Text size used for observation labels. Default \code{3}.
#' @param alpha Numeric [0, 1]. Transparency for points. Default \code{1}.
#' @param var.axes Logical. If \code{TRUE}, renders variable loading arrows. Default \code{TRUE}.
#' @param circle Logical. If \code{TRUE}, draws a correlation circle. Default \code{FALSE}.
#' @param circle.prob Numeric. Size of the correlation circle in Normal probability. Default \code{0.69}.
#' @param varname.size Numeric. Text size for variable names. Default \code{3}.
#' @param varname.adjust Numeric. Adjustment factor for variable name placement (\eqn{\ge 1} moves text away from arrow). Default \code{1.5}.
#' @param varname.abbrev Logical. If \code{TRUE}, abbreviates variable names. Default \code{FALSE}.
#' @param ... Additional arguments passed to \code{ggplot2}.
#'
#' @return A \code{\link[ggplot2]{ggplot}} object.
#'
#' @import ggplot2
#' @importFrom stats qchisq var
#' @export
ggbiplot <- function(pcobj,
                     choices = 1:2,
                     scale = 1,
                     pc.biplot = TRUE,
                     obs.scale = 1 - scale,
                     var.scale = scale,
                     groups = NULL,
                     ellipse = FALSE,
                     ellipse.prob = 0.68,
                     labels = NULL,
                     labels.size = 3,
                     alpha = 1,
                     var.axes = TRUE,
                     circle = FALSE,
                     circle.prob = 0.69,
                     varname.size = 3,
                     varname.adjust = 1.5,
                     varname.abbrev = FALSE,
                     ...) {
  
  stopifnot(length(choices) == 2L)
  
  # ── 1. Recover SVD Components ──────────────────────────────────────────────
  if (inherits(pcobj, "prcomp")) {
    nobs.factor <- sqrt(nrow(pcobj$x) - 1L)
    d <- pcobj$sdev
    u <- sweep(pcobj$x, 2L, 1 / (d * nobs.factor), FUN = "*")
    v <- pcobj$rotation
  } else if (inherits(pcobj, "princomp")) {
    nobs.factor <- sqrt(pcobj$n.obs)
    d <- pcobj$sdev
    u <- sweep(pcobj$scores, 2L, 1 / (d * nobs.factor), FUN = "*")
    v <- pcobj$loadings
  } else if (inherits(pcobj, "PCA")) {
    nobs.factor <- sqrt(nrow(pcobj$call$X))
    d <- unlist(sqrt(pcobj$eig)[1L])
    u <- sweep(pcobj$ind$coord, 2L, 1 / (d * nobs.factor), FUN = "*")
    v <- sweep(pcobj$var$coord, 2L, sqrt(pcobj$eig[seq_len(ncol(pcobj$var$coord)), 1L]), FUN = "/")
  } else if (inherits(pcobj, "lda")) {
    nobs.factor <- sqrt(pcobj$N)
    d <- pcobj$svd
    u <- predict(pcobj)$x / nobs.factor
    v <- pcobj$scaling
  } else {
    stop("Expected an object of class 'prcomp', 'princomp', 'PCA', or 'lda'.")
  }
  
  # ── 2. Scale Scores & Directions ───────────────────────────────────────────
  choices <- pmin(choices, ncol(u))
  df.u <- as.data.frame(sweep(u[, choices, drop = FALSE], 2L, d[choices]^obs.scale, FUN = "*"))
  
  v <- sweep(v, 2L, d^var.scale, FUN = "*")
  df.v <- as.data.frame(v[, choices, drop = FALSE])
  
  colnames(df.u) <- c("xvar", "yvar")
  colnames(df.v) <- c("xvar", "yvar")
  
  if (pc.biplot) {
    df.u <- df.u * nobs.factor
  }
  
  # Scale radius of correlation circle
  r <- sqrt(stats::qchisq(circle.prob, df = 2L)) * prod(colMeans(df.u^2))^(1/4)
  
  # Scale directions
  v.scale <- rowSums(v^2)
  df.v <- r * df.v / sqrt(max(v.scale))
  
  # ── 3. Axis Label Preparation ──────────────────────────────────────────────
  if (obs.scale == 0) {
    u.axis.labs <- paste0("standardized PC", choices)
  } else {
    u.axis.labs <- paste0("PC", choices)
  }
  
  if (!is.null(pcobj$sdev)) {
    var_exp <- 100 * (pcobj$sdev[choices]^2) / sum(pcobj$sdev^2)
    u.axis.labs <- paste0(u.axis.labs, sprintf(" (%0.1f%% explained var.)", var_exp))
  }
  
  # ── 4. Metadata Mapping ───────────────────────────────────────────────────
  if (!is.null(labels)) df.u$labels <- labels
  if (!is.null(groups)) df.u$groups <- as.factor(groups)
  
  if (varname.abbrev) {
    df.v$varname <- abbreviate(rownames(v))
  } else {
    df.v$varname <- rownames(v)
  }
  
  df.v$angle <- (180 / pi) * atan(df.v$yvar / df.v$xvar)
  df.v$hjust <- (1 - varname.adjust * sign(df.v$xvar)) / 2
  
  # ── 5. Build ggplot Base ───────────────────────────────────────────────────
  g <- ggplot2::ggplot(data = df.u, ggplot2::aes(x = .data$xvar, y = .data$yvar)) +
    ggplot2::xlab(u.axis.labs[1L]) +
    ggplot2::ylab(u.axis.labs[2L]) +
    ggplot2::coord_equal()
  
  # Correlation Circle & Vectors
  if (var.axes) {
    if (circle) {
      theta  <- c(seq(-pi, pi, length.out = 50L), seq(pi, -pi, length.out = 50L))
      circ_df <- data.frame(xvar = r * cos(theta), yvar = r * sin(theta))
      g <- g + ggplot2::geom_path(
        data = circ_df,
        color = "gray80",
        linewidth = 0.5,
        alpha = 0.5
      )
    }
    
    g <- g + ggplot2::geom_segment(
      data = df.v,
      ggplot2::aes(x = 0, y = 0, xend = .data$xvar, yend = .data$yvar),
      arrow = ggplot2::arrow(length = ggplot2::unit(1/2, "picas")),
      color = "red3"
    )
  }
  
  # Points / Labels
  if (!is.null(df.u$labels)) {
    if (!is.null(df.u$groups)) {
      g <- g + ggplot2::geom_text(
        ggplot2::aes(label = .data$labels, color = .data$groups),
        size = labels.size
      )
    } else {
      g <- g + ggplot2::geom_text(
        ggplot2::aes(label = .data$labels),
        size = labels.size
      )
    }
  } else {
    if (!is.null(df.u$groups)) {
      g <- g + ggplot2::geom_point(
        ggplot2::aes(color = .data$groups),
        alpha = alpha
      )
    } else {
      g <- g + ggplot2::geom_point(alpha = alpha)
    }
  }
  
  # Concentration Ellipses
  if (!is.null(df.u$groups) && ellipse) {
    theta <- c(seq(-pi, pi, length.out = 50L), seq(pi, -pi, length.out = 50L))
    circle_pts <- cbind(cos(theta), sin(theta))
    
    ell_list <- by(df.u, df.u$groups, function(sub_df) {
      if (nrow(sub_df) <= 2L) return(NULL)
      
      sigma <- stats::var(cbind(sub_df$xvar, sub_df$yvar))
      mu    <- c(mean(sub_df$xvar), mean(sub_df$yvar))
      ed    <- sqrt(stats::qchisq(ellipse.prob, df = 2L))
      
      e_df  <- as.data.frame(sweep(circle_pts %*% chol(sigma) * ed, 2L, mu, FUN = "+"))
      colnames(e_df) <- c("xvar", "yvar")
      e_df$groups    <- sub_df$groups[1L]
      return(e_df)
    })
    
    ell <- do.call(rbind, ell_list)
    if (!is.null(ell)) {
      g <- g + ggplot2::geom_path(
        data = ell,
        ggplot2::aes(x = .data$xvar, y = .data$yvar, color = .data$groups, group = .data$groups)
      )
    }
  }
  
  # Variable Text Labels
  if (var.axes) {
    g <- g + ggplot2::geom_text(
      data = df.v,
      ggplot2::aes(
        label = .data$varname,
        x = .data$xvar,
        y = .data$yvar,
        angle = .data$angle,
        hjust = .data$hjust
      ),
      color = "darkred",
      size = varname.size
    )
  }
  
  return(g)
}