R/converge_diag_param_fn.R

In BayesMoFo: Bayesian Mortality Forecasting

#' A function to assess convergence of the posterior sampling of fitted parameters for monitoring purposes
#'
#' Produce several convergence diagnostic tools (e.g. trace/density/acf plots and effective sample sizes) from the posterior samples of fitted parameters.
#' @param result object of type either "fit_result" or "BayesMoFo".
#' @param plot_params A vector of character strings specifying which set of parameters to plot for visualisation. If not specified, a random selection of the parameters will be included in the plots (see \code{fit_result$param} or \code{runBayesMoFo_result$result$best$param} for a full list) will be chosen. Note that a specific combination of alpha, beta, kappa, and gamma, is to be plotted, then users need to specify the exact indices of them, e.g. \code{plot_params="gamma[1,2]"}. Otherwise, only three randomly selected of them will be plotted. To see a complete list of parameters, i.e. \code{colnames(fit_result$post_sample[[1]])[!startsWith(colnames(fit_result$post_sample[[1]]),"q[")]} or \code{colnames(runBayesMoFo_result$result$best$post_sample[[1]])[!startsWith(colnames(runBayesMoFo_result$result$best$post_sample[[1]]),"q[")]}. 
#' @param trace A logical value to indicate if trace plots of posterior samples of death rates should be shown (default) or suppressed (e.g. to aid visibility).
#' @param density A logical value to indicate if density plots of posterior samples of death rates should be shown (default) or suppressed (e.g. to aid visibility).
#' @param acf_plot A logical value to indicate if auto-correlation plots should be shown or suppressed (default).
#' @param ESS_all A logical value indicating if effective sample sizes are to be computed for all parameters. The default is FALSE where only chosen parameters will be evaluated, if TRUE all parameters will be assessed.
#' @return Some convergence-related plots of posterior samples of fitted parameters.
#' \describe{
#'   \item{\code{ESS}}{The effective sample sizes of the chosen parameters.}
#' }
#' @keywords bayesian diagnostics visualization
#' @concept posterior samples
#' @concept parameters
#' @concept trace plots
#' @concept density plots
#' @concept autocorrelations
#' @concept effective sample sizes
#' @importFrom graphics legend lines par points boxplot
#' @importFrom stats dbinom dpois quantile sd acf
#' @export
#' @examples
#' \donttest{
#' #load and prepare data
#' data("dxt_array_product");data("Ext_array_product")
#' death<-preparedata_fn(dxt_array_product,strat_name = c("ACI","DB","SCI"),ages=35:65)
#' expo<-preparedata_fn(Ext_array_product,strat_name = c("ACI","DB","SCI"),ages=35:65)
#' 
#' #fit any mortality model
#' runBayesMoFo_result<-runBayesMoFo(death=death,expo=expo,models="APCI")
#' 
#' #default plot
#' converge_runBayesMoFo_result<-converge_diag_param_fn(runBayesMoFo_result)
#' 
#' #ESS
#' converge_runBayesMoFo_result$ESS
#' 
#' #plot specific parameters 
#' 
#' runBayesMoFo_result$result$best$param  #run this line to check parameters of the model
#' converge_diag_param_fn(runBayesMoFo_result,plot_params=c("rho","sigma2_kappa","beta")) 
#' #note only three betas were plotted
#' 
#' colnames(runBayesMoFo_result$result$best$post_sample[[1]])[!startsWith(
#'   colnames(runBayesMoFo_result$result$best$post_sample[[1]]),"q[")]  
#' #run the above line to check full list of parameters of the model
#' converge_diag_param_fn(runBayesMoFo_result,plot_params=c("beta[1,2]","gamma[3,2]")) 
#' 
#' #ACF plot 
#' converge_diag_param_fn(runBayesMoFo_result,plot_params=c("beta[1,2]","gamma[3,2]"),
#' trace=FALSE,density=FALSE,acf_plot=TRUE)
#' }

