R/tempo_diagnostics.R

In vlandau/tempo: Temporally-Explicit Models for Phenological Observations

#' Potential Scale Reduction Factor `(`PSRF`)`
#' 
#' Calculate the PSRF for parameters in \code{draws} output from 
#' \code{tempo_mcmc()}
#' 
#' @param draws The output from \code{tempo_mcmc(}). \code{draws} must contain
#' at least 3 chains in order to calculate the PSRF.
#' 
#' @return A named list of PSRF values for each parameter
#' 
#' @importFrom stats var
#' @rdname tempo_psrf
#' @name tempo_psrf
#' @export
tempo_psrf <- function(draws) {
  if (attr(draws, "rand_eff_bool") & !attr(draws, "monitor_random_effects")) {
    warning(
      "draws does not contain monitored random effects. To get the PSRF ",
      "for the random\n  effects (recommmended), rerun tempo_mcmc() with ",
      "monitor_random_effects = TRUE\n  and run tempo_psrf() on the resulting ",
      "object")
  }
  
  n_chains <- attr(draws, "n_chains")
  n_mcmc <- nrow(draws[[1]])
  
  if (n_chains < 3) {
    stop("draws must contain at least 3 chains in order to calculate the PSRF.")
  }
  
  # Convert list of matrices to list of lists of vectors
  param_names <- colnames(draws[[1]])
  
  params <- 
    lapply(1:length(param_names), FUN = function(i) lapply(draws, `[[`, i))
  names(params) <- param_names
  
  param_means <- lapply(params, tempo_chain_means)
  param_vars <- lapply(params, tempo_chain_variances)
  
  psrf <- mapply(param_means,
                 FUN = tempo_single_param_psrf,
                 param_vars,
                 MoreArgs = list(n_mcmc = n_mcmc, n_chains = n_chains))
  
  output <- sapply(psrf, list)
  
  return(output)
}

#' Deviance Information Criterion `(`DIC`)`
#' 
#' Calculate the DIC score for a model using the output from \code{tempo_mcmc()}
#' 
#' @return The DIC score with attributes for covariates included in the model,
#' and which (if any) regression coefficients were modeled with random effects,
#' and the effective number of parameters (\code{p_d}).
#' 
#' @param draws The output from \code{tempo_mcmc(}).
#' 
#' @param parallelize Boolean, should the function be run in parallel. Defaults
#' to \code{TRUE}.
#' 
#' @param n_cores Integer, number of cores to use for parallel processing. 
#' Defaults to \code{NULL}. If \code{parallelize = TRUE}, and 
#' \code{n_cores = NULL} (the defaults), \code{tempo_dic()} will use 
#' \code{parallel::detectCores() - 1} as the number of cores.
#' 
#' @importFrom parallel parSapply makeCluster stopCluster detectCores
#' @rdname tempo_dic
#' @name tempo_dic
#' @export
tempo_dic <- function(draws, parallelize = TRUE, n_cores = NULL) {
  rand_eff_bool <- attr(draws, "rand_eff_bool")
  if (rand_eff_bool & !attr(draws, "monitor_random_effects")) {
    stop("If random effects were used, draws must include monitored ",
         "group-level\n  effects in order to use this function. ",
         "Rerun tempo_mcmc() with \n  monitor_random_effects = TRUE")
  }
  
  n_chains <- attr(draws, "n_chains")
  n_mcmc <- nrow(draws[[1]])
  
  covs <- attr(draws, "covariates")
  y <- attr(draws, "y") # NAs already replaced
  group_ids <- attr(draws, "group_ids")
  
  # Need to calculate the DIC score for every iteration
  if (rand_eff_bool) {
    # combine random effects and parameters
    params_all <- tempo_calc_all_params_re(draws)
  } else {
    draws_all <- as.matrix(do.call(rbind, draws))
    
    params_all <- 
      draws_all[, str_detect(colnames(draws_all), pattern = "^beta"),
                drop = FALSE]
    params_all <- array(params_all, c(dim(params_all), 1))
    group_ids <- rep(1, nrow(y))
  }

  # define number of workers don't put inside of of (parallelize) control flow
  # for code coverage purposes
  if(!is.null(n_cores)) {
    n_workers <- n_cores 
  } else {
    n_workers <- detectCores() - 1
  }
  # make cluster
  if (parallelize) {
    cl <- makeCluster(max(1, n_workers))
    chunk_size <- (n_mcmc * n_chains) / (4 * length(cl))
  } else {
    cl <- makeCluster(1) 
    chunk_size <- n_mcmc * n_chains
  }

  # calculate deviance at every iter      
  deviance <- parSapply(cl = cl,
                        X = 1:(n_mcmc * n_chains), 
                        FUN = tempo_calc_deviance,
                        covs = covs,
                        y = y,
                        params_all = params_all,
                        group_ids = group_ids,
                        chunk.size = chunk_size)

  stopCluster(cl)
  
  # calculate deviance using posterior means
  deviance_of_mean <- 
    tempo_calc_deviance_of_mean(covs, y, params_all, group_ids)
  
  # effective number of parameters
  p_d <- mean(deviance) - deviance_of_mean
  
  # final dic calculation
  dic <- deviance_of_mean + 2 * p_d
  
  # save attributes
  attr(dic, "p_d") <- p_d
  attr(dic, "covariates") <- dimnames(covs)[[3]]
  if (rand_eff_bool) {
    attr(dic, "random_effects") <- dimnames(covs)[[3]][attr(draws, "rand_eff_ind")]
  } else {
    attr(dic, "random_effects") <- NA
  }
  
  dic
}

