R/filter-fit-locality.R

In rplzzz/CovMitigation: COVID-19 infection model for Virginia

### Input file location for local tests:
#ensemble_file_dir <- here::here('analysis', 'bayesian-filter-tests')

#' Run the filter model for a single locality
#'
#' This function requires a \code{metrosamp} save file with parameter samples
#' for the locality.  An ensemble of models is initialized from these samples,
#' and the filter is run for each ensemble member.
#'
#' The \code{rndselect} argument allows one to select random entries from the MCMC
#' samples, consistently across localities (i.e., each locality will get the same
#' parameters for each ensemble member).  The input for \code{rndselect} should be
#' a vector of \code{Nens} random numbers between 0 and 1.  These will be converted
#' to integer values in 1:N, where N is the total number of MCMC samples.
#'
#' @param locality Either the name of the locality, or the number in the master
#' list.  There are 133 localities.
#' @param inputdir Directory containing the input files generated by \code{metrosamp}.
#' The file name will be constructed by appending the locality name to the input file prefix.
#' @param Nens Number of ensemble members.
#' @param rndselect Vector of random ensemble members to select (see details).
#' @param inputpfx Stem of input file name.  Default is \code{'filter-ensemble-candidates-'}.
#' The name of the locality will be appended to this, followed by the input file suffix.
#' @param inputsfx Suffix appended to the input filename.  Default is \code{'.rds'}.
#' @param datemax Maximum date to use in the observed data.  Default is to exclude
#' the last week if it is incomplete (i.e., daily observations aren't current through
#' the end of the week).
#' @param history Flag indicating whether to keep full results from all ensemble members.
#' If \code{FALSE}, keep only median and confidence intervals at each time step,
#' plus final state
#' @export
filter_fit_locality <- function(locality,
                                inputdir,
                                Nens=100,
                                rndselect=NULL,
                                inputpfx = 'filter-ensemble-candidates-',
                                inputsfx = '.rds',
                                datemax=NULL, history=TRUE)
{
  ## constants
  jan01 <- as.Date('2020-01-01')
  filter_strtdate <- jan01 + 123     # May 3rd.  This gets us through the end of April
  localities <- sort(unique(vdhcovid::va_weekly_ntest_county$locality))

  ## check if we were passed a number
  lnum <- suppressWarnings(as.integer(locality))
  if(!is.na(lnum)) {
    locality <- localities[lnum]
  }

  ensemble_filename <-
    file.path(inputdir, paste0(inputpfx, locality, inputsfx))
  message('filename:  ', ensemble_filename)
  stopifnot(file.exists(ensemble_filename))
  ensemble_ms <- readRDS(ensemble_filename)

  if(inherits(ensemble_ms, 'metrosamp')) {
    ## Extract ensemble from metropolis structure
    if(is.null(rndselect)) {
      ensemble <- metrosamp::getsamples(ensemble_ms, Nens)
    }
    else {
      stopifnot(length(rndselect) == Nens)
      ensemble <- metrosamp::getsamples(ensemble_ms)
      ntot <- nrow(ensemble)
      rndselect <- as.integer(ceiling(rndselect*ntot))
      ensemble <- ensemble[rndselect, ]
    }
  }
  else {
    stop('Input file does not contain a metrosamp object.')
  }
  stopifnot(is.matrix(ensemble))

  if(!'import_cases' %in% colnames(ensemble)) {
    ensemble <- cbind(ensemble, import_cases=0)
  }
  betacol <- which(colnames(ensemble) == paste0('beta_',locality))
  colnames(ensemble)[betacol] <- 'beta'

  ## Run each of these parameter sets to get an initial state
  t0 <- locality_startdate(locality)
  totpop <- vdhcovid::getpop(locality)
  t1 <- as.numeric(filter_strtdate - jan01)
  tvals0 <- seq(t0, t1)

  ### TODO:  Refactor this to use initialize_parmset
  initrun <-
    function(i) {
      parms <- ensemble[i,]
      E0 <- parms[['I0']]
      vars0 <- c(time=t0, S=totpop-E0, E=E0, I=0, Is=0, R=0)
      st <- run_parmset(parms, vars0, tvals0)
      vars <- st[nrow(st), ]
      c(vars, parms)
    }

  initstates <- do.call(rbind,
                        lapply(seq(1, Nens), initrun))

  ## Get the observed data
  ### TODO:  make a way to get this data through the public interface
  obs <- get_obsdata()
  obsdata <- obs$obsdata
  obsdata <- obsdata[obsdata$locality == locality, ]
  hosp <- obs$hosp[obs$hosp$locality == locality, ]
  
  if(is.null(datemax)) {
    datemax <- max(vdhcovid::vaweeklytests$date)
  }
  obsdata <- obsdata[obsdata$date <= datemax, ]

  fit_filter(initstates, obsdata, t1, max(obsdata$time), history = TRUE)
}