R/LDA_Chart.R

In OceanTracer: Marine Source Unmixing Toolkit

#' Linear Discriminant Analysis Chart
#'
#' Performs a linear discriminant analysis (LDA) to project data into maximum-separation 
#' dimensions. Includes a comprehensive discrimination summary detailing variance explained, 
#' a leave-one-out cross-validation confusion matrix, and misclassification rates to 
#' rigorously assess source separability.
#'
#' @param data Data frame. Contains source and mixture data in standard format.
#' @param Plot3D Logical. If \code{TRUE}, renders a 3D scatter plot instead of a 2D projection. Default is \code{FALSE}.
#' @param text Logical. If \code{TRUE}, overlays the sample identification numbers on the plot. Default is \code{FALSE}.
#' @param colors Character vector. Custom colors for the source groups. Default is \code{NULL}.
#' @param interactive Logical. If \code{TRUE}, renders an interactive HTML widget for 3D plots. Default is \code{FALSE}.
#' @param seed Integer. Random seed for reproducibility. Default is \code{NULL}.
#' @param verbose Logical. Whether to print discrimination summaries and variance explained to the console. Default is \code{TRUE}.
#'
#' @return Invisibly returns a named list containing:
#'    \item{lda}{An object of class \code{\link[MASS]{lda}} representing the fitted linear discriminant model.}
#'    \item{confusion}{A cross-tabulation table of leave-one-out cross-validation predictions versus actual sources.}
#'    \item{prop_var}{A numeric vector indicating the percentage of variance explained by each discriminant axis.}
#'
#' @examples
#' data <- data.frame(
#'    ID = paste0("S", 1:12),
#'    Group = rep(c("SourceA", "SourceB", "SourceC", "Mixture"), each = 3),
#'    Tracer1 = c(10, 12, 11, 20, 22, 21, 30, 31, 32, 20, 21, 22),
#'    Tracer2 = c(5, 6, 5, 15, 14, 16, 25, 26, 24, 15, 15, 16),
#'    Tracer3 = c(1, 2, 1.5, 5, 6, 5.5, 9, 10, 9.5, 5, 5.5, 6)
#' )
#' LDA_Chart(data, seed = 123)
#'
#' @importFrom stats rnorm predict
#' @importFrom graphics plot
#' @importFrom MASS lda
#' @importFrom utils capture.output
#'
#' @export
LDA_Chart <- function(data, Plot3D = FALSE, text = FALSE, colors = NULL, interactive = FALSE, seed = NULL, verbose = TRUE) {
  
  # ── Internal helpers ────────────────────────────────────────────────────────
  
  is_averaged_internal <- function(d) {
    if (sum(grepl("^n$", tolower(colnames(d)))) != 1) return(FALSE)
    if ((ncol(d) - 3) %% 2 != 0) return(FALSE)
    ntracer <- (ncol(d) - 3) / 2
    nms <- colnames(d)[-c(1:2)][1:(2 * 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))
  }
  
  raw_dataset_internal <- function(d) {
    ntracer <- (ncol(d) - 3) / 2
    n_col   <- which(tolower(colnames(d)) == "n")
    d[, 2]  <- factor(d[, 2], levels = unique(d[, 2]))
    mixture_label <- levels(d[, 2])[nlevels(d[, 2])]
    is_mixture    <- d[, 2] == mixture_label
    raw_data <- data.frame()
    source_rows <- d[!is_mixture, ]
    for (i in seq_len(nrow(source_rows))) {
      n_samples <- source_rows[i, n_col]
      if (n_samples < 1) n_samples <- 1
      temp_data <- data.frame(ID    = rep(source_rows[i, 1], n_samples),
                              Group = rep(as.character(source_rows[i, 2]), n_samples))
      for (j in seq_len(ntracer)) {
        mean_val <- source_rows[i, 2 + j]
        sd_val   <- source_rows[i, 2 + ntracer + j]
        temp_data[colnames(source_rows)[2 + j]] <-
          if (n_samples > 1) stats::rnorm(n_samples, mean = mean_val, sd = sd_val)
        else               mean_val
      }
      raw_data <- rbind(raw_data, temp_data)
    }
    mix_rows <- d[is_mixture, ]
    for (i in seq_len(nrow(mix_rows))) {
      temp_data <- data.frame(ID    = mix_rows[i, 1],
                              Group = as.character(mix_rows[i, 2]))
      for (j in seq_len(ntracer)) {
        temp_data[colnames(mix_rows)[2 + j]] <- mix_rows[i, 2 + j]
      }
      raw_data <- rbind(raw_data, temp_data)
    }
    colnames(raw_data) <- gsub("^[Mm]ean_", "", colnames(raw_data))
    raw_data
  }
  
  if (!is.null(seed)) set.seed(seed)
  
  # ── Pre-processing ──────────────────────────────────────────────────────────
  if (is_averaged_internal(data)) {
    data <- raw_dataset_internal(data)
  }
  
  # ── Plotting ────────────────────────────────────────────────────────────────
  data[, 2] <- factor(data[, 2], levels = unique(data[, 2]))
  groups    <- data[, 2]
  mixture   <- levels(groups)[nlevels(groups)]
  sources   <- data[!groups == mixture, ]
  s_groups  <- droplevels(sources[, 2])
  
  data.lda1 <- sources[3:ncol(sources)]
  data.lda1$groups <- as.factor(s_groups)
  
  # Check for NA values and impute with group minimums using Base R
  if (any(is.na(data.lda1))) {
    if (verbose) message("NA values were filled with the minimum value of the same group.")
    data.lda1 <- do.call(rbind, lapply(split(data.lda1, data.lda1$groups), function(sub_df) {
      num_cols <- sapply(sub_df, is.numeric)
      for (col in names(sub_df)[num_cols]) {
        if (any(is.na(sub_df[[col]]))) {
          min_val <- min(sub_df[[col]], na.rm = TRUE)
          sub_df[[col]][is.na(sub_df[[col]])] <- min_val
        }
      }
      sub_df
    }))
  }
  
  data.lda1$groups <- as.factor(s_groups)
  data.lda      <- MASS::lda(groups ~ ., data = data.lda1, tol = 1e-10)
  
  data.lda.pred <- stats::predict(data.lda)
  data.lda.temp <- data.frame(data.lda.pred$x, Sources = data.lda1$groups)
  
  num_groups <- length(unique(data.lda.temp$Sources))
  if (is.null(colors)) {
    default_colors <- grDevices::rainbow(num_groups)
  } else {
    if (length(colors) < num_groups) stop("Insufficient number of colors provided.")
    default_colors <- colors
  }
  
  if (Plot3D) {
    if (interactive) {
      if (!requireNamespace("plotly", quietly = TRUE)) {
        stop("Package 'plotly' is required for interactive 3D plots. Please install it.")
      }
      plot_obj <- plotly::plot_ly(data.lda.temp, x = ~LD1, y = ~LD2, z = ~LD3,
                                  type = "scatter3d", mode = "markers",
                                  color = ~Sources, colors = default_colors,
                                  text = ~Sources, marker = list(size = 8, opacity = 0.7))
      print(plot_obj)
    } else {
      group_num <- as.numeric(data.lda.temp$Sources)
      graphics::plot(data.lda.temp$LD1, data.lda.temp$LD2,
                     col = default_colors[group_num],
                     pch = 19, xlab = "LD1", ylab = "LD2",
                     main = "LDA Projection")
    }
  } else {
    base_plot <- ggplot2::ggplot(
      data.lda.temp,
      ggplot2::aes(x = .data[["LD1"]], y = .data[["LD2"]],
                   colour = .data[["Sources"]], fill = .data[["Sources"]])
    ) +
      ggplot2::geom_point(ggplot2::aes(shape = .data[["Sources"]]), size = 3, alpha = 0.7) +
      ggplot2::geom_hline(yintercept = 0, colour = "black", linetype = "longdash") +
      ggplot2::geom_vline(xintercept = 0, colour = "black", linetype = "longdash") +
      ggplot2::stat_ellipse(type = "t", linewidth = 1, geom = "polygon", alpha = 0.2, level = 0.9) +
      ggplot2::stat_ellipse(type = "t", linewidth = 1, alpha = 0.7, level = 0.9) +
      ggplot2::ggtitle("LDA") +
      ggplot2::theme(plot.title = ggplot2::element_text(hjust = 0.5)) +
      ggplot2::scale_color_manual(values = default_colors) +
      ggplot2::scale_fill_manual(values = default_colors)
    
    if (text) {
      data.lda.temp$row_num <- seq_len(nrow(data.lda.temp))
      plot_obj <- base_plot +
        ggplot2::geom_text(data = data.lda.temp,
                           ggplot2::aes(label = .data[["row_num"]], hjust = 2, vjust = 0.5), size = 4)
    } else {
      plot_obj <- base_plot
    }
    print(plot_obj)
  }
  
  # ── Discrimination summary ──────────────────────────────────────────────────
  
  prop_var <- data.lda$svd^2 / sum(data.lda$svd^2) * 100
  ld_labels <- paste0("LD", seq_along(prop_var))
  
  if (verbose) {
    message("--- Variance explained by each discriminant axis ---")
    for (k in seq_along(prop_var)) {
      message(sprintf("  %-6s  %.1f%%", ld_labels[k], prop_var[k]))
    }
  }
  
  data.lda.cv <- MASS::lda(groups ~ ., data = data.lda1, CV = TRUE, tol = 1e-10)
  conf <- table(Actual = data.lda1$groups, Predicted = data.lda.cv$class)
  
  if (verbose) {
    message("--- Leave-one-out classification accuracy per source ---")
    acc_per_group <- diag(conf) / rowSums(conf) * 100
    for (grp in names(acc_per_group)) {
      n_correct <- diag(conf)[grp]
      n_total   <- rowSums(conf)[grp]
      message(sprintf("  %-20s  %d/%d  (%.0f%%)", grp, n_correct, n_total, acc_per_group[grp]))
    }
    overall <- sum(diag(conf)) / sum(conf) * 100
    message(sprintf("  %-20s  %d/%d  (%.0f%%)", "Overall", sum(diag(conf)), sum(conf), overall))
    
    if (any(conf[row(conf) != col(conf)] > 0)) {
      message("--- Confusion matrix (rows = actual, cols = predicted) ---")
      conf_str <- paste(utils::capture.output(print(conf)), collapse = "\n")
      message(conf_str)
      
      conf_offdiag <- conf
      diag(conf_offdiag) <- 0L
      if (sum(conf_offdiag) > 0) {
        worst_idx <- which(conf_offdiag == max(conf_offdiag), arr.ind = TRUE)[1, ]
        message(sprintf("  Most confused pair: '%s' misclassified as '%s' (%d time%s)",
                        rownames(conf)[worst_idx["Actual"]],
                        colnames(conf)[worst_idx["Predicted"]],
                        max(conf_offdiag),
                        if (max(conf_offdiag) == 1) "" else "s"))
      }
    } else {
      message("  Perfect classification - all sources fully separated.")
    }
  }
  
  invisible(list(lda = data.lda, confusion = conf, prop_var = prop_var))
}

# ── Aliases ──────────────────────────────────────────────────────────────────
#' @rdname LDA_Chart
#' @export
LDA_plot <- LDA_Chart

#' @rdname LDA_Chart
#' @export
LDAPlot <- LDA_Chart