R/bd_extract_param.R

In MichaelHoltonPrice/BayDem: Bayesian Demography

#' @title Extract the Bayesian samples for a Gaussian mixture model generated by bd_do_inference
#'
#' @description
#' The input fit is the result of a call to stan by \code{bd_do_inference}, of
#' class stanfit. Return a matrix TH with dimensions S x (3*K), where S is the
#' number of samples (across all chains, and excluding warmup), and K is the
#' number of mixtures. Optionally, a list of samples is returned by calling
#' \code{bd_convert_gauss_mix_param_vect_to_list} repeatedly for each sample.
#' The final column of as.matrix for class stanfit is the log-posterior, which
#' must be removed.
#'
#' @param fit The fit from stan, of class stanfit
#' @param asList (Default False) Whether to return the samples as a matrix or a list
#'
#' @return A matrix or list of samples
#' @export

bd_extract_param <- function(fit, asList = F) {
  if(class(fit) != 'stanfit') {
    stop(paste('Expected fit to be class stanfit, but it is',class(fit)))
  }
  # as.matrix is defined for class stanfit and excludes warmup samples
  TH <- as.matrix(fit)
  # Remove the final column, which is the log-posterior, not a parameter
  TH <- TH[,-ncol(TH)]

  if(!asList) {
    return(TH)
  } else {
    # The total number of (non-warmup) samples
    numSamp <- dim(TH)[1]
    samps <- list()
    for (ss in 1:numSamp) {
      samps[[ss]] <- bd_convert_gauss_mix_param_vect_to_list(TH[ss, ])
    }
    return(samps)
  }
}