R/plot_MCMC.R

In BayLum: Chronological Bayesian Models Integrating Optically Stimulated Luminescence and Radiocarbon Age Dating

#' @title Plot MCMC trajectories and posterior distributions
#'
#' @description
#' This function uses the output of [rjags::jags.model] to visualise the traces of the MCMC and the
#' corresponding densities. In particular it displays the posterior distributions of the age, if it is calculated,
#' palaeodose and the equivalent dose dispersion parameters of the sample. The function output is very
#' similar to plot output produced with the 'coda' package, but tailored to meet the needs in
#' the context of the `'BayLum'` package.
#'
#' @details
#' The function is used in the function [Age_Computation], [AgeS_Computation]
#' and [Palaeodose_Computation], but can be used also as standalone plot function.
#'
#' @param object [coda::mcmc.list] or [coda::mcmc] (**required**): Output generated by [rjags::jags.model],
#' e.g., in [Age_Computation]. Alternatively, limited support is provided for `BayLum.list` object input
#'
#' @param sample_names [character] (optional): Names of the used samples. This argument overrides the optional
#' argument `mtext`.
#'
#' @param variables [character] (with default): Variables in your [coda::mcmc] object to be plotted.
#'
#' @param axes_labels [character] (with default): Axes labels used for the trace and density plots. The labels should
#' be provided as named [character] [vector] with the parameter names as the names used to assign the axes labelling.
#' The labelling for the x-axis (trace plots) and y-axis (density plot) cannot be modified.
#'
#' @param plot_single [logical] (with default): Enables/disables the single plot mode of the function, i.e.
#' if set to `TRUE` every plot is returned in a single plot and own [par] settings can be applied.
#'
#' @param n.chains [numeric] (optional): Set the number of chains to visualise, if nothing is provided the
#' number of chains is determined from the input object
#'
#' @param n.iter [integer] (with default): Set the number of iterations to be visualised in the trace plots, regardless
#' of the size of the input dataset as long as the real number of iterations is > `n.iter`.
#' Please note that large numbers impact the plot performance.
#'
#' @param smooth [logical] (with default): Enable/disables smooth of trace plots using [stats::smooth]
#'
#' @param rug [logical] (with default): Enable/disables [rug] under density plots
#'
#' @param ... further arguments that can be passed to modify the plot output. Supported arguments are
#' `lwd`, `lty`, `col`, `type`, `cex`,`mtext`, cf. [mtext] for `mtext` and [plot.default] for the other
#' arguments.
#'
#'
#' @return
#' Two plots: Traces of the MCMC chains and the corresponding density plots. This plots
#' are similar to [coda::traceplot] and [coda::densplot].
#'
#' @section Function version: 0.1.5
#'
#' @keywords dplot
#'
#' @author Sebastian Kreutzer, Institute of Geography, Ruprecht-Karl University of Heidelberg (Germany). This function
#' is a re-written version of the function 'MCMC_plot()' by Claire Christophe
#'
#' @seealso [Age_Computation], [AgeS_Computation], [Palaeodose_Computation],
#' [rjags::coda.samples] and [rjags] packages.
#'
#' @examples
#' data(MCMCsample,envir = environment())
#' object <- coda::as.mcmc(MCMCsample)
#' plot_MCMC(object)
#'
#' @md
#' @export
plot_MCMC <- function(
  object,
  sample_names = NULL,
  variables = c("A", "D", "sD"),
  axes_labels = c("A" = "Age (ka)", "D" = "D (Gy)", "sD" = "sD (Gy)"),
  n.chains = NULL,
  n.iter = 1000,
  smooth = FALSE,
  rug = TRUE,
  plot_single = FALSE,
  ...
){

  # Integrity tests ---------------------------------------------------------------------------
  if(!inherits(object, "mcmc.list")){
    if(inherits(object, "mcmc")){
      object <- coda::as.mcmc.list(object)

    }else if(inherits(object, "BayLum.list")){
      if(!is.null(attributes(object)$originator)){
        ##select what to do for different functions
        switch(attributes(object)$originator,
               AgeS_Computation = object <- object$Sampling,
               Age_OSLC14 = object <- object$Sampling
        )

      }
    }
  }

  if(!inherits(object, "mcmc.list"))
    stop("[plot_MCMC()] 'sample' has to be of class 'mcmc.list' or 'mcmc'!", call. = FALSE)

  # Extract wanted parameters -------------------------------------------------------------------
  if(!all(gsub(coda::varnames(object), pattern = "\\[.+\\]" ,replacement = "") %in% variables)){
    sel <- which(gsub(coda::varnames(object), pattern = "\\[.+\\]" ,replacement = "") %in% variables)

    if(length(sel) == 0){
      sel <- unique(gsub(coda::varnames(object), pattern = "\\[.\\]", replacement = ""))
      warning(paste0("[plot_MCMC()] Invalid 'variables'. Set to found variables from your dataset: ",
                  paste(sel, collapse = ", "), "."), call. = FALSE)
    }

    ##create new object
    object <- as.mcmc.list(lapply(object, function(x){
      x[,sel, drop = FALSE]

    }))

  }

  # Plot preparation ----------------------------------------------------------------------------
  ##set number of chains ... why we are doing it here and not in the argument preset?
  ##We first have to evalulate the object to make sure that it is of class 'mcmc.list'
  if(is.null(n.chains))
    n.chains <- coda::nchain(object)

  ##treat the argument n.iter, which is rather fragile and easly misused
  if(length(n.iter) > 1)
    warning("[plot_MCMC()] length 'n.iter' > 1, using only the first value", call. = FALSE, immediate. = TRUE)

  if(n.iter[1] > coda::niter(object) || n.iter[1] < 2){
    warning("[plot_MCMC()] 'n.iter' out of range, reset to number of observations", call. = FALSE, immediate. = TRUE)
    n.iter <- coda::niter(object)
  }

  #now set n.iter object
  n.iter <- floor(seq(1,coda::niter(object), length.out = abs(n.iter)))

  ##set plot defaults
  plot_settings <- list(
    mtext = "",
    cex = 1,
    type = "l",
    col = 1:n.chains,
    lty = 1:n.chains,
    lwd = 1

  )

  ##modify default on request
  plot_settings <- modifyList(x = plot_settings, val = list(...))

  ##prepare the plot object, this needs to be done manually, due to special requests[]
    ##(1) trace plots
    ##extract all chain for each variable and combine them in list of matrices
    traces_list <- lapply(coda::varnames(object), function(v){
      vapply(1:n.chains, function(x){
        if(smooth){
          smooth(as.numeric((object[[x]][,v]))[n.iter])

        }else{
          as.numeric((object[[x]][,v]))[n.iter]

        }

      }, numeric(length(n.iter)))

    })

    ##(1.1) grep information in mcpar, they should be the same. However, if not we can check here
    ##later
    mcpar_list <- lapply(object, function(x){
      attributes(x)$mcpar

    })

    ##re-assign var names to list
    names(traces_list) <- coda::varnames(object)

    ##create a list of ylab for the traces using a lookup table
    ##if something is not yet in our lookup table it will produce NA
    ylab_traces <- as.character(axes_labels[gsub(coda::varnames(object), pattern = "\\[.+\\]" ,replacement = "")])

    ##(2) density plots
    density_list <- lapply(traces_list, function(v){
      temp <- lapply(1:ncol(v), function(x){
        density(v[,x])

      })

    })

    ##the trace plot ylab becomes the density xlab
    xlab_density <- ylab_traces

  # Plotting ------------------------------------------------------------------------------------
  ##extract real variable names
  ##in our case, e.g, A[1] and A[2] ... two samples
  ##or just A, D, sD for one sample
  if(any(grepl(coda::varnames(object), pattern = "[", fixed = TRUE))){
    sample_info <- unlist(regmatches(coda::varnames(object), gregexpr('\\(?[0-9,.]+', coda::varnames(object))))

    ##check whether the sample_names are correct in length
    if(!is.null(sample_names)){
        if(length(sample_names) > max(as.numeric(sample_info))){
          warning("[plot_MCMC()] 'sample_names' have more entries than your data, take only the first two names!",
                  call. = FALSE, immediate. = TRUE)
          sample_names <- sample_names[1:2]
        }
    }

  }else{
    sample_info <- coda::varnames(object)
    sample_names <- sample_names[1]

  }


  ##order output according to the sample names (here number, e.g., A[1], A[2], ...)
  if(suppressWarnings(!any(is.na(as.numeric(sample_info))))){
    o <- order(as.numeric(sample_info))

  }else{
    o <- order(sample_info)

  }


  ##expand sample_names the right length
  if(!is.null(sample_names)){
  sample_names <- rep(
      sample_names,
      length.out = length(sample_info) * 2)

  }

  ##make sure that we reset the plotting parameters on exit
  par.default <- par()$mfrow
  on.exit(par(mfrow = par.default))

  ##set mfrow if no single plot is wanted
  if(!plot_single)
    par(mfrow = c(length(unique(gsub(coda::varnames(object), pattern = "\\[.+\\]" ,replacement = ""))), 2))

  ##plot everything in a loop
  for(v in o){
    ##traces
    matplot(
      x = seq(mcpar_list[[1]][1], mcpar_list[[1]][2], length.out = length(n.iter)),
      y = traces_list[[v]],
      type = plot_settings$type,
      lty = plot_settings$lty,
      lwd = plot_settings$lwd,
      col = plot_settings$col,
      xlab = "Iterations",
      ylab = ylab_traces[v],
      main = coda::varnames(object)[v],
      sub = paste0("(orig. thin. = ",mcpar_list[[1]][3] ," | iter. shown = ",length(n.iter),")"),
      cex = plot_settings$cex
    )

    if(is.null(sample_names)){
      mtext(side = 3, plot_settings$mtext, cex = plot_settings$cex * 0.8)
    }else{
      mtext(side = 3, sample_names[v], cex = plot_settings$cex * 0.8)

    }


    ##density plot

    ##set matrix row number
    m_nrow <- length(density_list[[v]][[1]]$x)

    ##extract needed matrices
    mx <- vapply(density_list[[v]], function(mx) {
      mx$x
    }, FUN.VALUE = numeric(m_nrow))

    my <- vapply(density_list[[v]], function(my) {
      my$y
    }, FUN.VALUE = numeric(m_nrow))

    matplot(
      x = mx,
      y = my,
      type = plot_settings$type,
      lty = plot_settings$lty,
      lwd = plot_settings$lwd,
      col = plot_settings$col,
      ylab = "Density",
      xlab = xlab_density[v],
      main = coda::varnames(object)[v],
      cex = plot_settings$cex
    )

    ##add rug
    if(rug)
      rug(as.numeric(traces_list[[v]]), col = rgb(0,0,0,0.3), quiet = TRUE)

    if(is.null(sample_names)){
      mtext(side = 3, plot_settings$mtext, cex = plot_settings$cex * 0.8)
    }else{
      mtext(side = 3, sample_names[v], cex = plot_settings$cex * 0.8)

    }
  }
}