converge_diag_param_fn<-function(result,plot_params=NULL,trace=TRUE,density=TRUE,acf_plot=FALSE,ESS_all=FALSE){
  
  if ("BayesMoFo_obj" %in% names(result)){
    fit_result<-result$result$best
  } else {
    fit_result<-result
  }
  
  mcmc_object<-coda::as.mcmc(as.matrix(fit_result$post_sample))
  mcmc_object<-mcmc_object[,!startsWith(colnames(mcmc_object),"q[")]
  invisible(gc())

  n<-dim(mcmc_object)[1]
  
  output<-list()
  
  if (is.null(plot_params)){
    param<-fit_result$param
  } else {
    param<-plot_params
  }
  
  nonhyper_param<-param[startsWith(param,"alpha") | startsWith(param,"beta") | startsWith(param,"Beta") | startsWith(param,"kappa") | startsWith(param,"Kappa") | startsWith(param,"gamma")]
  hyper_param<-param[!(param %in% nonhyper_param)] 
  hyper_param_length<-length(hyper_param)
  
  ESS_hyper<-NULL
  
  hyper_param_name<-NULL
  
  if (hyper_param_length>0){
    for (jj in 1:hyper_param_length){
      hyper_param_name<-c(hyper_param_name,colnames(mcmc_object)[startsWith(colnames(mcmc_object),hyper_param[jj])])
    }
    hyper_param_name<-unique(hyper_param_name)  
  hyper_param_mat_subset<-mcmc_object[,hyper_param_name,drop=FALSE]
  ESS_hyper<-coda::effectiveSize(hyper_param_mat_subset)
  
  #trace and density plots
  
  plot(hyper_param_mat_subset,trace=trace,density=density)
  if (interactive() & (trace | density) & (length(nonhyper_param)>0)) {
    readline(prompt = "Press [Enter] to continue to predictor-related parameters...")
  }
  }
  
  k<-length(nonhyper_param)
  ESS_nonhyper<-NULL
  
  if (k>0){
  for (l in 1:k){
    
    if (grepl("[",nonhyper_param[l],fixed=TRUE)) {
      nonhyper_param_mat_subset<-mcmc_object[,nonhyper_param[l],drop=FALSE]
      ESS_nonhyper<-c(ESS_nonhyper,coda::effectiveSize(nonhyper_param_mat_subset))
      plot(nonhyper_param_mat_subset,trace=trace,density=density)
      
      if (interactive() & (trace | density) & l < k) {
        readline(prompt = "Press [Enter] to see the next plot...")
      }
      
    } else {
    
    nonhyper_param_l_names<-colnames(mcmc_object)[startsWith(colnames(mcmc_object),nonhyper_param[l])]
    plot_max<-min(3,length(nonhyper_param_l_names))
    nonhyper_param_l_names_sub<-sample(nonhyper_param_l_names,plot_max,replace=FALSE)
    nonhyper_param_mat_subset<-mcmc_object[,nonhyper_param_l_names_sub,drop=FALSE]
    ESS_nonhyper<-c(ESS_nonhyper,coda::effectiveSize(nonhyper_param_mat_subset))
    
    #trace and density plots
    
    plot(nonhyper_param_mat_subset,trace=trace,density=density)
    if (length(nonhyper_param_l_names)>3 & (trace | density)){insight::print_colour(paste0("NOTE: Only showing three randomly selected ",nonhyper_param[l],"."),"red")}
    if (interactive() & (trace | density) & l < k) {
      readline(prompt = "Press [Enter] to see the next plot...")
    }
    invisible(gc())
    }
  }
  }
  
  #ACF plots
  if (acf_plot){
    # for hyper-parameters
    total_length<-sum(startsWith(colnames(mcmc_object),hyper_param))
    if (total_length>0){
    m<-ceiling(total_length/3)
    remainder<-total_length %% 3
    if (m>1){
      for (ll in 1:(m-1)){
        acf(hyper_param_mat_subset[,(1:3)+(ll-1)*3,drop=FALSE])
        if (interactive()){readline(prompt = "Press [Enter] to see the next plot...")}
      }
    }
    if (remainder==0){
      acf(hyper_param_mat_subset[,(1:3)+(m-1)*3,drop=FALSE])
      if (interactive()){readline(prompt = "Press [Enter] to see the next plot...")}
    }
    if (remainder!=0){
      acf(hyper_param_mat_subset[,(1:remainder)+(m-1)*3,drop=FALSE])
      if (interactive()){readline(prompt = "Press [Enter] to see the next plot...")}
    }
    }
    
    # for non hyper-parameters
    if (k>0){
    for (l in 1:k){
      if (grepl("[",nonhyper_param[l],fixed=TRUE)){
        nonhyper_param_mat_subset<-mcmc_object[,nonhyper_param[l],drop=FALSE]
        acf(nonhyper_param_mat_subset)
        if (interactive()){readline(prompt = "Press [Enter] to see the next plot...")}
      } else {
      nonhyper_param_l_names<-colnames(mcmc_object)[startsWith(colnames(mcmc_object),nonhyper_param[l])]
      plot_max<-min(3,length(nonhyper_param_l_names))
      nonhyper_param_l_names_sub<-sample(nonhyper_param_l_names,plot_max,replace=FALSE)
      nonhyper_param_mat_subset<-mcmc_object[,nonhyper_param_l_names_sub,drop=FALSE]
      
      total_length<-length(nonhyper_param_l_names_sub)
      m<-ceiling(total_length/3)
      remainder<-total_length %% 3
      if (m>1){
      for (l in 1:(m-1)){
          acf(nonhyper_param_mat_subset[,(1:3)+(l-1)*3,drop=FALSE])
          if (interactive()){readline(prompt = "Press [Enter] to see the next plot...")}
        }
      }
      if (remainder==0){
        acf(nonhyper_param_mat_subset[,(1:3)+(m-1)*3,drop=FALSE])
        if (interactive()){readline(prompt = "Press [Enter] to see the next plot...")}
      }
      if (remainder!=0){
        acf(nonhyper_param_mat_subset[,(1:remainder)+(m-1)*3,drop=FALSE])
        if (interactive()){readline(prompt = "Press [Enter] to see the next plot...")}
      }
      if (length(nonhyper_param_l_names)>3){insight::print_colour(paste0("NOTE: Only showing three randomly selected parameters ",nonhyper_param[l],"."),"red")}
      if (interactive()){readline(prompt = "Press [Enter] to see the next plot...")}
      invisible(gc())
      }
    }
    }
  }
  
  if (ESS_all){
    output[["ESS"]]<-coda::effectiveSize(mcmc_object)
  } else {
    output[["ESS"]]<-c(ESS_hyper,ESS_nonhyper)
  }
  return(output)
  invisible(gc())
}