R/unmix_mixsiar.R

In OceanTracer: Marine Source Unmixing Toolkit

#' Bayesian Source Apportionment using MixSIAR
#'
#' A wrapper function that seamlessly formats an \code{OceanTracer} dataset (raw or averaged) 
#' and passes it to the \code{MixSIAR} Bayesian unmixing engine.
#'
#' @param data Data frame. Contains source and mixture data in standard OceanTracer format.
#' @param tracers Character vector. Names of the tracers to include in the model. Default is \code{NULL} (auto-detected).
#' @param prior_type Character. \code{"uninformative"} (default, alpha=1) or \code{"informative"}.
#' @param prior_values Numeric vector. Required if \code{prior_type = "informative"}. Must match the number of sources.
#' @param run Character. MCMC run length. Options: \code{"test"}, \code{"very short"}, \code{"short"}, \code{"normal"} (default), \code{"long"}, \code{"very long"}, \code{"extreme"}.
#' @param process_err Logical. Include process error in the JAGS model? Default \code{FALSE}.
#' @param resid_err Logical. Include residual error in the JAGS model? Default \code{TRUE}.
#'
#' @return A named list containing:
#'   \item{apportionment}{A data frame with posterior draw proportions for each source, formatted for \code{\link{plot_apportionment}}.}
#'   \item{bayesian_summary}{A data frame summarizing MCMC parameter posteriors (mean, SD, quantiles, Rhat).}
#'   \item{jags_model}{The raw \code{rjags} output object returned by MixSIAR.}
#'
#' @examples
#' \donttest{
#' # Requires the external JAGS C++ library and MixSIAR R package installed
#' if (requireNamespace("MixSIAR", quietly = TRUE)) {
#'   
#'   # Sample dataset in OceanTracer raw format
#'   raw_data <- data.frame(
#'     ID = paste0("Sample_", 1:9),
#'     Group = rep(c("Source_A", "Source_B", "Mix_1"), each = 3),
#'     d13C = c(-24, -25, -24.5, -12, -11.5, -12.2, -18, -18.2, -17.9),
#'     d15N = c(5, 5.5, 5.2, 12, 11.8, 12.1, 8.5, 8.7, 8.4)
#'   )
#'   
#'   # Run MixSIAR test chain (1,000 iterations)
#'   try({
#'     res <- unmix_mixsiar(
#'       data = raw_data,
#'       tracers = c("d13C", "d15N"),
#'       run = "test"
#'     )
#'     
#'     # Plot posterior proportions
#'     plot_apportionment(res$apportionment, title = "MixSIAR Apportionment")
#'   }, silent = TRUE)
#' }
#' }
#'
#' @export
unmix_mixsiar <- function(data, 
                          tracers = NULL, 
                          prior_type = "uninformative", 
                          prior_values = NULL, 
                          run = "normal", 
                          process_err = FALSE, 
                          resid_err = TRUE) {
  
  # Ensure required packages are installed
  if (!requireNamespace("MixSIAR", quietly = TRUE)) {
    stop("The 'MixSIAR' package is required. Please install it using install.packages('MixSIAR').", call. = FALSE)
  }
  
  # ── Auto-detect tracers if not provided ───────────────────────────────────
  if (is.null(tracers)) {
    is_avg <- any(grepl("^[Mm]ean_", names(data)))
    if (is_avg) {
      mean_cols <- grep("^[Mm]ean_", names(data), value = TRUE)
      tracers <- gsub("^[Mm]ean_", "", mean_cols)
    } else {
      tracers <- names(data)[3:ncol(data)]
    }
  }
  
  group_col <- names(data)[2]
  is_avg <- any(grepl("^[Mm]ean_", names(data)))
  
  if (!is_avg) {
    # ── RAW DATA: Calculate Means, SDs, and n using Base R ─────────────────
    data[[group_col]] <- factor(data[[group_col]], levels = unique(data[[group_col]]))
    mix_label <- levels(data[[group_col]])[nlevels(data[[group_col]])]
    
    sources_raw <- data[data[[group_col]] != mix_label, , drop = FALSE]
    mix_df <- data[data[[group_col]] == mix_label, c(group_col, tracers), drop = FALSE]
    
    # Base R Aggregation for Means and SDs
    means_mat <- stats::aggregate(sources_raw[, tracers, drop = FALSE], 
                                  by = list(sources_raw[[group_col]]), 
                                  FUN = mean, na.rm = TRUE)
    sds_mat   <- stats::aggregate(sources_raw[, tracers, drop = FALSE], 
                                  by = list(sources_raw[[group_col]]), 
                                  FUN = sd, na.rm = TRUE)
    counts    <- as.data.frame(table(sources_raw[[group_col]]))
    
    source_df <- data.frame(Group = means_mat[[1]])
    for (t in tracers) {
      source_df[[paste0("Mean", t)]] <- means_mat[[t]]
      source_df[[paste0("SD", t)]]   <- sds_mat[[t]]
    }
    source_df$n <- counts$Freq[match(source_df$Group, counts$Var1)]
    colnames(source_df)[1] <- group_col
    source_names <- as.character(source_df[[1]])
    
  } else {
    # ── AVERAGED DATA: Standardize Column Names ───────────────────────────
    mean_cols <- paste0("Mean_", tracers)
    sd_cols   <- paste0("SD_", tracers)
    
    if (!any(mean_cols %in% names(data))) {
      mean_cols <- paste0("mean_", tracers)
      sd_cols   <- paste0("sd_", tracers)
    }
    
    data_unmix <- data[, c("ID", group_col, mean_cols, sd_cols, "n")]
    data_unmix[[group_col]] <- factor(data_unmix[[group_col]], levels = unique(data_unmix[[group_col]]))
    mix_label <- levels(data_unmix[[group_col]])[nlevels(data_unmix[[group_col]])]
    
    source_df <- data_unmix[data_unmix[[group_col]] != mix_label, -1, drop = FALSE]
    mix_df    <- data_unmix[data_unmix[[group_col]] == mix_label, -1, drop = FALSE]
    
    names(source_df)[names(source_df) %in% mean_cols] <- paste0("Mean", tracers)
    names(source_df)[names(source_df) %in% sd_cols]   <- paste0("SD", tracers)
    source_names <- as.character(source_df[[1]])
    
    mix_df <- mix_df[, !names(mix_df) %in% c(sd_cols, "n"), drop = FALSE]
    names(mix_df)[names(mix_df) %in% mean_cols] <- tracers
  }
  
  # ── Build TEF (Discrimination) Table ──────────────────────────────────────
  TEF_df <- source_df[, c(names(source_df)[1], paste0("Mean", tracers))]
  for (t in tracers) TEF_df[[paste0("Mean", t)]] <- 0
  for (t in tracers) TEF_df[[paste0("SD", t)]]   <- 0
  
  # ── Temporary CRAN-Safe Files ─────────────────────────────────────────────
  temp_dir   <- tempdir()
  src_file   <- file.path(temp_dir, "temp_source.csv")
  mix_file   <- file.path(temp_dir, "temp_mix.csv")
  tef_file   <- file.path(temp_dir, "temp_tef.csv")
  model_file <- file.path(temp_dir, "temp_model.txt")
  
  utils::write.table(source_df, src_file, sep = ",", row.names = FALSE)
  utils::write.table(mix_df, mix_file, sep = ",", row.names = FALSE)
  utils::write.table(TEF_df, tef_file, sep = ",", row.names = FALSE)
  
  on.exit(unlink(c(src_file, mix_file, tef_file, model_file)), add = TRUE)
  
  # ── Load into MixSIAR ─────────────────────────────────────────────────────
  message("Formatting data for MixSIAR...")
  suppressMessages({
    mix_obj <- MixSIAR::load_mix_data(
      filename = mix_file, 
      iso_names = tracers, 
      factors = NULL, fac_random = NULL, fac_nested = NULL, cont_effects = NULL
    )
    
    source_obj <- MixSIAR::load_source_data(
      filename = src_file, 
      source_factors = NULL, conc_dep = FALSE, data_type = "means", mix_obj
    )
    
    discr_obj <- MixSIAR::load_discr_data(filename = tef_file, mix_obj)
  })
  
  # ── Priors ────────────────────────────────────────────────────────────────
  if (prior_type == "uninformative") {
    alpha <- rep(1, source_obj$n.sources)
  } else if (prior_type == "informative") {
    if (is.null(prior_values)) stop("Provide 'prior_values' for informative prior.", call. = FALSE)
    alpha <- prior_values
  }
  alpha <- alpha * length(alpha) / sum(alpha)
  alpha[alpha == 0] <- 0.01
  
  # ── Run JAGS ──────────────────────────────────────────────────────────────
  message(paste0("Running MixSIAR MCMC chain (Length: ", run, ")... This may take a while."))
  MixSIAR::write_JAGS_model(model_file, process_err, resid_err, mix_obj, source_obj)
  
  jags_out <- MixSIAR::run_model(
    run = run, mix_obj, source_obj, discr_obj, model_file, 
    alpha.prior = alpha, process_err, resid_err
  )
  
  # ── Extract Proportions ───────────────────────────────────────────────────
  p_global <- as.data.frame(jags_out$BUGSoutput$sims.list$p.global)
  colnames(p_global) <- source_names
  
  mix_id <- as.character(mix_df[[group_col]][1])
  
  apportionment_df <- data.frame(
    ID = mix_id,
    GOF = NA
  )
  apportionment_df <- cbind(apportionment_df, p_global)
  
  message("MixSIAR run complete.")
  
  return(list(
    apportionment = apportionment_df,
    bayesian_summary = as.data.frame(jags_out$BUGSoutput$summary),
    jags_model = jags_out
  ))
}