##### unexported utils for the above functions
# param is a named list of vectors (one vector per chain) for a given parameter
tempo_chain_means <- function(param) {
  unlist(lapply(param, mean))
}

# param is a named list of vectors (one vector per chain) for a given parameter
tempo_chain_variances <- function(param) {
  unlist(lapply(param, var))
}

# param_means is vector of per-chain posterior means for a single parameter
# param_vars is vector of per-chain posterior variances for a single parameter
# n_mcmc is the number of post-burn-in iters per chain, nrow(draws[[1]])
# n_chains is the number of chains, attr(draws, "n_chains")
tempo_single_param_psrf <- function(param_means, param_vars, n_mcmc, n_chains) {
  b <- (n_mcmc / (n_chains - 1)) * sum((param_means - mean(param_means))^2)
  w <- mean(param_vars)
  
  v_hat <- (1 - 1/n_mcmc) * w + b / n_mcmc
  
  return(sqrt(v_hat/w))
}


#' Deviance
#' 
#' Calculate deviance for a single MCMC iteration for random effects models
#' 
#' @param k the MCMC iteration for which to calculate deviance.
#' @param covs the covariate data from \code{tempo_mcmc()} output, \code{draws},
#' \code{attr(draws, "covariates")}
#' @param y the original data (same as the y arg to tempo_mcmc())
#' @param params_all output from \code{tempo_calc_all_params_re()}
#' @param group_ids is attr(draws, "group_ids"), a vector of group IDs 
#' corresponding for the rows in \code{y}
#' @importFrom tvgeom dtvgeom
#' @importFrom boot inv.logit
tempo_calc_deviance <- function(k, covs, y, params_all, group_ids) {
  params <- params_all[k, , ]
  
  if (is.vector(params)) { # because params (above) is vector if no rand eff
                           # needs to be matrix, t(t()) converts to 1-col matrix
    params <- t(t(params))
  }

  # rows = observation, cols = time step
  prob_matrix <- t(do.call(cbind, 
                           lapply(1:nrow(y),
                                  FUN = function(n, covs, params, group_ids){
                                    inv.logit(covs[n, ,] %*% 
                                                params[, group_ids[n]])
                                  },
                                  covs = covs,
                                  params = params,
                                  group_ids = group_ids))
  )
  rownames(prob_matrix) <- dimnames(covs)[[1]]
  
  #### calculate deviance from prob_matrix
  probs_df <- as.data.frame(t(prob_matrix)) # do this so it can be fed to mapply
  deviance <- sum(mapply(y$c1, y$c2, probs_df, FUN = tempo_single_obs_dev))
  
  deviance
}

# calc deviance for the posterior parameter means
tempo_calc_deviance_of_mean <- function(covs, y, params_all, group_ids) {
  params <- apply(params_all, c(2,3), mean)

  # rows = observation, cols = time step
  prob_matrix <- t(do.call(cbind, 
                           lapply(1:nrow(y),
                                  FUN = function(n, covs, params, group_ids){
                                    inv.logit(covs[n, ,] %*% 
                                                params[, group_ids[n]])
                                  },
                                  covs = covs,
                                  params = params,
                                  group_ids = group_ids))
  )
  rownames(prob_matrix) <- dimnames(covs)[[1]]
  
  #### calculate deviance from prob_matrix
  probs_df <- as.data.frame(t(prob_matrix)) # do this so it can be fed to mapply
  deviance <- sum(mapply(y$c1, y$c2, probs_df, FUN = tempo_single_obs_dev))
  
  deviance
}

tempo_single_obs_dev <- function(c1, c2, prob) {
  start <- c1 + 1
  end <- c2
  dens <- sum(dtvgeom(start:end, prob)) # total probability of observation 
                                        # falling in interval
  -2 * log(dens)
}