R/plot_gMAP.R

In RBesT: R Bayesian Evidence Synthesis Tools

#' Diagnostic plots for gMAP analyses
#'
#' @param x [gMAP()] object
#' @param size Controls line sizes of traceplots and forest plot.
#' @param ... Ignored.
#'
#' @details Creates MCMC diagnostics and a forest plot (including
#' model estimates) for a [gMAP()] analysis. For a
#' customized forest plot, please use the dedicated function
#' [forest_plot()].
#'
#' @template plot-help
#'
#' @return The function returns a list of [ggplot2::ggplot()]
#' objects.
#'
#' @method plot gMAP
#' @export
plot.gMAP <- function(x, size = NULL, ...) {
  pl <- list()

  draws_all <- rstan::extract(x$fit, permuted = FALSE, inc_warmup = TRUE)
  thin <- attr(x$fit, "sim")$thin
  n_warmup <- floor(attr(x$fit, "sim")$warmup / thin)
  draws <- draws_all[-(1:n_warmup), , ]
  nuts_diag <- bayesplot::nuts_params(x$fit, inc_warmup = FALSE)

  ## by default we return a small set of plots only
  plot_verbose <- getOption("RBesT.verbose", FALSE)

  div_opts <- list()
  if (sum(nuts_diag$Value[nuts_diag$Parameter == "divergent__"]) > 0) {
    div_opts$np <- bayesplot::nuts_params(x$fit, inc_warmup = TRUE)
    plot_verbose <- TRUE ## if any divergent transition happens,
    ## then we plot verbose in any case
  }

  tau_pars <- grep("tau\\[", dimnames(draws)$parameters, value = TRUE)
  beta_pars <- grep("beta\\[", dimnames(draws)$parameters, value = TRUE)

  tau_log_trans <- as.list(rep("log", length(tau_pars)))
  names(tau_log_trans) <- tau_pars

  if (plot_verbose) {
    ## traces are only shown if in verbose mode...
    pl$traceBeta <- do.call(
      bayesplot::mcmc_trace,
      c(
        list(
          x = draws_all,
          pars = beta_pars,
          size = size,
          n_warmup = n_warmup,
          facet_args = list(labeller = ggplot2::label_parsed)
        ),
        div_opts
      )
    ) +
      ggtitle(expression(paste("Trace of Regression Coefficient ", beta))) +
      bayesplot::facet_text(length(beta_pars) != 1) +
      xlab("Iteration")
    pl$traceTau <- do.call(
      bayesplot::mcmc_trace,
      c(
        list(
          x = draws_all,
          pars = tau_pars,
          size = size,
          n_warmup = n_warmup,
          facet_args = list(labeller = ggplot2::label_parsed)
        ),
        div_opts
      )
    ) +
      bayesplot::facet_text(length(tau_pars) != 1) +
      ggtitle(expression(paste("Trace of Heterogeneity Parameter ", tau))) +
      xlab("Iteration")
    pl$traceLogTau <- do.call(
      bayesplot::mcmc_trace,
      c(
        list(
          x = draws_all,
          pars = tau_pars,
          size = size,
          n_warmup = n_warmup,
          facet_args = list(labeller = ggplot2::label_parsed),
          transformations = tau_log_trans
        ),
        div_opts
      )
    ) +
      bayesplot::facet_text(length(tau_pars) != 1) +
      ggtitle(expression(paste(
        "Trace of Heterogeneity Parameter ",
        tau,
        " on log-scale"
      ))) +
      xlab("Iteration")

    ## ... as well as auxilary model parameters
    pl$densityBeta <- bayesplot::mcmc_dens_overlay(
      x = draws,
      pars = beta_pars,
      facet_args = list(
        labeller = ggplot2::label_parsed,
        strip.position = "bottom"
      )
    ) +
      ggtitle(expression(paste("Density of Regression Coefficient ", beta)))
    pl$densityTau <- bayesplot::mcmc_dens_overlay(
      x = draws,
      pars = tau_pars,
      facet_args = list(
        labeller = ggplot2::label_parsed,
        strip.position = "bottom"
      )
    ) +
      ggtitle(expression(paste("Density of Heterogeneity Parameter ", tau)))
    pl$densityLogTau <- bayesplot::mcmc_dens_overlay(
      x = draws,
      pars = tau_pars,
      facet_args = list(
        labeller = ggplot2::label_parsed,
        strip.position = "bottom"
      ),
      transformations = tau_log_trans
    ) +
      ggtitle(expression(paste(
        "Density of Heterogeneity Parameter ",
        tau,
        " on log-scale"
      )))
  }

  if (x$has_intercept) {
    pl$densityThetaStar <- bayesplot::mcmc_dens_overlay(
      x = draws,
      pars = "theta_resp_pred"
    ) +
      xlab(expression(theta[symbol("*")])) +
      bayesplot::facet_text(FALSE) +
      ggtitle(expression(paste("Density of MAP Prior ", theta[symbol("*")])))
    pl$densityThetaStarLink <- bayesplot::mcmc_dens_overlay(
      x = draws,
      pars = "theta_pred"
    ) +
      xlab(expression(theta[symbol("*")])) +
      bayesplot::facet_text(FALSE) +
      ggtitle(expression(paste(
        "Density of MAP Prior ",
        theta[symbol("*")],
        " (link scale)"
      )))

    pl$forest_model <- forest_plot(
      x,
      model = "both",
      size = if (is.null(size)) 1.25 else size
    )
  } else {
    message("No intercept defined.")
  }

  return(pl)
}