R/sf_output_month_prev.R

In AurMad/STOCfree: STOC free model: prediction of probability of freedom from longitudinal data

## function that checks that a dataset contains
## MCMC samples for monthly prevalences
## generated by the STOCfree_model() function
validate_STOCfree_month_prev <- function(x){

  ## input should be a data.frame
  if(!is.data.frame(x)) stop("x should be a data.frame")

  ## columns associated with mcmc draws should be present
  columns_x <- colnames(x)
  mcmc_cols <- match(c(".chain", ".iteration", ".draw"), columns_x)

  if(length(mcmc_cols) != 3) stop("wrong type of input")

  ## month columns
  month_cols <- columns_x[-mcmc_cols]

  month_abb <- format(seq.Date(as.Date("2020-01-01"),
                               as.Date("2020-12-01"), by= "1 month"),
                      "%b")

  check_month_col <- function(z){

    z_split <- strsplit(z, "_")

    month_abb <- format(seq.Date(as.Date("2020-01-01"),
                                 as.Date("2020-12-01"), by= "1 month"),
                        "%b")

    if(!z_split[[1]][1] %in% month_abb) stop("wrong type of input")

    if(as.integer(z_split[[1]][2]) < 0) stop("wrong type of input")

  }

  sapply(month_cols, check_month_col)

  invisible()

}

## Creation of STOCfree_month_prev objects
new_STOCfree_month_prev <- function(x){

  validate_STOCfree_month_prev(x)

  class(x) <- c("STOCfree_month_prev", class(x))

  return(x)

  }

#' Extracting draws from the posterior distributions for monthly prevalences
#'
#' @param STOCfree_model_output output from a call to STOCfree_model()
#'
#' @return
#' @export
extract_STOCfree_month_prev <- function(STOCfree_model_output, STOCfree_data = NULL){

  ## reconstructing the list of all months
  months_all <- data.frame(
    date = seq(as.Date(paste0(attr(STOCfree_data, "month first test"), "-01")),
               as.Date(paste0(attr(STOCfree_data, "month last test"), "-01")),
               by = "1 month"),
    month_abb = NA,
    tidybayes_names = NA
  )

  months_all <- months_all[-nrow(months_all), ]

  ## month names abbreviated to get meaningful column names
  months_all$month_abb <- format(months_all$date, "%b_%Y")


  if("CmdStanMCMC" %in% class(STOCfree_model_output)){ # if Stan model

    month_prev <- STOCfree_model_output$draws("month_prev")
    month_prev <- posterior::as_draws_df(month_prev)

    ## reordering columns for consistency with JAGS output
    id_cols1 <- match(c(".chain", ".iteration", ".draw"),
                      colnames(month_prev))
    id_cols2 <- (1:length(month_prev))[-id_cols1]

    month_prev <- month_prev[, c(id_cols1, id_cols2)]

    colnames(month_prev)[4:length(month_prev)] <- months_all$month_abb

  } else { # if not Stan model

  ## extracting MCMC draws from the STOCfree_model
  samples <- STOCfree_model_output$mcmc

  ## tidying the results for monthly prevalences
  ## model output made tidy
  month_prev <- tidybayes::spread_draws(samples,
                                        month_prev[..])

  ## month names in the month_prev dataset
  months_all$tidybayes_names <- paste0("month_prev.", 1:nrow(months_all))

  ## columns reordered
  cln <- match(months_all$tidybayes_names, colnames(month_prev))
  month_prev <- month_prev[, c(1:3, cln)]

  ## column names changed
  cln1 <- match(months_all$tidybayes_names, colnames(month_prev))
  colnames(month_prev)[cln1] <- months_all$month_abb

  }

  month_prev <- new_STOCfree_month_prev(month_prev)

  return(month_prev)

  }


#' Importing output files for predicted monthly prevalences generated by the STOC free model
#'
#' @param out_path
#'
#' @return
#' @export
read_STOCfree_month_prev <- function(out_path = "STOCfree_files"){

  nm <- paste0(out_path, "/month_prev.csv")
  month_prev <- new_STOCfree_month_prev(read.csv(nm))

  return(month_prev)

}


#' Boxplot of monthly prevalences predicted by the STOC free model
#'
#' @param month_prev an object of class STOCfree_month_prev
#'
#' @return
#' @export
plot.STOCfree_month_prev <- function(month_prev){

  mnth <- tidyr::pivot_longer(month_prev, cols = -(1:3), names_to = "month", values_to = "Prevalence")
  mnth <- dplyr::mutate(mnth, date = as.Date(paste0("15-", month), format = "%d-%b_%Y"))

  ggplot2::ggplot(mnth, ggplot2::aes(x = factor(date), y = Prevalence)) +
    ggplot2::geom_boxplot() +
    ggplot2::theme(axis.text.x = ggplot2::element_text(angle = 90, vjust = 0.5, hjust=1)) +
    ggplot2::xlab("")

}