R/ISMCMC.R

In JMacDonaldPhD/REpi: Pseudo-Marginal Analysis of partially observed Epidemics

#' @name ISMCMC
#' @title Independence Sampler/RWM MCMC
#' @description 
#' Generates dependent samples from target distribution
#' using a Independence Sampler/RWM MCMC proposal scheme.
#' @param init initial value of epidemic parameters
#' @param U_init initial value of random variables used to construct
#'               epidemic trajectory.
#' @param epiModel epidemic model
#' @param obsFrame Generator function for observational model.
#' @param epiSample Observed epidemic data.
#' @param I0 Initial state of the epidemic, which assumed to be known.
#' @param alpha Observational parameters to be passed on to the log-likelihood function
#'              generated by `obsFrame(X_sim)`
#' @param logPrior Functions which calculate log-density of the assumed priors of the 
#'                 epidemic parameters.
#' @param lambda Random Walk Metropolis (RWM) proposal scale parameter.
#' @param V A square matrix containing the covariance values for RWM proposal. Shape
#'          of matrix should match the length of the init parameter.
#' @param S Number of non-centered random variables to be 'refreshed' in Independence
#'          Sampler proposal step.
#' @param noIts Number of iterations of MCMC sampler scheme to carry out.
#' @return 
#' A matrix of MCMC samples and proposal acceptance rate of the sample.
#' 
#' @export
ISMCMC <- function(init, U_init, epiModel, obsFrame, epiSample, I0, alpha, logPrior, lambda, V, S, noIts){
  # Set up functions
  k <- length(init)
  r <- length(U_init)
  # Estimate Likelihood for initial parameters
  epi_init <- epiModel$sim(list(I0, init, U_init))
  obsModelInit <- obsFrame(epi_init)
  logLikeCurr <- obsModel$llh(epiSample, alpha)
  while(is.infinite(logLikeCurr)){
    U_init <- matrix(runif(r, 0, 1), ncol = ncol(U_init), nrow = nrow(U_init))
    epi_init <- epiModel$sim(list(I0, init, U_init))
    obsModelInit <- obsFrame(epi_init)
    logLikeCurr <- obsModelInit$llh(epiSample, alpha)
  }
  curr <- init
  U_curr <- U_init
  epi_curr <- epi_init
  obsModelCurr <- obsModelInit
  
  accept.theta <- 0
  accept.S <- 0
  draws <- matrix(ncol = k + 1, nrow = noIts)
  for(i in 1:noIts){
    # Propose new epidemic parameters
    prop <- abs(curr + mvnfast::rmvn(1, mu = rep(0, k), sigma = lambda*V))

    # Estimate Likelihood
    
    epiProp <- epiModel$sim(list(I0, prop, U_curr))
    obsModelProp <- obsFrame(epiProp)
    logLikeProp <- obsModelProp$llh(epiSample, alpha)
    
    
    if(!is.infinite(logLikeProp)){
      logAccProb <- (logLikeProp + logPrior(prop)) - (logLikeCurr + logPrior(curr)) 
      #print(logAccPRob)
      if(log(runif(1, 0, 1)) < logAccProb){
        curr <- prop
        logLikeCurr <- logLikeProp 
        accept.theta <- accept.theta + 1
      }
    }
    
    U_prop <- U_curr
    U_prop[sample(1:r, S, replace = F)] <- runif(S, 0, 1)
    
    # Estimate Likelihood
    epiProp <- epiModel$sim(list(I0, curr, U_prop))
    obsModelProp <- obsFrame(epiProp)
    logLikeProp <- obsModelProp$llh(epiSample, alpha)
    
    if(!is.infinite(logLikeProp)){
      logAccProb <- (logLikeProp) - (logLikeCurr) 
      #print(logAccPRob)
      if(log(runif(1, 0, 1)) < logAccProb){
        U_curr <- U_prop
        logLikeCurr <- logLikeProp 
        accept.S <- accept.S + 1
      }
    }
    draws[i, ] <- c(curr, logLikeCurr)
  }
  return(list(draws = draws, acceptRate = c(accept.theta/noIts, accept.S/noIts)))
}