R/converge_diag_rates_fn.R

In BayesMoFo: Bayesian Mortality Forecasting

#' A function to assess convergence of the posterior sampling of death rates for monitoring purposes
#'
#' Produce several convergence diagnostic tools (e.g. trace/density/acf plots and effective sample sizes) from the posterior samples of death rates.
#' @param result object of type either "fit_result" or "BayesMoFo".
#' @param plot_strata A vector of character strings specifying which strata to plot for visualisation. If not specified, a random selection of the strata used to fit the model (i.e. \code{fit_result$death$strat_name}) will be chosen. 
#' @param plot_ages A numeric vector specifying which range of ages to plot for visualisation. If not specified, a random selection of ages that were used to fit the model (i.e. \code{fit_result$death$ages}) will be chosen. 
#' @param plot_years A numeric vector specifying which range of years to plot for visualisation. If not specified, a random selection of years that were used to fit the model (i.e. \code{fit_result$death$years}) will be chosen. 
#' @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 rates. The default is FALSE where only chosen rates will be evaluated, if TRUE all rates will be assessed.
#' @return Some convergence-related plots of posterior samples of mortality rates.
#' \describe{
#'   \item{\code{ESS}}{The effective sample sizes of the chosen parameters.}
#' }
#' @keywords bayesian diagnostics visualization
#' @concept posterior samples
#' @concept death rates
#' @concept trace plots
#' @concept density plots
#' @concept autocorrelations
#' @concept effective sample sizes
#' @importFrom graphics legend lines par points
#' @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_rates_fn(runBayesMoFo_result)
#' 
#' #ESS
#' converge_runBayesMoFo_result$ESS
#' 
#' #plot by age and changing pre-specified arguments 
#' converge_diag_rates_fn(runBayesMoFo_result,plot_ages=c(40,50,60),plot_years=c(2017,2020))
#' 
#' #ACF plot 
#' converge_diag_rates_fn(runBayesMoFo_result,plot_ages=c(40,50,60),plot_years=c(2017,2020),
#' trace=FALSE,density=FALSE,acf_plot=TRUE)
#' }

converge_diag_rates_fn<-function(result,plot_ages=NULL,plot_years=NULL,plot_strata=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
  }
  
  output<-list()
  
  death<-fit_result$death
  expo_data<-fit_result$expo$data
  if (fit_result$family=="binomial"){expo_data<-round(expo_data+0.5*death$data)}
  
  mcmc_object<-coda::as.mcmc(as.matrix(fit_result$post_sample))
  mcmc_object<-mcmc_object[,startsWith(colnames(mcmc_object),"q[")]
  invisible(gc())
  
  p<-death$n_strat
  A<-death$n_ages
  T<-death$n_years  
  h<-fit_result$h
  
  forecast<-fit_result$forecast
  
  n<-dim(mcmc_object)[1]
  
  p_names<-dimnames(death$data)[[1]]
  ages_names<-dimnames(death$data)[[2]]
  years_names<-dimnames(death$data)[[3]]

  if (is.null(plot_strata)){plot_strata=sample(death$strat_name,size=min(3,p),replace=FALSE)}
  
  if (is.null(plot_ages)){plot_ages=sample(death$ages,size=min(3,A),replace=FALSE)}
  
  if (is.null(plot_years)){
    if (forecast){plot_years=sample(c(death$years,death$years[T]+(1:h)),size=min(3,T),replace=FALSE)} else {
      plot_years=sample(death$years,size=min(3,T),replace=FALSE)}
  }
  
  strat_index<-(1:p)[death$strat_name %in% plot_strata]
  age_index<-(1:A)[death$ages %in% plot_ages]
  if (forecast){
    year_index<-(1:(T+h))[(c(death$years,death$years[T]+(1:h))) %in% plot_years]} else {
      year_index<-(1:T)[death$years %in% plot_years]
    }  

    subset_name<-NULL
    
    for (i in strat_index){
      for(j in age_index){
        for(k in year_index){
          index<-paste0("q[",i,",",j,",",k,"]")
          subset_name<-c(subset_name,index)
        }
      }
    }

    rates_mat_subset<-mcmc_object[,subset_name]
    
    if (ESS_all){
      output[["ESS"]]<-coda::effectiveSize(mcmc_object)
    } else {
      output[["ESS"]]<-coda::effectiveSize(rates_mat_subset)
    }
  
  #trace and density plots
    
    plot(rates_mat_subset,trace=trace,density=density,ask=TRUE)

  #ACF plots
  if (acf_plot){
    total_length<-length(strat_index)*length(age_index)*length(year_index)
    m<-ceiling(total_length/3)
    remainder<-total_length %% 3
    for (l in 1:(m-1)){
      acf(rates_mat_subset[,(1:3)+(l-1)*3])
      if (interactive()){readline(prompt = "Press [Enter] to see the next plot...")}
    }
    acf(rates_mat_subset[,(1:remainder)+(m-1)*3])
    if (interactive()){readline(prompt = "Press [Enter] to see the next plot...")}
  }
  invisible(gc())
  return(output)
}