R/estimates.R

In rferrali/rogali: A series of functions for the corrupt_audit project

#' \code{estimates} and \code{estimate} objects and methods
#' 
#' @param burn number or percent of iterations to discard
#' @param thin thinning interval. Keep every \code{thin}-th iteration.
#' @param what returns either parameter values or posterior probabilities Pr(Z = 1)
#' @param ... use to select variables a la \code{\link[dplyr]{select}} 
#' @param details a \code{logical} indicating whether the returned 
#' \code{tibble} should contain additional columns such as the iteration number 
#' and and id variable for the chain. 
#' @details An \code{estimates} object is merely a \code{list} of 
#' \code{estimate} objects. Each \code{estimate} object contains parameter 
#' samples for a given chain. They have a fairly complicated structure that 
#' I describe in more details below. Those objects should never be used 
#' directly, but rather be converted to more familiar data types. 
#' In particular, the \code{as.mcmc} and \code{as.mcmc.list} methods 
#' should be used to perform additional diagnostics using the \code{coda} 
#' package. 
#' @return An \code{estimate} object contains the following fields: 
#' \describe{
#'   \item{\code{beta}}{a \code{kBeta x 2 x nSamples} \code{array} with draws of beta coefficients}
#'   \item{\code{gamma}}{a \code{kGamma x nSamples} \code{matrix} with draws of gamma coefficients}
#'   \item{\code{pZ}}{a \code{nMun x nSamples} \code{matrix} with posterior Pr(Z = 1)}
#' }
#' @name estimatesmethods
NULL


bt <- function(m, burn, thin) {
  burn <- as.numeric(burn)
  thin <- as.integer(thin)
  if(length(burn) > 1 | burn < 0) stop("burn must be a non-negative scalar")
  if(length(thin) > 1 | thin < 1) stop("thin must be an integer >= 1")
  if(burn > 0) {
    m <- as.matrix(m)
    if(burn < 1) {
      m <- m[-(1:(round(burn * nrow(m)))),]
    } else {
      m <- m[-(1:burn),]
    }
  }
  if(thin > 1) {
    m <- as.matrix(m)
    m <- m[1:nrow(m) %% thin == 0,]
  }
  m
}

#' @rdname estimatesmethods
#' @export
as.matrix.estimate <- function(x, burn = 0, thin = 1, what = c("params", "probs")) {
  if(what[1] == "params") {
    m <- cbind(t(x$gamma), t(x$beta[,1,]), t(x$beta[,2,]))
    cnames <- c(
      paste0("gamma-", rownames(x$gamma)), 
      paste0("beta1-", dimnames(x$beta)[[1]]), 
      paste0("beta2-", dimnames(x$beta)[[1]]) 
    )
  } else {
    m <- t(x$pZ)
    cnames <- colnames(m)
  }
  m <- bt(m, burn, thin)
  colnames(m) <- cnames
  m
}


#' @rdname estimatesmethods
#' @export
as.matrix.estimates <- function(x, burn = 0, thin = 1, what = c("params", "probs")) {
  x <- purrr::map(x, as.matrix, burn = burn, thin = thin, what = what)
  do.call(rbind, x)
}

#' @rdname estimatesmethods
#' @export
as.mcmc.estimate <- function(x, ..., burn = 0, thin = 1, what = c("params", "probs")) {
  x <- as.matrix.estimate(x, what = what)
  end <- nrow(x)
  x <- bt(x, burn, 1)
  start <- end - nrow(x) + 1
  x <- bt(x, 0, thin)
  if(...length() > 0) {
    x <- select(as.data.frame(x), ... = ...)
    x <- as.matrix(x)
  }
  coda::mcmc(x, start = start, end = end, thin = thin)
}

#' @rdname estimatesmethods
#' @export
as.mcmc.list.estimates <- function(x, ..., burn = 0, thin = 1, what = c("params", "probs")) {
  x <- purrr::map(x, as.mcmc.estimate, burn = burn, thin = thin, what = what, ... = ...)
  coda::as.mcmc.list(x)
}

#' @rdname estimatesmethods
#' @export
as_tibble.estimate <- function(x, ..., burn = 0, thin = 1, what = c("params", "probs"), details = T) {
  if(details) it <- tibble::tibble(iteration = bt(1:ncol(x$gamma), burn, thin))
  x <- as.matrix.estimate(x, burn, thin, what = what)
  x <- tibble::as_tibble(x) 
  if(...length() > 0) x <- dplyr::select(x, ...)
  if(details) x <- bind_cols(it, x)
  x
}

#' @rdname estimatesmethods
#' @export
as_tibble.estimates <- function(x, ..., burn = 0, thin = 1, what = c("params", "probs"), details = T) {
  if(details) {
    id <- "chain"
  } else {
    id <- NULL
  }
  x <- purrr::map_dfr(
    x, as_tibble.estimate, 
    burn = burn, thin = thin, 
    details = details, ... = ..., 
    what = what, .id = id
  )
  if(details) x$chain <- as.integer(x$chain)
  x
}