Merge pull request #163 from mrc-ide/nimue_fitting

OJWatson · web-flow · commit 7b828c5b3439 · 2021-03-12T13:47:06.000Z
v0.6.4 changes for nimue vaccine fitting
diff --git a/DESCRIPTION b/DESCRIPTION
@@ -1,7 +1,7 @@
 Package: squire
 Type: Package
 Title: SEIR transmission model of COVID-19
-Version: 0.6.3
+Version: 0.6.4
 Authors@R: c(
   person("OJ", "Watson", email = "o.watson15@imperial.ac.uk", role = c("aut", "cre")),
   person("Patrick", "Walker", email = "patrick.walker06@imperial.ac.uk", role = c("aut")),
diff --git a/NEWS.md b/NEWS.md
@@ -1,3 +1,7 @@
+# squire 0.6.4
+
+* `pmcmc` changes to adapt for `nimue` class models and vaccines
+
 # squire 0.6.3
 
 * `beta_est` accepts `nimue` models
diff --git a/R/models.R b/R/models.R
@@ -13,10 +13,50 @@ explicit_model <- function() {
     compare_output(model, pars_obs, data, type=model_class)
   }
 
+  parameter_func <-  function(..., tt_vaccine, max_vaccine) {
+
+    # build model parameters with no vaccine being passed through as
+    # this is not the vaccine model
+    parameters_explicit_SEEIR(...)
+
+  }
+
+  run_func <-  function(country,
+                        contact_matrix_set,
+                        tt_contact_matrix,
+                        hosp_bed_capacity,
+                        tt_hosp_beds,
+                        ICU_bed_capacity,
+                        tt_ICU_beds,
+                        max_vaccine,
+                        tt_vaccine,
+                        population,
+                        replicates,
+                        day_return,
+                        time_period,
+                        ...) {
+
+    # build model run with no vaccine being passed through as
+    # this is not the vaccine model
+    run_explicit_SEEIR_model(country = country,
+                             contact_matrix_set = contact_matrix_set,
+                             tt_contact_matrix = tt_contact_matrix,
+                             hosp_bed_capacity = hosp_bed_capacity,
+                             tt_hosp_beds = tt_hosp_beds,
+                             ICU_bed_capacity = ICU_bed_capacity,
+                             tt_ICU_beds = tt_ICU_beds,
+                             population = population,
+                             replicates = replicates,
+                             day_return = day_return,
+                             time_period = time_period,
+                             ...)
+  }
+
+
   explicit_model <- list(odin_model = explicit_SEIR,
                          generate_beta_func = beta_est_explicit,
-                         parameter_func = parameters_explicit_SEEIR,
-                         run_func = run_explicit_SEEIR_model,
+                         parameter_func = parameter_func,
+                         run_func = run_func,
                          compare_model = compare_model)
   class(explicit_model) <- c(model_class, "stochastic", "squire_model")
   explicit_model
@@ -61,10 +101,48 @@ deterministic_model <- function() {
     compare_output(model, pars_obs, data, type=model_class)
   }
 
+  parameter_func <-  function(..., tt_vaccine, max_vaccine) {
+
+    # build model parameters with no vaccine being passed through as
+    # this is not the vaccine model
+    parameters_explicit_SEEIR(...)
+  }
+
+  run_func <-  function(country,
+                        contact_matrix_set,
+                        tt_contact_matrix,
+                        hosp_bed_capacity,
+                        tt_hosp_beds,
+                        ICU_bed_capacity,
+                        tt_ICU_beds,
+                        max_vaccine,
+                        tt_vaccine,
+                        population,
+                        replicates,
+                        day_return,
+                        time_period,
+                        ...) {
+
+    # build model run with no vaccine being passed through as
+    # this is not the vaccine model
+    run_deterministic_SEIR_model(country = country,
+                             contact_matrix_set = contact_matrix_set,
+                             tt_contact_matrix = tt_contact_matrix,
+                             hosp_bed_capacity = hosp_bed_capacity,
+                             tt_hosp_beds = tt_hosp_beds,
+                             ICU_bed_capacity = ICU_bed_capacity,
+                             tt_ICU_beds = tt_ICU_beds,
+                             population = population,
+                             replicates = replicates,
+                             day_return = day_return,
+                             time_period = time_period,
+                             ...)
+  }
+
   det_model <- list(odin_model = explicit_SEIR_deterministic,
                        generate_beta_func = beta_est_explicit,
-                       parameter_func = parameters_explicit_SEEIR,
-                       run_func = run_deterministic_SEIR_model,
+                       parameter_func = parameter_func,
+                       run_func = run_func,
                        compare_model = compare_model)
   class(det_model) <- c(model_class, "deterministic", "squire_model")
   det_model
diff --git a/R/pmcmc.R b/R/pmcmc.R
@@ -33,6 +33,10 @@
 #' @param output_proposals Logical indicating whether proposed parameter jumps should be output along with results
 #' @param n_chains number of MCMC chains to run
 #' @inheritParams calibrate
+#' @param date_vaccine_change Date that vaccine doses per day change.
+#'   Default = NULL.
+#' @param baseline_max_vaccine Baseline vaccine doses per day. Default = NULL
+#' @param max_vaccine Time varying maximum vaccine doeses per day. Default = NULL.
 #' @param Rt_args List of arguments to be passed to \code{evaluate_Rt_pmcmc} for calculating Rt.
 #'   Current arguments are available in \code{Rt_args_list}
 #' @param burnin number of iterations to discard from the start of MCMC run when sampling from the posterior for trajectories
@@ -161,6 +165,9 @@ pmcmc <- function(data,
                   ICU_bed_capacity = NULL,
                   baseline_ICU_bed_capacity = NULL,
                   date_ICU_bed_capacity_change = NULL,
+                  date_vaccine_change = NULL,
+                  baseline_max_vaccine = NULL,
+                  max_vaccine = NULL,
                   Rt_args = NULL,
                   burnin = 0,
                   replicates = 100,
@@ -180,6 +187,11 @@ pmcmc <- function(data,
   # assertions & checks
   #--------------------
 
+  # if nimue keep to 1 step per day
+  if(inherits(squire_model, "nimue_model")) {
+    steps_per_day <- 1
+  }
+
   # we work with pars_init being a list of inital conditions for starting
   if(any(c("start_date", "R0") %in% names(pars_init))) {
     pars_init <- list(pars_init)
@@ -312,9 +324,6 @@ pmcmc <- function(data,
     if(as.Date(tail(date_contact_matrix_set_change,1)) > as.Date(tail(data$date, 1))) {
       stop("Last date in date_contact_matrix_set_change is greater than the last date in data")
     }
-    if(as.Date(pars_max$start_date) >= as.Date(head(date_contact_matrix_set_change, 1))) {
-      stop("First date in date_contact_matrix_set_change is earlier than maximum start date allowed in pars search")
-    }
 
     # Get in correct format
     if(is.matrix(baseline_contact_matrix)) {
@@ -343,9 +352,6 @@ pmcmc <- function(data,
     if(as.Date(tail(date_ICU_bed_capacity_change,1)) > as.Date(tail(data$date, 1))) {
       stop("Last date in date_ICU_bed_capacity_change is greater than the last date in data")
     }
-    if(as.Date(pars_max$start_date) >= as.Date(head(date_ICU_bed_capacity_change, 1))) {
-      stop("First date in date_ICU_bed_capacity_change is earlier than maximum start date of epidemic")
-    }
 
     tt_ICU_beds <- c(0, seq_len(length(date_ICU_bed_capacity_change)))
     ICU_bed_capacity <- c(baseline_ICU_bed_capacity, ICU_bed_capacity)
@@ -355,6 +361,33 @@ pmcmc <- function(data,
     ICU_bed_capacity <- baseline_ICU_bed_capacity
   }
 
+  # handle vaccine changes
+  if(!is.null(date_vaccine_change)) {
+
+    assert_date(date_vaccine_change)
+    assert_vector(max_vaccine)
+    assert_numeric(max_vaccine)
+    assert_numeric(baseline_max_vaccine)
+
+    if(is.null(baseline_max_vaccine)) {
+      stop("baseline_max_vaccine can't be NULL if date_vaccine_change is provided")
+    }
+    if(as.Date(tail(date_vaccine_change,1)) > as.Date(tail(data$date, 1))) {
+      stop("Last date in date_vaccine_change is greater than the last date in data")
+    }
+
+    tt_vaccine <- c(0, seq_len(length(date_vaccine_change)))
+    max_vaccine <- c(baseline_max_vaccine, max_vaccine)
+
+  } else {
+    tt_vaccine <- 0
+    if(!is.null(baseline_max_vaccine)) {
+    max_vaccine <- baseline_max_vaccine
+    } else {
+      max_vaccine <- 0
+    }
+  }
+
   # handle hosp bed changed
   if(!is.null(date_hosp_bed_capacity_change)) {
 
@@ -369,9 +402,6 @@ pmcmc <- function(data,
     if(as.Date(tail(date_hosp_bed_capacity_change,1)) > as.Date(tail(data$date, 1))) {
       stop("Last date in date_hosp_bed_capacity_change is greater than the last date in data")
     }
-    if(as.Date(pars_max$start_date) >= as.Date(head(date_hosp_bed_capacity_change, 1))) {
-      stop("First date in date_hosp_bed_capacity_change is earlier than maximum start date of epidemic")
-    }
 
     tt_hosp_beds <- c(0, seq_len(length(date_hosp_bed_capacity_change)))
     hosp_bed_capacity <- c(baseline_hosp_bed_capacity, hosp_bed_capacity)
@@ -403,6 +433,8 @@ pmcmc <- function(data,
                                               tt_hosp_beds = tt_hosp_beds,
                                               ICU_bed_capacity = ICU_bed_capacity,
                                               tt_ICU_beds = tt_ICU_beds,
+                                              max_vaccine = max_vaccine,
+                                              tt_vaccine = tt_vaccine,
                                               ...)
 
   # collect interventions for odin model likelihood
@@ -413,7 +445,9 @@ pmcmc <- function(data,
                         date_ICU_bed_capacity_change = date_ICU_bed_capacity_change,
                         ICU_bed_capacity = ICU_bed_capacity,
                         date_hosp_bed_capacity_change = date_hosp_bed_capacity_change,
-                        hosp_bed_capacity = hosp_bed_capacity)
+                        hosp_bed_capacity = hosp_bed_capacity,
+                        date_vaccine_change = date_vaccine_change,
+                        max_vaccine = max_vaccine)
 
   #----------------..
   # Collect Odin and MCMC Inputs
@@ -609,10 +643,13 @@ pmcmc <- function(data,
                              tt_hosp_beds = tt_hosp_beds,
                              ICU_bed_capacity = ICU_bed_capacity,
                              tt_ICU_beds = tt_ICU_beds,
+                             max_vaccine = max_vaccine,
+                             tt_vaccine = tt_vaccine,
                              population = population,
                              replicates = 1,
                              day_return = TRUE,
-                             time_period = nrow(pmcmc_samples$trajectories))
+                             time_period = nrow(pmcmc_samples$trajectories),
+                             ...)
 
   # and add the parameters that changed between each simulation, i.e. posterior draws
   r$replicate_parameters <- pmcmc_samples$sampled_PMCMC_Results
@@ -1195,6 +1232,7 @@ calc_loglikelihood <- function(pars, data, squire_model, model_params,
   date_contact_matrix_set_change <- interventions$date_contact_matrix_set_change
   date_ICU_bed_capacity_change <- interventions$date_ICU_bed_capacity_change
   date_hosp_bed_capacity_change <- interventions$date_hosp_bed_capacity_change
+  date_vaccine_change <- interventions$date_vaccine_change
 
   # change betas
   if (is.null(date_R0_change)) {
@@ -1245,6 +1283,19 @@ calc_loglikelihood <- function(pars, data, squire_model, model_params,
     model_params$hosp_beds <- tt_list$change
   }
 
+  # and vaccine coverage
+  if (is.null(date_vaccine_change)) {
+    tt_vaccine <- 0
+    max_vaccine <- 0
+  } else {
+    tt_list <- intervention_dates_for_odin(dates = sort(unique(c(start_date,date_vaccine_change))),
+                                           change = interventions$max_vaccine,
+                                           start_date = start_date,
+                                           steps_per_day = round(1/model_params$dt))
+    model_params$tt_vaccine <- tt_list$tt
+    model_params$max_vaccine <- tt_list$change
+  }
+
   #--------------------..
   # update new R0s based on R0_change and R0_date_change, and Meff_date_change
   #--------------------..
diff --git a/R/projections.R b/R/projections.R
@@ -107,7 +107,7 @@ projections <- function(r,
   # ----------------------------------------------------------------------------
   ## assertion checks on parameters
   # ----------------------------------------------------------------------------
-  assert_custom_class(r, "squire_simulation")
+  # assert_custom_class(r, "squire_simulation")
   # TODO future asserts if these are our "end" classes
   if (is.null(r$output) & is.null(r$scan_results) & is.null(r$pmcmc_results)) {
     stop("Model must have been produced either with Squire Default, Scan Grid (calibrate), or pMCMC (pmcmc) Approach")
diff --git a/man/pmcmc.Rd b/man/pmcmc.Rd
diff --git a/tests/testthat/test-pmcmc.R b/tests/testthat/test-pmcmc.R
