R/unmixing.R

In OceanTracer: Marine Source Unmixing Toolkit

#' Unmix sediment or marine particle mixtures
#'
#' Estimates the fractional contribution of each potential source to one or
#' more mixtures. Users can seamlessly switch between a lightning-fast Frequentist 
#' least-squares approach (with optional Linear Variability Propagation) or a 
#' standard Bayesian MCMC approach.
#'
#' Both raw (individual observations) and averaged (\code{Mean_}/\code{SD_}
#' columns) input formats are accepted. When multiple mixture samples are
#' present, the model processes each independently.
#'
#' @param data Data frame in the standard OceanTracer format.
#' @param model Character. \code{"frequentist"} (default) or \code{"bayesian"}.
#' @param iter Integer. Number of iterations for frequentist Monte Carlo/LVP.
#'   Default \code{1000L}.
#' @param variability Character. Uncertainty metric: \code{"SEM"} (Standard
#'   Error of the Mean, default) or \code{"SD"} (Standard Deviation).
#' @param lvp Logical. \code{TRUE} (default) uses Linear Variability
#'   Propagation; \code{FALSE} uses classical Monte Carlo resampling.
#' @param seed Integer. Random seed for reproducibility. Default \code{123456L}.
#' @param ... Additional arguments passed to the Bayesian engine (e.g., \code{run = "test"}).
#'
#' @return For the \code{"frequentist"} model, a data frame with one row per iteration 
#'   and mixture. For the \code{"bayesian"} model, a named list containing the 
#'   formatted apportionment data frame, the MixSIAR summary statistics, and the 
#'   raw JAGS model object.
#'
#' @references
#' Latorre, B., Lizaga, I., Gaspar, L., & Navas, A. (2025). Evaluating the
#' impact of high source variability and extreme contributing sources on
#' sediment fingerprinting models. \emph{Water Resources Management}.
#' \doi{10.1007/s11269-025-04169-8}
#'
#' @export
unmixing <- function(data, 
                     model = "frequentist", 
                     iter = 1000L, 
                     variability = "SEM",
                     lvp = TRUE, 
                     seed = 123456L, 
                     ...) {
  
  # ── Internal helper: detect averaged format (Mean_/SD_ or mean_/sd_) ──────
  is_averaged_internal <- function(d) {
    cnames <- colnames(d)
    if (sum(grepl("^n$", tolower(cnames))) != 1L) return(FALSE)
    if ((ncol(d) - 3L) %% 2L != 0L) return(FALSE)
    
    ntracer <- (ncol(d) - 3L) / 2L
    nms <- cnames[-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))
  }
  
  # Extract mixture samples
  mix_samples <- inputSample(data)
  mixture_n   <- nrow(mix_samples)
  
  # ── Multiple mixtures: process each individually ──────────────────────────
  if (mixture_n > 1L) {
    if (tolower(model) == "bayesian") {
      message("----------------------------------------------------------------")
      message("WARNING: Multiple mixture samples detected for Bayesian unmixing.")
      message("MixSIAR MCMC chains will run sequentially for each mixture.")
      message("Depending on your 'run' length, this may take a considerable amount of time.")
      message("----------------------------------------------------------------")
    } else {
      warning("Multiple mixture samples detected. Each sample will be processed individually.")
    }
    
    # Extract source rows once to avoid repeated filtering issues
    groups      <- factor(data[, 2L], levels = unique(data[, 2L]))
    mix_label   <- levels(groups)[nlevels(groups)]
    source_data <- data[data[, 2L] != mix_label, , drop = FALSE]
    mix_data    <- data[data[, 2L] == mix_label, , drop = FALSE]
    
    results <- NULL
    
    for (i in seq_len(mixture_n)) {
      # Safely construct dataset with all sources and EXACTLY ONE mixture row
      data_one <- rbind(source_data, mix_data[i, , drop = FALSE])
      
      # Vary seed per iteration to preserve independent sampling across mixtures
      iter_seed <- seed + (i - 1L)
      
      res_i <- unmixing(data = data_one, model = model, iter = iter, 
                        variability = variability, lvp = lvp, seed = iter_seed, ...)
      
      if (tolower(model) == "bayesian") {
        if (is.null(results)) {
          results <- res_i
        } else {
          results$apportionment    <- rbind(results$apportionment, res_i$apportionment)
          results$bayesian_summary <- rbind(results$bayesian_summary, res_i$bayesian_summary)
          results$jags_model       <- res_i$jags_model 
        }
      } else {
        results <- if (is.null(results)) res_i else rbind(results, res_i)
      }
    }
    return(results)
  }
  
  # ── Single mixture: run unmixing ──────────────────────────────────────────
  if (tolower(model) == "frequentist") {
    
    means <- is_averaged_internal(data)
    
    if (!lvp) {
      results <- frequentist_unconstrained(data, variability = variability,
                                           iter = iter, means = means, seed = seed)
    } else {
      results <- frequentist_unconstrained_lvp(data, variability = variability,
                                               iter = iter, means = means, seed = seed)
    }
    
    colnames(results)[1L:2L] <- c("ID", "GOF")
    
    # Correct ID identification directly from data
    groups       <- factor(data[, 2L], levels = unique(data[, 2L]))
    mixture_name <- levels(groups)[nlevels(groups)]
    actual_id    <- as.character(data[data[, 2L] == mixture_name, 1L][1L])
    
    results$ID <- actual_id
    
    return(results)
    
  } else if (tolower(model) == "bayesian") {
    
    message("----------------------------------------------------------------")
    message("Running Bayesian apportionment using the MixSIAR framework.")
    message("Powered by JAGS (Markov Chain Monte Carlo).")
    message("----------------------------------------------------------------")
    
    results <- unmix_mixsiar(data = data, ...)
    return(results)
    
  } else {
    stop("The 'model' argument must be either 'frequentist' or 'bayesian'.")
  }
}