R/Correlation_Chart.R

In OceanTracer: Marine Source Unmixing Toolkit

#' Correlation matrix chart
#'
#' Generates a correlation matrix scatterplot of dataset properties, grouped by source, 
#' to aid in collinearity diagnosis and optimal tracer selection.
#'
#' @param data Data frame. Contains source and mixture data in the standard fingerprinting format.
#' @param tracers Numeric vector. Indices of the tracers to include in the matrix. Default \code{1:(ncol(data) - 2)}.
#' @param mixtures Logical. If TRUE, includes mixture samples in the correlation plots. Default TRUE.
#' @param nmixtures Integer. Number of mixture samples present in the dataset. Default 2.
#' @param colors Character vector. Custom colors for the scatterplot groups. Default NULL.
#' @param seed Integer. Random seed for reproducibility. Default NULL.
#'
#' @return A GGally::ggpairs object.
#'
#' @export
Correlation_Chart <- function(data,
                              tracers   = 1:(ncol(data) - 2),
                              mixtures  = TRUE,
                              nmixtures = 2,
                              colors    = NULL,
                              seed      = NULL) {
  
  if (!requireNamespace("GGally", quietly = TRUE)) {
    stop("Package 'GGally' is required for Correlation_Chart. Please install it.")
  }
  
  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))
  }
  
  raw_dataset_internal <- function(d) {
    ntracer <- (ncol(d) - 3L) / 2L
    n_col   <- which(tolower(colnames(d)) == "n")
    
    d[, 2L] <- factor(d[, 2L], levels = unique(d[, 2L]))
    mixture_label <- levels(d[, 2L])[nlevels(d[, 2L])]
    is_mixture    <- d[, 2L] == mixture_label
    
    out_list    <- list()
    source_rows <- d[!is_mixture, , drop = FALSE]
    
    # Expand source distributions into synthetic sample rows
    for (i in seq_len(nrow(source_rows))) {
      n_samples <- source_rows[i, n_col]
      if (n_samples < 1L) n_samples <- 1L
      
      temp_df <- data.frame(
        ID    = rep(source_rows[i, 1L], n_samples),
        Group = rep(as.character(source_rows[i, 2L]), n_samples),
        stringsAsFactors = FALSE
      )
      
      for (j in seq_len(ntracer)) {
        mean_val <- source_rows[i, 2L + j]
        sd_val   <- source_rows[i, 2L + ntracer + j]
        temp_df[[colnames(source_rows)[2L + j]]] <- if (n_samples > 1L) {
          stats::rnorm(n_samples, mean = mean_val, sd = sd_val)
        } else {
          mean_val
        }
      }
      out_list[[length(out_list) + 1L]] <- temp_df
    }
    
    # Append mixture rows directly
    mix_rows <- d[is_mixture, , drop = FALSE]
    for (i in seq_len(nrow(mix_rows))) {
      temp_df <- data.frame(
        ID    = mix_rows[i, 1L],
        Group = as.character(mix_rows[i, 2L]),
        stringsAsFactors = FALSE
      )
      for (j in seq_len(ntracer)) {
        temp_df[[colnames(mix_rows)[2L + j]]] <- mix_rows[i, 2L + j]
      }
      out_list[[length(out_list) + 1L]] <- temp_df
    }
    
    raw_data <- do.call(rbind, out_list)
    colnames(raw_data) <- gsub("^[Mm]ean_", "", colnames(raw_data))
    raw_data
  }
  
  # Custom correlation display function that handles zero-variance safely
  custom_cor <- function(data, mapping, ...) {
    x_var <- ggplot2::quo_name(mapping$x)
    y_var <- ggplot2::quo_name(mapping$y)
    
    x_val <- data[[x_var]]
    y_val <- data[[y_var]]
    
    ok <- complete.cases(x_val, y_val) & is.finite(x_val) & is.finite(y_val)
    
    if (sum(ok) < 3L || stats::sd(x_val[ok]) == 0 || stats::sd(y_val[ok]) == 0) {
      cor_txt <- "NA"
    } else {
      res <- tryCatch(
        stats::cor.test(x_val[ok], y_val[ok]),
        error = function(e) NULL
      )
      cor_txt <- if (!is.null(res)) {
        sprintf("%.2f", res$estimate)
      } else {
        "NA"
      }
    }
    
    GGally::ggally_text(
      label = paste0("Corr: ", cor_txt),
      mapping = ggplot2::aes(color = NULL),
      xP = 0.5, yP = 0.5,
      ...
    )
  }
  
  # ── Data Pre-processing ─────────────────────────────────────────────────────
  if (is_averaged_internal(data)) {
    data <- raw_dataset_internal(data)
  }
  
  data[, 2L] <- factor(data[, 2L], levels = unique(data[, 2L]))
  data1      <- data
  colnames(data1)[2L] <- "Sources"
  
  if (!mixtures) {
    mix_label <- levels(data1$Sources)[nlevels(data1$Sources)]
    data1     <- data1[data1$Sources != mix_label, , drop = FALSE]
    data1$Sources <- factor(data1$Sources)
  }
  
  target_cols <- tracers + 2L
  
  # ── Chart Construction ──────────────────────────────────────────────────────
  p <- GGally::ggpairs(
    data    = data1,
    columns = target_cols,
    upper   = list(continuous = GGally::wrap(custom_cor)),
    lower   = list(combo = "facetdensity"),
    title   = "Correlation Matrix",
    mapping = ggplot2::aes(color = Sources, fill = Sources, alpha = 0.25),
    legend  = c(1L, 1L)
  ) +
    ggplot2::theme_bw() +
    ggplot2::theme(legend.position = "bottom")
  
  if (!is.null(colors)) {
    plot_colors <- colors
    if (!is.null(names(plot_colors)) && !all(names(plot_colors) %in% levels(data1$Sources))) {
      plot_colors <- unname(plot_colors)
    }
    for (i in seq_len(p$nrow)) {
      for (j in seq_len(p$ncol)) {
        p[i, j] <- p[i, j] +
          ggplot2::scale_color_manual(values = plot_colors) +
          ggplot2::scale_fill_manual(values = plot_colors)
      }
    }
  }
  
  p
}

# ── Aliases ──────────────────────────────────────────────────────────────────
#' @rdname Correlation_Chart
#' @export
correlationPlot <- Correlation_Chart

#' @rdname Correlation_Chart
#' @export
correlation_plot <- Correlation_Chart