R/aa_baseline_correction.R

In OceanTracer: Marine Source Unmixing Toolkit

#' Amino Acid Isotopic Baseline Correction
#'
#' Corrects isotopic ratios by subtracting a source amino acid baseline 
#' (e.g., Phe or Lys) to isolate trophic enrichment. Automatically handles 
#' both raw data and averaged (Mean/SD) data, applying exact quadrature 
#' error propagation for standard deviations.
#'
#' @param data A data frame containing isotopic ratios.
#' @param references Character vector of the amino acid(s) to use as the baseline. 
#'   If multiple are provided (e.g., \code{c("Phe", "Lys")}), their row-wise mean is used.
#' @param invert Character vector of amino acids that undergo inverse trophic 
#'   fractionation (typically \code{"Thr"}). The baseline will be added to these instead of subtracted.
#' @param drop_ref Logical. If \code{TRUE}, the reference AA columns are removed 
#'   from the final dataset. Default is \code{TRUE}.
#'
#' @return A \code{data.frame} of the same class as \code{data} with baseline-corrected 
#'   isotopic values. For raw datasets, ratios are updated in place; for averaged datasets, 
#'   means and standard deviations are updated with propagated quadrature errors. 
#'   If \code{drop_ref = TRUE}, columns corresponding to \code{references} are omitted.
#' 
#' @examples
#' # Example using raw isotopic ratio data
#' raw_df <- data.frame(
#'   R_Phe = c(5.1, 5.3, 5.0),
#'   R_Glu = c(15.2, 16.0, 15.5),
#'   R_Thr = c(-2.1, -2.5, -2.0)
#' )
#' aa_baseline_correction(raw_df, references = "Phe", invert = "Thr")
#'
#' @export
aa_baseline_correction <- function(data, references = "Phe", invert = "Thr", drop_ref = TRUE) {
  
  # Check if data is in averaged format (uses the same logic as conservative_balance)
  nms <- colnames(data)
  is_avg <- (any(grepl("^[Mm]ean_", nms)) && any(grepl("^[Ss][Dd]_", nms))) ||
    (any(grepl("^mean_", nms))    && any(grepl("^sd_", nms)))
  
  # Normalise prefixes if averaged
  if (is_avg) {
    colnames(data) <- sub("^mean_", "Mean_", colnames(data))
    colnames(data) <- sub("^sd_", "SD_", colnames(data))
  }
  
  # ─── 1. RAW DATA FORMAT ─────────────────────────────────────────────────────
  if (!is_avg) {
    ref_cols <- paste0("R_", references)
    if (!all(ref_cols %in% colnames(data))) stop("One or more reference AAs not found. Check names.")
    
    # Calculate baseline
    if (length(references) == 1) {
      baseline <- data[[ref_cols]]
    } else {
      baseline <- rowMeans(data[, ref_cols, drop = FALSE], na.rm = TRUE)
    }
    
    all_r_cols <- grep("^R_", colnames(data), value = TRUE)
    
    for (col in all_r_cols) {
      aa_name <- sub("^R_", "", col)
      if (aa_name %in% references) next 
      
      if (aa_name %in% invert) {
        data[[col]] <- data[[col]] + baseline
      } else {
        data[[col]] <- data[[col]] - baseline
      }
    }
    
    if (drop_ref) {
      cols_to_drop <- c(ref_cols, paste0("Conc_", references))
      data <- data[, !(colnames(data) %in% cols_to_drop), drop = FALSE]
    }
    
    # ─── 2. AVERAGED DATA FORMAT (Mean / SD) ────────────────────────────────────
  } else {
    # Isolate ratio columns (ignoring concentration/content columns)
    mean_cols <- grep("^Mean_", colnames(data), value = TRUE)
    mean_ratios <- mean_cols[!grepl("Conc|cont|conc", mean_cols, ignore.case = TRUE)]
    
    sd_cols <- grep("^SD_", colnames(data), value = TRUE)
    sd_ratios <- sd_cols[!grepl("Conc|cont|conc", sd_cols, ignore.case = TRUE)]
    
    # Identify the specific reference columns
    ref_pattern <- paste0("(", paste(references, collapse = "|"), ")$")
    mean_ref_cols <- grep(ref_pattern, mean_ratios, value = TRUE)
    sd_ref_cols   <- grep(ref_pattern, sd_ratios, value = TRUE)
    
    if (length(mean_ref_cols) != length(references)) stop("Reference AA means not found or ambiguous.")
    
    # Calculate Baseline Mean and SD
    if (length(references) == 1) {
      baseline_mean <- data[[mean_ref_cols]]
      baseline_sd   <- data[[sd_ref_cols]]
    } else {
      baseline_mean <- rowMeans(data[, mean_ref_cols, drop = FALSE], na.rm = TRUE)
      # Error propagation for the mean of multiple items
      sum_variances <- rowSums(data[, sd_ref_cols, drop = FALSE]^2, na.rm = TRUE)
      baseline_sd   <- sqrt(sum_variances) / length(references)
    }
    
    for (mean_col in mean_ratios) {
      # Extract AA name to check against references/inversions
      aa_name <- sub("^Mean_R_", "", mean_col)
      aa_name <- sub("^Mean_", "", aa_name)
      
      if (aa_name %in% references) next
      
      sd_col <- sub("^Mean_", "SD_", mean_col)
      
      # Adjust Means
      if (aa_name %in% invert) {
        data[[mean_col]] <- data[[mean_col]] + baseline_mean
      } else {
        data[[mean_col]] <- data[[mean_col]] - baseline_mean
      }
      
      # Propagate SDs (Quadrature addition applies to both subtraction and addition)
      data[[sd_col]] <- sqrt(data[[sd_col]]^2 + baseline_sd^2)
    }
    
    if (drop_ref) {
      # Drop any column containing the reference names (Means, SDs, and Contents)
      cols_to_drop <- grep(ref_pattern, colnames(data), value = TRUE)
      data <- data[, !(colnames(data) %in% cols_to_drop), drop = FALSE]
    }
  }
  
  message(sprintf("Baseline correction successfully applied using: %s", paste(references, collapse = ", ")))
  if (drop_ref) message("Reference columns were removed from the dataset.")
  
  return(data)
}