R/adapt_particleMCMC.R

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

#' @name adapt_particleMCMC
#' @title Particle MCMC Adaptation
#' @description 
#' Adapts proposal parameters of Particle MCMC scheme
#' to give near-optimal performance.
#' @param init initial value of epidemic parameters
#' @param epiModel epidemic model
#' @param obsFrame Generator function for observational model.
#' @param y Observed epidemic data.
#' @param X0 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 lambda0  Initial value of the Random Walk Metropolis (RWM) proposal scale parameter
#' @param V0 A square matrix containing the initial covariance values for RWM proposal. Shape
#'           of matrix should match the length of the init parameter.
#' @param K  Number of particles to be used in particle filter. This can be optimised using
#'           `adapt_BS_PF()`.
#' @param noIts Number of iterations of MCMC sampler scheme to carry out.
#' @param delta Probability that the initial proposal parameters are used as opposed to using 
#'              the current adapted proposal parameters.
#' @return 
#' A list of the adapted proposal parameters.
#' 
#' @export
adapt_particleMCMC <- function(init, epiModel, obsFrame, y, X0, alpha, logPrior, lambda0, V0, 
                               delta = 0.05, K, noIts){
  
  # Set up functions
  particleFilter <- VBS_PF(y, X0, obs = obsFrame, epiModel = epiModel)

  lambda <- lambda0
  # Estimate Likelihood for initial parameters
  logLikeCurr <- -Inf
  while(is.infinite(logLikeCurr)){
    logLikeCurr <- particleFilter(K, init, alpha)$logLikeEst
  }
  curr <- init
  accept <- 0
  lambda_vec <- lambda
  k <- length(init)
  draws <- matrix(ncol = k + 1, nrow = noIts)
  for(i in 1:noIts){
    
    adapt <- accept > 10 & (runif(1, 0, 1) > delta)
    # Propose new parameter
    
    if(adapt){
      V <- var(draws[1:(i-1), 1:k])
      prop <- abs(curr + mvtnorm::rmvnorm(1, mean = rep(0, k), sigma = (lambda^2)*V))
    } else{
      prop <- abs(curr + mvtnorm::rmvnorm(1, mean = rep(0, k), sigma = (lambda0^2)*V0))
    }
    
    # Estimate Likelihood
    logLikeProp <- particleFilter(K, prop, alpha)$logLikeEst

    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 <- accept + 1
        
        if(adapt){
          lambda <- lambda + 9*(0.01*lambda0/(i)^(1/3))
        }
      } else{
        if(adapt){
          lambda <- lambda - 1*(0.01*lambda0/(i)^(1/3))
        }
      }
    } else{
      if(adapt){
        lambda <- lambda - 1*(0.01*lambda0/(i)^(1/3))
      }
    }
    draws[i, ] <- c(curr, logLikeCurr)
    lambda_vec[i + 1] <- lambda
    #print(lambda)
  }
  par(mfrow = c(1,1))
  plot(lambda_vec, type = 'l')
  return(list(lambda = lambda, lambda_vec = lambda_vec, V = var(draws[,1:k]),
              curr = curr))
}