From fdbb841148eb6c9c7c43a4ce611d77b1a55b70e1 Mon Sep 17 00:00:00 2001
From: Paul Lietar <pl2113@ic.ac.uk>
Date: Tue, 13 Feb 2024 17:26:26 +0000
Subject: [PATCH 1/2] Allow simulation to be restored with new interventions.
 (#286)

This uses a few of individual's new features that make restoring more
flexible. It also fixes a bug when restoring the mosquitto solvers, by
correctly restoring the timestep, which is needed to model seasonality.

The intervention events use the new `restore = FALSE` flag to make sure
their schedule can be modified when resuming. Instead of having the
events re-schedule themselves everytime they fire, we setup the entire
schedule upfront when initialising the simulation.

It adds end-to-end testing of this feature, across a range of scenarios.
For each scenario, the outcomes of the simulation with and without
restoring are compared and we make sure they are equivalent.
---
 DESCRIPTION                  |   2 +-
 R/RcppExports.R              |   4 +-
 R/compartmental.R            |   8 +-
 R/correlation.R              |   9 +-
 R/events.R                   |  18 ++--
 R/lag.R                      |   2 +-
 R/mda_processes.R            |   7 --
 R/model.R                    |  17 ++--
 R/pev.R                      |   5 -
 R/stateful.R                 |  47 ----------
 R/tbv.R                      |   5 -
 R/utils.R                    |  24 +++++
 man/CorrelationParameters.Rd |   7 +-
 src/RcppExports.cpp          |   9 +-
 src/solver.cpp               |   3 +-
 tests/testthat/test-mda.R    |  15 ---
 tests/testthat/test-pev.R    |   7 --
 tests/testthat/test-resume.R | 174 +++++++++++++++++++++++++++++------
 18 files changed, 217 insertions(+), 146 deletions(-)
 delete mode 100644 R/stateful.R

diff --git a/DESCRIPTION b/DESCRIPTION
index 76eb8165..74445497 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -72,7 +72,7 @@ Remotes:
 Additional_repositories:
   https://mrc-ide.r-universe.dev
 Imports:
-    individual (>= 0.1.15),
+    individual (>= 0.1.16),
     malariaEquilibrium (>= 1.0.1),
     Rcpp,
     statmod,
diff --git a/R/RcppExports.R b/R/RcppExports.R
index 678a50bf..ad658f2c 100644
--- a/R/RcppExports.R
+++ b/R/RcppExports.R
@@ -53,8 +53,8 @@ solver_get_states <- function(solver) {
     .Call(`_malariasimulation_solver_get_states`, solver)
 }
 
-solver_set_states <- function(solver, state) {
-    invisible(.Call(`_malariasimulation_solver_set_states`, solver, state))
+solver_set_states <- function(solver, t, state) {
+    invisible(.Call(`_malariasimulation_solver_set_states`, solver, t, state))
 }
 
 solver_step <- function(solver) {
diff --git a/R/compartmental.R b/R/compartmental.R
index 7df83a6e..4e16c7c1 100644
--- a/R/compartmental.R
+++ b/R/compartmental.R
@@ -155,8 +155,8 @@ Solver <- R6::R6Class(
     save_state = function() {
       solver_get_states(private$.solver)
     },
-    restore_state = function(state) {
-      solver_set_states(private$.solver, state)
+    restore_state = function(t, state) {
+      solver_set_states(private$.solver, t, state)
     }
   )
 )
@@ -173,7 +173,7 @@ AquaticMosquitoModel <- R6::R6Class(
     # state of the ODE is stored separately). We still provide these methods to
     # conform to the expected interface.
     save_state = function() { NULL },
-    restore_state = function(state) { }
+    restore_state = function(t, state) { }
   )
 )
 
@@ -187,7 +187,7 @@ AdultMosquitoModel <- R6::R6Class(
     save_state = function() {
       adult_mosquito_model_save_state(self$.model)
     },
-    restore_state = function(state) {
+    restore_state = function(t, state) {
       adult_mosquito_model_restore_state(self$.model, state)
     }
   )
diff --git a/R/correlation.R b/R/correlation.R
index 458a3015..2a59fdbc 100644
--- a/R/correlation.R
+++ b/R/correlation.R
@@ -125,11 +125,16 @@ CorrelationParameters <- R6::R6Class(
     },
 
     #' @description Restore the correlation state.
+    #'
     #' Only the randomly drawn weights are restored. The object needs to be
     #' initialized with the same rhos.
+    #'
+    #' @param timestep the timestep at which simulation is resumed. This
+    #' parameter's value is ignored, it only exists to conform to a uniform
+    #' interface.
     #' @param state a previously saved correlation state, as returned by the
-    #'   save_state method.
-    restore_state = function(state) {
+    #' save_state method.
+    restore_state = function(timestep, state) {
       private$.mvnorm <- state$mvnorm
     }
   )
diff --git a/R/events.R b/R/events.R
index d7ad3416..de0c7c1f 100644
--- a/R/events.R
+++ b/R/events.R
@@ -1,11 +1,11 @@
 create_events <- function(parameters) {
   events <- list(
     # MDA events
-    mda_administer = individual::Event$new(),
-    smc_administer = individual::Event$new(),
+    mda_administer = individual::Event$new(restore=FALSE),
+    smc_administer = individual::Event$new(restore=FALSE),
 
     # TBV event
-    tbv_vaccination = individual::Event$new(),
+    tbv_vaccination = individual::Event$new(restore=FALSE),
 
     # Bednet events
     throw_away_net = individual::TargetedEvent$new(parameters$human_population)
@@ -21,7 +21,7 @@ create_events <- function(parameters) {
       seq_along(parameters$mass_pev_booster_spacing),
       function(.) individual::TargetedEvent$new(parameters$human_population)
     )
-    events$mass_pev <- individual::Event$new()
+    events$mass_pev <- individual::Event$new(restore=FALSE)
     events$mass_pev_doses <- mass_pev_doses
     events$mass_pev_boosters <- mass_pev_boosters
   }
@@ -63,16 +63,16 @@ initialise_events <- function(events, variables, parameters) {
 
   # Initialise scheduled interventions
   if (!is.null(parameters$mass_pev_timesteps)) {
-    events$mass_pev$schedule(parameters$mass_pev_timesteps[[1]] - 1)
+    events$mass_pev$schedule(parameters$mass_pev_timesteps - 1)
   }
   if (parameters$mda) {
-    events$mda_administer$schedule(parameters$mda_timesteps[[1]] - 1)
+    events$mda_administer$schedule(parameters$mda_timesteps - 1)
   }
   if (parameters$smc) {
-    events$smc_administer$schedule(parameters$smc_timesteps[[1]] - 1)
+    events$smc_administer$schedule(parameters$smc_timesteps - 1)
   }
   if (parameters$tbv) {
-    events$tbv_vaccination$schedule(parameters$tbv_timesteps[[1]] - 1)
+    events$tbv_vaccination$schedule(parameters$tbv_timesteps - 1)
   }
 }
 
@@ -158,7 +158,6 @@ attach_event_listeners <- function(
   if (parameters$mda == 1) {
     events$mda_administer$add_listener(create_mda_listeners(
       variables,
-      events$mda_administer,
       parameters$mda_drug,
       parameters$mda_timesteps,
       parameters$mda_coverages,
@@ -174,7 +173,6 @@ attach_event_listeners <- function(
   if (parameters$smc == 1) {
     events$smc_administer$add_listener(create_mda_listeners(
       variables,
-      events$smc_administer,
       parameters$smc_drug,
       parameters$smc_timesteps,
       parameters$smc_coverages,
diff --git a/R/lag.R b/R/lag.R
index ca858ac7..de8d1653 100644
--- a/R/lag.R
+++ b/R/lag.R
@@ -20,7 +20,7 @@ LaggedValue <- R6::R6Class(
       timeseries_save_state(private$history)
     },
 
-    restore_state = function(state) {
+    restore_state = function(t, state) {
       timeseries_restore_state(private$history, state)
     }
   )
diff --git a/R/mda_processes.R b/R/mda_processes.R
index 42960be4..72cbe779 100644
--- a/R/mda_processes.R
+++ b/R/mda_processes.R
@@ -1,6 +1,5 @@
 #' @title Create listeners for MDA events
 #' @param variables the variables available in the model
-#' @param administer_event the event schedule for drug administration
 #' @param drug the drug to administer
 #' @param timesteps timesteps for each round
 #' @param coverages the coverage for each round
@@ -14,7 +13,6 @@
 #' @noRd
 create_mda_listeners <- function(
   variables,
-  administer_event,
   drug,
   timesteps,
   coverages,
@@ -78,11 +76,6 @@ create_mda_listeners <- function(
       variables$drug$queue_update(drug, to_move)
       variables$drug_time$queue_update(timestep, to_move)
     }
-
-    # Schedule next round
-    if (time_index < length(timesteps)) {
-      administer_event$schedule(timesteps[[time_index + 1]] - timestep)
-    }
   }
 }
 
diff --git a/R/model.R b/R/model.R
index def6db23..67f79758 100644
--- a/R/model.R
+++ b/R/model.R
@@ -135,16 +135,19 @@ run_resumable_simulation <- function(
   lagged_eir <- create_lagged_eir(variables, solvers, parameters)
   lagged_infectivity <- create_lagged_infectivity(variables, parameters)
 
-  stateful_objects <- unlist(list(
+  stateful_objects <- list(
     RandomState$new(restore_random_state),
     correlations,
     vector_models,
     solvers,
     lagged_eir,
-    lagged_infectivity))
+    lagged_infectivity)
 
   if (!is.null(initial_state)) {
-    restore_state(initial_state$malariasimulation, stateful_objects)
+    individual::restore_object_state(
+      initial_state$timesteps,
+      stateful_objects,
+      initial_state$malariasimulation)
   }
 
   individual_state <- individual::simulation_loop(
@@ -161,16 +164,16 @@ run_resumable_simulation <- function(
       timesteps
     ),
     variables = variables,
-    events = unlist(events),
+    events = events,
     timesteps = timesteps,
     state = initial_state$individual,
     restore_random_state = restore_random_state
   )
 
   final_state <- list(
-    timesteps=timesteps,
-    individual=individual_state,
-    malariasimulation=save_state(stateful_objects)
+    timesteps = timesteps,
+    individual = individual_state,
+    malariasimulation = individual::save_object_state(stateful_objects)
   )
 
   data <- renderer$to_dataframe()
diff --git a/R/pev.R b/R/pev.R
index bc27fc9f..e460e5c3 100644
--- a/R/pev.R
+++ b/R/pev.R
@@ -130,11 +130,6 @@ create_mass_pev_listener <- function(
       parameters,
       events$mass_pev_doses
     )
-    if (time_index < length(parameters$mass_pev_timesteps)) {
-      events$mass_pev$schedule(
-        parameters$mass_pev_timesteps[[time_index + 1]] - timestep
-      )
-    }
   }
 }
 
diff --git a/R/stateful.R b/R/stateful.R
deleted file mode 100644
index 05c614a0..00000000
--- a/R/stateful.R
+++ /dev/null
@@ -1,47 +0,0 @@
-#' @title Save the state of a list of \emph{stateful objects}
-#' @description The state of each element is saved and stored into a single
-#' object, representing them in a way that can be exported and re-used later.
-#' @param objects A list of stateful objects to be saved. Stateful objects are
-#' instances of R6 classes with a pair of \code{save_state} and
-#' \code{restore_state} methods.
-#' @noRd
-save_state <- function(objects) {
-  lapply(objects, function(o) o$save_state())
-}
-
-#' @title Restore the state of a collection of stateful objects
-#' @description This is the counterpart of \code{save_state}. Calling it
-#' restores the collection of objects back into their original state.
-#' @param state A state object, as returned by the \code{save_state} function.
-#' @param objects A collection of stateful objects to be restored.
-#' @noRd
-restore_state <- function(state, objects) {
-  stopifnot(length(state) == length(objects))
-  for (i in seq_along(state)) {
-    objects[[i]]$restore_state(state[[i]])
-  }
-}
-
-#' @title a placeholder class to save the random number generator class.
-#' @description the class integrates with the simulation loop to save and
-#' restore the random number generator class when appropriate.
-#' @noRd
-RandomState <- R6::R6Class(
-  'RandomState',
-  private = list(
-    restore_random_state = NULL
-  ),
-  public = list(
-    initialize = function(restore_random_state) {
-      private$restore_random_state <- restore_random_state
-    },
-    save_state = function() {
-      random_save_state()
-    },
-    restore_state = function(state) {
-      if (private$restore_random_state) {
-        random_restore_state(state)
-      }
-    }
-  )
-)
diff --git a/R/tbv.R b/R/tbv.R
index 27c29dfb..d0f75de1 100644
--- a/R/tbv.R
+++ b/R/tbv.R
@@ -75,11 +75,6 @@ create_tbv_listener <- function(variables, events, parameters, correlations, ren
         to_vaccinate
       )
     }
-    if (time_index < length(parameters$tbv_timesteps)) {
-      events$tbv_vaccination$schedule(
-        parameters$tbv_timesteps[[time_index + 1]] - timestep
-      )
-    }
   }
 }
 
diff --git a/R/utils.R b/R/utils.R
index f25244bd..2262c939 100644
--- a/R/utils.R
+++ b/R/utils.R
@@ -64,3 +64,27 @@ rtexp <- function(n, m, t) { itexp(runif(n), m, t) }
 match_timestep <- function(ts, t) {
   min(sum(ts <= t), length(ts))
 }
+
+#' @title a placeholder class to save the random number generator class.
+#' @description the class integrates with the simulation loop to save and
+#' restore the random number generator class when appropriate.
+#' @noRd
+RandomState <- R6::R6Class(
+  'RandomState',
+  private = list(
+    restore_random_state = NULL
+  ),
+  public = list(
+    initialize = function(restore_random_state) {
+      private$restore_random_state <- restore_random_state
+    },
+    save_state = function() {
+      random_save_state()
+    },
+    restore_state = function(t, state) {
+      if (private$restore_random_state) {
+        random_restore_state(state)
+      }
+    }
+  )
+)
diff --git a/man/CorrelationParameters.Rd b/man/CorrelationParameters.Rd
index 2898aee5..8bb43d70 100644
--- a/man/CorrelationParameters.Rd
+++ b/man/CorrelationParameters.Rd
@@ -121,15 +121,20 @@ Save the correlation state.
 \if{latex}{\out{\hypertarget{method-CorrelationParameters-restore_state}{}}}
 \subsection{Method \code{restore_state()}}{
 Restore the correlation state.
+
 Only the randomly drawn weights are restored. The object needs to be
 initialized with the same rhos.
 \subsection{Usage}{
-\if{html}{\out{<div class="r">}}\preformatted{CorrelationParameters$restore_state(state)}\if{html}{\out{</div>}}
+\if{html}{\out{<div class="r">}}\preformatted{CorrelationParameters$restore_state(timestep, state)}\if{html}{\out{</div>}}
 }
 
 \subsection{Arguments}{
 \if{html}{\out{<div class="arguments">}}
 \describe{
+\item{\code{timestep}}{the timestep at which simulation is resumed. This
+parameter's value is ignored, it only exists to conform to a uniform
+interface.}
+
 \item{\code{state}}{a previously saved correlation state, as returned by the
 save_state method.}
 }
diff --git a/src/RcppExports.cpp b/src/RcppExports.cpp
index 439cb12b..e2a6b9d6 100644
--- a/src/RcppExports.cpp
+++ b/src/RcppExports.cpp
@@ -203,13 +203,14 @@ BEGIN_RCPP
 END_RCPP
 }
 // solver_set_states
-void solver_set_states(Rcpp::XPtr<Solver> solver, std::vector<double> state);
-RcppExport SEXP _malariasimulation_solver_set_states(SEXP solverSEXP, SEXP stateSEXP) {
+void solver_set_states(Rcpp::XPtr<Solver> solver, double t, std::vector<double> state);
+RcppExport SEXP _malariasimulation_solver_set_states(SEXP solverSEXP, SEXP tSEXP, SEXP stateSEXP) {
 BEGIN_RCPP
     Rcpp::RNGScope rcpp_rngScope_gen;
     Rcpp::traits::input_parameter< Rcpp::XPtr<Solver> >::type solver(solverSEXP);
+    Rcpp::traits::input_parameter< double >::type t(tSEXP);
     Rcpp::traits::input_parameter< std::vector<double> >::type state(stateSEXP);
-    solver_set_states(solver, state);
+    solver_set_states(solver, t, state);
     return R_NilValue;
 END_RCPP
 }
@@ -363,7 +364,7 @@ static const R_CallMethodDef CallEntries[] = {
     {"_malariasimulation_rainfall", (DL_FUNC) &_malariasimulation_rainfall, 5},
     {"_malariasimulation_exponential_process_cpp", (DL_FUNC) &_malariasimulation_exponential_process_cpp, 2},
     {"_malariasimulation_solver_get_states", (DL_FUNC) &_malariasimulation_solver_get_states, 1},
-    {"_malariasimulation_solver_set_states", (DL_FUNC) &_malariasimulation_solver_set_states, 2},
+    {"_malariasimulation_solver_set_states", (DL_FUNC) &_malariasimulation_solver_set_states, 3},
     {"_malariasimulation_solver_step", (DL_FUNC) &_malariasimulation_solver_step, 1},
     {"_malariasimulation_create_timeseries", (DL_FUNC) &_malariasimulation_create_timeseries, 2},
     {"_malariasimulation_timeseries_at", (DL_FUNC) &_malariasimulation_timeseries_at, 3},
diff --git a/src/solver.cpp b/src/solver.cpp
index 0a5f7401..4c3f182f 100644
--- a/src/solver.cpp
+++ b/src/solver.cpp
@@ -48,7 +48,8 @@ std::vector<double> solver_get_states(Rcpp::XPtr<Solver> solver) {
 }
 
 //[[Rcpp::export]]
-void solver_set_states(Rcpp::XPtr<Solver> solver, std::vector<double> state) {
+void solver_set_states(Rcpp::XPtr<Solver> solver, double t, std::vector<double> state) {
+    solver->t = t;
     solver->state = state;
 }
 
diff --git a/tests/testthat/test-mda.R b/tests/testthat/test-mda.R
index bf92ae3a..1ee440f6 100644
--- a/tests/testthat/test-mda.R
+++ b/tests/testthat/test-mda.R
@@ -81,11 +81,8 @@ test_that('MDA moves the diseased and non-diseased population correctly', {
     drug = mock_double(c(1, 2, 1, 2))
   )
   
-  events$mda_administer <- mock_event(events$mda_administer)
-  
   listener <- create_mda_listeners(
     variables,
-    events$mda_administer,
     parameters$mda_drug,
     parameters$mda_timesteps,
     parameters$mda_coverages,
@@ -141,12 +138,6 @@ test_that('MDA moves the diseased and non-diseased population correctly', {
     50,
     c(3, 4)
   )
-  
-  mockery::expect_args(
-    events$mda_administer$schedule,
-    1,
-    100
-  )
 })
 
 test_that('MDA moves the diseased and non-diseased population correctly - second round, varying age range', {
@@ -176,11 +167,8 @@ test_that('MDA moves the diseased and non-diseased population correctly - second
     drug = mock_double(c(1, 2, 1, 2))
   )
   
-  events$mda_administer <- mock_event(events$mda_administer)
-  
   listener <- create_mda_listeners(
     variables,
-    events$mda_administer,
     parameters$mda_drug,
     parameters$mda_timesteps,
     parameters$mda_coverages,
@@ -265,11 +253,8 @@ test_that('MDA ignores non-detectable asymptomatics', {
     drug = mock_double(c(1, 2, 1, 2))
   )
   
-  events$mda_administer <- mock_event(events$mda_administer)
-  
   listener <- create_mda_listeners(
     variables,
-    events$mda_administer,
     parameters$mda_drug,
     parameters$mda_timesteps,
     parameters$mda_coverages,
diff --git a/tests/testthat/test-pev.R b/tests/testthat/test-pev.R
index 88b95b35..09c4cf7d 100644
--- a/tests/testthat/test-pev.R
+++ b/tests/testthat/test-pev.R
@@ -203,7 +203,6 @@ test_that('Mass vaccinations update vaccination time', {
     c(-1, -1, -1, 50, 50)
   )
 
-  events$mass_pev <- mock_event(events$mass_pev)
   events$mass_pev_doses <- lapply(events$mass_pev_doses, mock_event)
 
   listener <- create_mass_pev_listener(
@@ -241,12 +240,6 @@ test_that('Mass vaccinations update vaccination time', {
     c(1, 3),
     parameters$pev_doses[[3]]
   )
-
-  mockery::expect_args(
-    events$mass_pev$schedule,
-    1,
-    365
-  )
 })
 
 test_that('Mass vaccinations ignore EPI individuals', {
diff --git a/tests/testthat/test-resume.R b/tests/testthat/test-resume.R
index 83bada35..4d4a572d 100644
--- a/tests/testthat/test-resume.R
+++ b/tests/testthat/test-resume.R
@@ -1,41 +1,161 @@
-test_that('Simulation can be resumed', {
-  initial_timesteps <- 50
-  total_timesteps <- 100
-
-  parameters <- get_parameters()
+#' Test simulation saving and restoring for a given parameter set.
+#'
+#' This function runs the simulation twice. A first, continuous and uninterrupted
+#' run of the simulation is used as a reference. The second run is split into
+#' two phases. Between the two phases, the simulation state is saved and
+#' restored. Optionally, the initial warmup phase can use a different set of
+#' parameters, by specifying a value for warmup_parameters.
+test_resume <- function(
+  parameters,
+  timesteps = 200,
+  warmup_parameters = parameters,
+  warmup_timesteps = 50
+  ) {
+  set.seed(123)
+  uninterrupted_run <- run_simulation(timesteps, parameters=parameters)
 
-  set.seed(1)
-  first_phase <- run_resumable_simulation(initial_timesteps, parameters)
+  set.seed(123)
+  first_phase <- run_resumable_simulation(warmup_timesteps, warmup_parameters)
   second_phase <- run_resumable_simulation(
-    total_timesteps,
+    timesteps,
     parameters,
     initial_state=first_phase$state,
     restore_random_state=TRUE)
 
-  set.seed(1)
-  uninterrupted_run <- run_simulation(total_timesteps, parameters=parameters)
+  expect_equal(nrow(first_phase$data), warmup_timesteps)
+  expect_equal(nrow(second_phase$data), timesteps - warmup_timesteps)
+
+  expect_mapequal(
+    second_phase$data,
+    uninterrupted_run[-(1:warmup_timesteps),])
+
+  invisible(second_phase$data)
+}
+
+test_that('Simulation can be resumed', {
+  test_resume(get_parameters(overrides=list(
+    individual_mosquitoes = FALSE,
+    model_seasonality = TRUE
+  )))
+  test_resume(get_parameters(overrides=list(
+    individual_mosquitoes = TRUE,
+    model_seasonality = TRUE
+  )))
+  test_resume(get_parameters(overrides=list(
+    individual_mosquitoes = FALSE,
+    model_seasonality = TRUE
+  )))
+  test_resume(get_parameters(overrides=list(
+    individual_mosquitoes = TRUE,
+    model_seasonality = TRUE
+  )))
+})
+
+test_that('PEV intervention can be added when resuming', {
+  set_default_mass_pev <- function(parameters, timesteps) {
+    n <- length(timesteps)
+    set_mass_pev(
+      parameters,
+      profile = rtss_profile,
+      timesteps = timesteps,
+      coverages = rep(0.5, n),
+      min_wait = 5,
+      min_ages = 365*10,
+      max_ages = 365*60,
+      booster_spacing = NULL,
+      booster_coverage = NULL,
+      booster_profile = NULL)
+  }
+  base <- get_parameters(overrides=list(pev_doses=c(0, 45, 90)))
+
+  data <- test_resume(
+    warmup_parameters = base,
+    parameters = base %>% set_default_mass_pev(100))
+  expect_equal(data[data$n_pev_mass_dose_1 > 0, "timestep"], 100)
+  expect_equal(data[data$n_pev_mass_dose_2 > 0, "timestep"], 145)
+  expect_equal(data[data$n_pev_mass_dose_3 > 0, "timestep"], 190)
+
+  # Add a second mass PEV intervention when resuming the simulation.
+  data <- test_resume(
+    warmup_parameters = base %>% set_default_mass_pev(25),
+    parameters = base %>% set_default_mass_pev(c(25, 100)))
 
-  expect_equal(nrow(first_phase$data), initial_timesteps)
-  expect_equal(nrow(second_phase$data), total_timesteps - initial_timesteps)
-  expect_equal(rbind(first_phase$data, second_phase$data), uninterrupted_run)
+  # The first dose, at time step 25, happens during the warmup and is not
+  # returned by test_resume, hence why we don't see it in the data. Follow-up
+  # doses do show up, even though they we scheduled during warmup.
+  expect_equal(data[data$n_pev_mass_dose_1 > 0, "timestep"], c(100))
+  expect_equal(data[data$n_pev_mass_dose_2 > 0, "timestep"], c(70, 145))
+  expect_equal(data[data$n_pev_mass_dose_3 > 0, "timestep"], c(115, 190))
 })
 
-test_that('Intervention parameters can be changed when resuming', {
-  initial_timesteps <- 50
-  total_timesteps <- 100
-  tbv_timesteps <- 70
+test_that("TBV intervention can be added when resuming", {
+  set_default_tbv <- function(parameters, timesteps) {
+    set_tbv(
+      parameters,
+      timesteps=timesteps,
+      coverage=rep(1, length(timesteps)),
+      ages=5:60)
+  }
+
+  base <- get_parameters()
+
+  data <- test_resume(
+    warmup_parameters = base,
+    parameters = base %>% set_default_tbv(100))
+  expect_equal(data[!is.na(data$n_vaccinated_tbv), "timestep"], 100)
+
+  data <- test_resume(
+    warmup_parameters = base %>% set_default_tbv(25),
+    parameters = base %>% set_default_tbv(c(25, 100)))
+  expect_equal(data[!is.na(data$n_vaccinated_tbv), "timestep"], 100)
+})
+
+test_that("MDA intervention can be added when resuming", {
+  set_default_mda <- function(parameters, timesteps) {
+    parameters %>% set_drugs(list(SP_AQ_params)) %>% set_mda(
+      drug = 1,
+      timesteps = timesteps,
+      coverages = rep(0.8, length(timesteps)),
+      min_ages = rep(0, length(timesteps)),
+      max_ages = rep(60*365, length(timesteps)))
+  }
+
+  base <- get_parameters()
+
+  data <- test_resume(
+    warmup_parameters = base,
+    parameters = base %>% set_default_mda(100))
+  expect_equal(data[data$n_mda_treated > 0, "timestep"], 100)
+
+  data <- test_resume(
+    warmup_parameters = base %>% set_default_mda(25),
+    parameters = base %>% set_default_mda(c(25, 100)))
+  expect_equal(data[data$n_mda_treated > 0, "timestep"], 100)
+})
 
-  # Because of how event scheduling works, we must enable TBV even in the inital phase.
-  # We set a coverage of 0 to act as-if it was disabled.
-  initial_parameters <- get_parameters() |> set_tbv(timesteps=tbv_timesteps, coverage=0, ages=5:60)
+test_that("Bednets intervention can be added when resuming", {
+  set_default_bednets <- function(parameters, timesteps) {
+    n <- length(timesteps)
+    set_bednets(
+      parameters,
+      timesteps = timesteps,
+      coverages = rep(0.5, n),
+      retention = 25,
+      dn0 = matrix(rep(0.533, n), ncol=1),
+      rn = matrix(rep(0.56, n), ncol=1),
+      rnm = matrix(rep(0.24, n), ncol=1),
+      gamman = rep(2.64 * 365, n))
+  }
 
-  tbv_parameters <- initial_parameters |>
-    set_tbv(timesteps=tbv_timesteps, coverage=1, ages=5:60)
+  base <- get_parameters()
 
-  initial_run <- run_resumable_simulation(initial_timesteps, initial_parameters)
-  control_run <- run_resumable_simulation(total_timesteps, initial_parameters, initial_state = initial_run$state)
-  tbv_run <- run_resumable_simulation(total_timesteps, tbv_parameters, initial_state = initial_run$state)
+  data <- test_resume(
+    warmup_parameters = base,
+    parameters = base %>% set_default_bednets(100))
+  expect_equal(data[diff(data$n_use_net) > 0, "timestep"], 100)
 
-  expect_equal(control_run$data$n_vaccinated_tbv[tbv_timesteps - initial_timesteps], 0)
-  expect_gt(tbv_run$data$n_vaccinated_tbv[tbv_timesteps - initial_timesteps], 0)
+  data <- test_resume(
+    warmup_parameters = base %>% set_default_bednets(25),
+    parameters = base %>% set_default_bednets(c(25, 100)))
+  expect_equal(data[diff(data$n_use_net) > 0, "timestep"], 100)
 })

From ac3703b8c5150fc0fa54bcb246614b881cf9f55d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Paul=20Li=C3=A9tar?= <pl2113@ic.ac.uk>
Date: Tue, 30 Apr 2024 15:20:12 +0100
Subject: [PATCH 2/2] Allow correlations to be extended when resuming. (#293)

When resuming a simulation, it is possible to add new intervention.
The correlation of the new intervention may need configuring, both
relative to itself and to existing interventions.

The correlation parameters work by sampling, at the start of the
simulation, relative weights for each individual and intervention. The
weight are later used to select which individuals are targeted by the a
given intervention. There weights are drawn from a multivariate normal
distribution, whose variance-covariance matrix depends on the configured
correlation.

When adding interventions, new corresponding columns need to be added to
the weights matrix. A fresh mvnorm distribution cannot be used, since it
would override the already drawn weigths for the existing interventions.
The solution instead is to use a conditional mvnorm distribution on the
new columns, using the existing values as the conditions. This yields
new weigths which follow the determined variance-covariance matrix while
preserving the existing values.
---
 R/correlation.R                   | 162 +++++++++++++++++++---
 man/CorrelationParameters.Rd      |  31 ++++-
 tests/testthat/test-correlation.R | 218 ++++++++++++++++++++++++++++++
 tests/testthat/test-resume.R      | 150 +++++++++++++++++++-
 4 files changed, 535 insertions(+), 26 deletions(-)

diff --git a/R/correlation.R b/R/correlation.R
index 2a59fdbc..bb98fb3b 100644
--- a/R/correlation.R
+++ b/R/correlation.R
@@ -9,7 +9,26 @@ INTS <- c(
 )
 
 #' Class: Correlation parameters
-#' Describes an event in the simulation
+#'
+#' This class implements functionality that allows interventions to be
+#' correlated, positively or negatively. By default, interventions are applied
+#' independently and an individual's probability of receiving two interventions
+#' (either two separate interventions or two rounds of the same one) is the
+#' product of the probability of receiving each one.
+#'
+#' By setting a positive correlation between two interventions, we can make it
+#' so that the individuals that receive intervention A are more likely to
+#' receive intervention B. Conversely, a negative correlation will make it such
+#' that individuals that receive intervention A are less likely to also receive
+#' intervention B.
+#'
+#' Broadly speaking, the implementation works by assigning at startup a weight
+#' to each individual and intervention pair, reflecting how likely an individual
+#' is to receive that intervention. Those weights are derived stochastically
+#' from the configured correlation parameters.
+#'
+#' For a detailed breakdown of the calculations, see Protocol S2 of
+#' Griffin et al. (2010).
 CorrelationParameters <- R6::R6Class(
   'CorrelationParameters',
   private = list(
@@ -19,7 +38,40 @@ CorrelationParameters <- R6::R6Class(
     rho_matrix = NULL,
     rho = function() diag(private$rho_matrix),
     .sigma = NULL,
-    .mvnorm = NULL
+    .mvnorm = NULL,
+
+    #' Derive the mvnorm from the configured correlations.
+    #'
+    #' If a \code{restored_mvnorm} is specified, its columns (corresponding to
+    #' restored interventions) will be re-used as is. Missing columns (for new
+    #' interventions) are derived in accordance with the restored data.
+    calculate_mvnorm = function(restored_mvnorm = matrix(ncol=0, nrow=private$population)) {
+      sigma <- self$sigma()
+      V <- outer(sigma, sigma) * private$rho_matrix
+      diag(V) <- sigma ^ 2
+
+      restored_interventions <- match(colnames(restored_mvnorm), private$interventions)
+      new_interventions <- setdiff(seq_along(private$interventions), restored_interventions)
+
+      mvnorm <- matrix(
+        nrow = private$population,
+        ncol = length(private$interventions),
+        dimnames = list(NULL, private$interventions)
+      )
+      mvnorm[,restored_interventions] <- restored_mvnorm
+      if (length(new_interventions) > 0) {
+        mvnorm[,new_interventions] <- rcondmvnorm(
+          private$population,
+          mean = rep(0, length(private$interventions)),
+          sigma = V,
+          given = restored_mvnorm,
+          dependent.ind = new_interventions,
+          given.ind = restored_interventions
+        )
+      }
+
+      mvnorm
+    }
   ),
   public = list(
 
@@ -45,6 +97,8 @@ CorrelationParameters <- R6::R6Class(
     #' @param rho value between 0 and 1 representing the correlation between rounds of
     #' the intervention
     inter_round_rho = function(int, rho) {
+      stopifnot(is.null(private$.sigma) && is.null(private$.mvnorm))
+
       if (!(int %in% private$interventions)) {
         stop(paste0('invalid intervention name: ', int))
       }
@@ -55,8 +109,6 @@ CorrelationParameters <- R6::R6Class(
         rho <- 1 - .Machine$double.eps
       }
       private$rho_matrix[[int, int]] <- rho
-      private$.sigma <- NULL
-      private$.mvnorm <- NULL
     },
 
     #' @description Add rho between interventions
@@ -66,6 +118,8 @@ CorrelationParameters <- R6::R6Class(
     #' @param rho value between -1 and 1 representing the correlation between rounds of
     #' the intervention
     inter_intervention_rho = function(int_1, int_2, rho) {
+      stopifnot(is.null(private$.sigma) && is.null(private$.mvnorm))
+
       if (!(int_1 %in% private$interventions)) {
         stop(paste0('invalid intervention name: ', int_1))
       }
@@ -83,8 +137,6 @@ CorrelationParameters <- R6::R6Class(
       }
       private$rho_matrix[[int_1, int_2]] <- rho
       private$rho_matrix[[int_2, int_1]] <- rho
-      private$.sigma <- NULL
-      private$.mvnorm <- NULL
     },
 
     #' @description Standard deviation of each intervention between rounds
@@ -98,18 +150,9 @@ CorrelationParameters <- R6::R6Class(
     },
 
     #' @description multivariate norm draws for these parameters
-    #' @importFrom MASS mvrnorm
     mvnorm = function() {
       if (is.null(private$.mvnorm)) {
-        sigma <- self$sigma()
-        V <- outer(sigma, sigma) * private$rho_matrix
-        diag(V) <- sigma ^ 2
-        private$.mvnorm <- mvrnorm(
-          private$population,
-          rep(0, length(private$interventions)),
-          V
-        )
-        dimnames(private$.mvnorm)[[2]] <- private$interventions
+        private$.mvnorm <- private$calculate_mvnorm()
       }
       private$.mvnorm
     },
@@ -121,7 +164,7 @@ CorrelationParameters <- R6::R6Class(
       # otherwise we would be drawing a new, probably different, value.
       # The rest of the object is derived deterministically from the parameters
       # and does not need saving.
-      list(mvnorm=private$.mvnorm)
+      list(mvnorm=self$mvnorm())
     },
 
     #' @description Restore the correlation state.
@@ -135,7 +178,8 @@ CorrelationParameters <- R6::R6Class(
     #' @param state a previously saved correlation state, as returned by the
     #' save_state method.
     restore_state = function(timestep, state) {
-      private$.mvnorm <- state$mvnorm
+      stopifnot(is.null(private$.sigma) && is.null(private$.mvnorm))
+      private$.mvnorm <- private$calculate_mvnorm(state$mvnorm)
     }
   )
 )
@@ -205,3 +249,85 @@ sample_intervention <- function(target, intervention, p, correlations) {
   z <- rnorm(length(target))
   u0 + correlations$mvnorm()[target, intervention] + z < 0
 }
+
+#' Simulate from a conditional multivariate normal distribution.
+#'
+#' Given a multidimensional variable Z which follows a multivariate normal
+#' distribution, this function allows one to draw samples for a subset of Z,
+#' while putting conditions on the values of the rest of Z.
+#'
+#' This effectively allows one to grow a MVN distributed matrix (with columns as
+#' the dimensions and a row per sampled vector), adding new dimensions after the
+#' fact. The existing columns are used as the condition set on the distribution,
+#' and the values returned by this function are used as the new dimensions.
+#'
+#' The maths behind the implementation are described in various online sources:
+#' - https://statproofbook.github.io/P/mvn-cond.html
+#' - https://www.stats.ox.ac.uk/~doucet/doucet_simulationconditionalgaussian.pdf
+#' - https://en.wikipedia.org/wiki/Multivariate_normal_distribution#Conditional_distributions
+#'
+#' @param n the number of samples to simulate
+#' @param mean the mean vector of the distribution, including both given and
+#' dependent variables
+#' @param sigma the variance-covariance matrix of the distribution, including
+#' both given and dependent variables
+#' @param given a matrix of given values used as conditions when simulating the
+#' distribution. The matrix should have \code{n} rows, each one specifying a
+#' different set of values for the given variables.
+#' @param dependent.ind the indices within \code{mean} and \code{sigma} of the
+#' variables to simulate.
+#' @param given.ind the indices within \code{mean} and \code{sigma} of the
+#' variables for which conditions are given. The length of this vector must be
+#' equal to the number of columns of the \code{given} matrix. If empty or NULL,
+#' this function is equivalent to simulating from an unconditional multivariate
+#' normal distribution.
+#' @return a matrix with \code{n} rows and \code{length(dependent.ind)} columns,
+#' containing the simulated value.
+#' @importFrom MASS mvrnorm
+#' @noRd
+rcondmvnorm <- function(n, mean, sigma, given, dependent.ind, given.ind) {
+  stopifnot(length(mean) == nrow(sigma))
+  stopifnot(length(mean) == ncol(sigma))
+  stopifnot(nrow(given) == n)
+  stopifnot(ncol(given) == length(given.ind))
+
+  sigma11 <- sigma[dependent.ind, dependent.ind, drop=FALSE]
+  sigma12 <- sigma[dependent.ind, given.ind, drop=FALSE]
+  sigma21 <- sigma[given.ind, dependent.ind, drop=FALSE]
+  sigma22 <- sigma[given.ind, given.ind, drop=FALSE]
+
+  if (all(sigma22 == 0)) {
+    # This covers two cases: there were no given variables and therefore their
+    # variance-covariance matrix is empty, or there were given variables but
+    # they had a variance of zero. The general formula can't support the latter
+    # case since it tries to invert the matrix, but we can safely ignore the
+    # values since they are all equal to their mean and don't influence the
+    # dependent variables.
+    #
+    # In both cases we revert to a standard MVN with no condition.
+    mvrnorm(n, mean[dependent.ind], sigma11)
+  } else {
+    # Available implementations of the conditional multivariate normal assume
+    # every sample is drawn using the same condition on the given variables.
+    # This is not true in our usecase, where every individual has already had an
+    # independent vector of values drawn for the given variable. We are
+    # effectively drawing from as many different distributions as there are
+    # individuals. Thankfully the same conditional covariance matrix can be
+    # used for all the distributions, only the mean vector needs to be
+    # different. We draw the underlying samples from the MVN at mean 0, and
+    # offset that later on a per-individual basis.
+    #
+    # To work over all the vectors directly they need to be as columns, which
+    # is why we start by transposing `given`. R will recycle the `m` matrix and
+    # `mean` vectors across all the columns. The last step is to transpose the
+    # result back into the expected configuration.
+
+    m <- sigma12 %*% solve(sigma22)
+    residual <- t(given) - mean[given.ind]
+    cond_mu <- t(m %*% residual + mean[dependent.ind])
+    cond_sigma <- sigma11 - m %*% sigma21
+
+    samples <- mvrnorm(n, rep(0, length(dependent.ind)), cond_sigma)
+    samples + cond_mu
+  }
+}
diff --git a/man/CorrelationParameters.Rd b/man/CorrelationParameters.Rd
index 8bb43d70..0aebf688 100644
--- a/man/CorrelationParameters.Rd
+++ b/man/CorrelationParameters.Rd
@@ -2,14 +2,37 @@
 % Please edit documentation in R/correlation.R
 \name{CorrelationParameters}
 \alias{CorrelationParameters}
-\title{Class: Correlation parameters
-Describes an event in the simulation}
+\title{Class: Correlation parameters}
 \description{
 Class: Correlation parameters
-Describes an event in the simulation
 
 Class: Correlation parameters
-Describes an event in the simulation
+}
+\details{
+This class implements functionality that allows interventions to be
+correlated, positively or negatively. By default, interventions are applied
+independently and an individual's probability of receiving two interventions
+(either two separate interventions or two rounds of the same one) is the
+product of the probability of receiving each one.
+
+By setting a positive correlation between two interventions, we can make it
+so that the individuals that receive intervention A are more likely to
+receive intervention B. Conversely, a negative correlation will make it such
+that individuals that receive intervention A are less likely to also receive
+intervention B.
+
+Broadly speaking, the implementation works by assigning at startup a weight
+to each individual and intervention pair, reflecting how likely an individual
+is to receive that intervention. Those weights are derived stochastically
+from the configured correlation parameters.
+
+For a detailed breakdown of the calculations, see Protocol S2 of
+Griffin et al. (2010).
+Derive the mvnorm from the configured correlations.
+
+If a \code{restored_mvnorm} is specified, its columns (corresponding to
+restored interventions) will be re-used as is. Missing columns (for new
+interventions) are derived in accordance with the restored data.
 }
 \section{Methods}{
 \subsection{Public methods}{
diff --git a/tests/testthat/test-correlation.R b/tests/testthat/test-correlation.R
index 882831f8..5653daaa 100644
--- a/tests/testthat/test-correlation.R
+++ b/tests/testthat/test-correlation.R
@@ -134,3 +134,221 @@ test_that('-1 correlation between interventions gives sensible samples', {
     pop * (pev_coverage + mda_coverage),
     tolerance=.1)
 })
+
+test_that('correlation between rounds is preserved when adding interventions', {
+  pop <- 1e6
+  target <- seq(pop)
+
+  pev_coverage_1 <- .2
+  pev_coverage_2 <- .4
+
+  initial <- CorrelationParameters$new(pop, c('pev'))
+  initial$inter_round_rho('pev', 1)
+
+  restored <- CorrelationParameters$new(pop, c('pev', 'mda'))
+  restored$inter_round_rho('pev', 1)
+  restored$inter_round_rho('mda', 1)
+  restored$inter_intervention_rho('pev', 'mda', 1)
+  restored$restore_state(0, initial$save_state())
+
+  round_1 <- sample_intervention(target, 'pev', pev_coverage_1, initial)
+  round_2 <- sample_intervention(target, 'pev', pev_coverage_2, restored)
+
+  expect_equal(sum(round_1), pop * pev_coverage_1, tolerance=.1)
+  expect_equal(sum(round_2), pop * pev_coverage_2, tolerance=.1)
+
+  expect_equal(
+    sum(round_1 & round_2),
+    pop * min(pev_coverage_1, pev_coverage_2),
+    tolerance=.1)
+
+  expect_equal(
+    sum(round_1 | round_2),
+    pop * max(pev_coverage_1, pev_coverage_2),
+    tolerance=.1)
+})
+
+test_that('1 correlation between interventions gives sensible samples when restored', {
+  pop <- 1e6
+  target <- seq(pop)
+
+  pev_coverage <- .2
+  mda_coverage <- .2
+
+  initial <- CorrelationParameters$new(pop, c('pev'))
+  initial$inter_round_rho('pev', 1)
+
+  restored <- CorrelationParameters$new(pop, c('pev', 'mda'))
+  restored$inter_round_rho('pev', 1)
+  restored$inter_round_rho('mda', 1)
+  restored$inter_intervention_rho('pev', 'mda', 1)
+  restored$restore_state(0, initial$save_state())
+
+  pev_sample <- sample_intervention(target, 'pev', pev_coverage, initial)
+  mda_sample <- sample_intervention(target, 'mda', mda_coverage, restored)
+
+  expect_equal(sum(pev_sample), pop * pev_coverage, tolerance=.1)
+  expect_equal(sum(mda_sample), pop * mda_coverage, tolerance=.1)
+
+  expect_equal(
+    sum(pev_sample & mda_sample),
+    pop * min(pev_coverage, mda_coverage),
+    tolerance=.1)
+
+  expect_equal(
+    sum(pev_sample | mda_sample),
+    pop * max(pev_coverage, mda_coverage),
+    tolerance=.1)
+})
+
+test_that('0 correlation between interventions gives sensible samples when restored', {
+  pop <- 1e6
+  target <- seq(pop)
+
+  pev_coverage <- .2
+  mda_coverage <- .2
+
+  initial <- CorrelationParameters$new(pop, c('pev'))
+  initial$inter_round_rho('pev', 1)
+
+  restored <- CorrelationParameters$new(pop, c('pev', 'mda'))
+  restored$inter_round_rho('pev', 1)
+  restored$inter_round_rho('mda', 1)
+  restored$inter_intervention_rho('pev', 'mda', 0)
+  restored$restore_state(0, initial$save_state())
+
+  pev_sample <- sample_intervention(target, 'pev', pev_coverage, initial)
+  mda_sample <- sample_intervention(target, 'mda', mda_coverage, restored)
+
+  expect_equal(sum(pev_sample), pop * pev_coverage, tolerance=.1)
+  expect_equal(sum(mda_sample), pop * mda_coverage, tolerance=.1)
+
+  expect_equal(
+    sum(pev_sample & mda_sample),
+    pop * pev_coverage * mda_coverage,
+    tolerance=.1)
+
+  expect_equal(
+    sum(pev_sample | mda_sample),
+    pop * (pev_coverage + mda_coverage - (pev_coverage * mda_coverage)),
+    tolerance=.1)
+})
+
+test_that('-1 correlation between interventions gives sensible samples when restored', {
+  pop <- 1e6
+  target <- seq(pop)
+  pev_coverage <- .2
+  mda_coverage <- .2
+
+  initial <- CorrelationParameters$new(pop, c('pev'))
+  initial$inter_round_rho('pev', 1)
+
+  restored <- CorrelationParameters$new(pop, c('pev', 'mda'))
+  restored$inter_round_rho('pev', 1)
+  restored$inter_round_rho('mda', 1)
+  restored$inter_intervention_rho('pev', 'mda', -1)
+  restored$restore_state(0, initial$save_state())
+
+  pev_sample <- sample_intervention(target, 'pev', pev_coverage, initial)
+  mda_sample <- sample_intervention(target, 'mda', mda_coverage, restored)
+
+  expect_equal(sum(pev_sample), pop * pev_coverage, tolerance=.1)
+  expect_equal(sum(mda_sample), pop * mda_coverage, tolerance=.1)
+
+  expect_equal(sum(pev_sample & mda_sample), 0, tolerance=.1)
+  expect_equal(
+    sum(pev_sample | mda_sample),
+    pop * (pev_coverage + mda_coverage),
+    tolerance=.1)
+})
+
+test_that("rcondmvnorm has correct distribution", {
+  set.seed(123)
+
+  pop <- 1e6
+
+  # These are completely arbitrary values. The statistics from the simulated
+  # sample will get compared back to this.
+  means <- c(-5, 7, 0, 0.3)
+  variance <- c(0.3, 0.6, 0.9, 0.2)
+  correlation <- matrix(0, ncol=4, nrow=4)
+  correlation[1,2] <- correlation[2,1] <- -0.6
+  correlation[1,3] <- correlation[3,1] <- 0.4
+  correlation[1,4] <- correlation[4,1] <- 1
+  correlation[2,3] <- correlation[3,2] <- 0
+  correlation[2,4] <- correlation[4,2] <- 0.1
+  correlation[3,4] <- correlation[4,3] <- 0.5
+
+  covariance <- outer(variance, variance) * correlation
+  diag(covariance) <- variance
+
+  # These are indices of variables that get simulated together.
+  wy.ind <- c(1, 3)
+  xz.ind <- c(2, 4)
+
+  # Simulate from the MVN for a subset of the dimensions. We intentionally pass
+  # in a subset of the mean and covariance matrices, as the rest of the
+  # parameters are not needed and may not be known. `dependent.ind` is relative
+  # to the subsetted mean and covariance, and therefore is set the first two
+  # indices as opposed to `wy.ind`.
+  wy <- rcondmvnorm(pop, means[wy.ind], covariance[wy.ind, wy.ind], given=NULL,
+                    dependent.ind=c(1,2), given.ind=NULL)
+  expect_equal(dim(wy), c(pop, 2))
+  expect_equal(apply(wy, 2, mean), means[wy.ind], tolerance=0.001)
+  expect_equal(cov(wy), covariance[wy.ind, wy.ind], tolerance=0.001)
+
+  # Now simulate some more, but conditional on the existing values.
+  # The call to rcondmvnorm needs all the mean and covariance parameters,
+  # including the ones that have already been simulated.
+  xz <- rcondmvnorm(pop, means, covariance, given=wy,
+                    dependent.ind=xz.ind, given.ind=wy.ind)
+  expect_equal(dim(xz), c(pop, 2))
+  expect_equal(apply(xz, 2, mean), means[xz.ind], tolerance=0.001)
+  expect_equal(cov(xz), covariance[xz.ind, xz.ind], tolerance=0.001)
+
+  # Stitch the variables back together and make sure the covariance across the
+  # separately simulated values match the expected one.
+  values <- cbind(wy[,1], xz[,1], wy[,2], xz[,2])
+  expect_equal(apply(values, 2, mean), means, tolerance=0.001)
+  expect_equal(cov(values), covariance, tolerance=0.001)
+})
+
+# This would usually be considered an uninteresting edge case, since null
+# variance just yields a constant vector. The way we use the function in the
+# simulation though, a null variance is actually the default and most common
+# case, as it is used where the different rounds of an intervention have no
+# correlation.
+#
+# We need to make sure the other variables are still correctly simulated.
+test_that("rcondmvnorm allows null variance of given variables", {
+  set.seed(123)
+
+  pop <- 1e6
+
+  means <- c(-5, 7, 0)
+  variance <- c(0.3, 0, 0.9)
+  correlation <- matrix(0, ncol=3, nrow=3)
+  correlation[1,3] <- 0.4
+  correlation[3,1] <- 0.4
+
+  covariance <- outer(variance, variance) * correlation
+  diag(covariance) <- variance
+
+  # These are indices of variables that get simulated together.
+  y.ind <- 2
+  xz.ind <- c(1,3)
+
+  # Simulate one dimension. Because its variance was null, the result is a
+  # constant vector.
+  y <- rcondmvnorm(pop, means, covariance, given=NULL,
+                   dependent.ind=y.ind, given.ind=NULL)
+  expect_equal(as.vector(y), rep(means[y.ind], pop))
+
+  # Simulate the rest. The original dimension has no influence on the result.
+  xz <- rcondmvnorm(pop, means, covariance, given=y,
+                    dependent.ind=xz.ind, given.ind=y.ind)
+
+  xyz <- cbind(xz[,1], y, xz[,2])
+  expect_equal(apply(xyz, 2, mean), means, tolerance=0.001)
+  expect_equal(cov(xyz), covariance, tolerance=0.001)
+})
diff --git a/tests/testthat/test-resume.R b/tests/testthat/test-resume.R
index 4d4a572d..1fb7ff87 100644
--- a/tests/testthat/test-resume.R
+++ b/tests/testthat/test-resume.R
@@ -11,20 +11,39 @@ test_resume <- function(
   warmup_parameters = parameters,
   warmup_timesteps = 50
   ) {
+
+  # This function is only used with null correlations. However a null
+  # correlation involves sampling random numbers during initialization, which
+  # disrupts the global RNG and affects the reproducibility if the size of the
+  # matrix is not always the same.
+  #
+  # We use a single correlation object, that we initialize eagerly, such that
+  # the simulation can run undisturbed.
+  correlations <- get_correlation_parameters(parameters)
+  correlations$mvnorm()
+
   set.seed(123)
-  uninterrupted_run <- run_simulation(timesteps, parameters=parameters)
+  uninterrupted_run <- run_simulation(
+    timesteps,
+    parameters = parameters,
+    correlations = correlations)
 
   set.seed(123)
-  first_phase <- run_resumable_simulation(warmup_timesteps, warmup_parameters)
+  first_phase <- run_resumable_simulation(
+    warmup_timesteps,
+    warmup_parameters,
+    correlations = correlations)
   second_phase <- run_resumable_simulation(
     timesteps,
     parameters,
-    initial_state=first_phase$state,
-    restore_random_state=TRUE)
+    initial_state = first_phase$state,
+    restore_random_state = TRUE)
 
   expect_equal(nrow(first_phase$data), warmup_timesteps)
   expect_equal(nrow(second_phase$data), timesteps - warmup_timesteps)
 
+  # The order of columns isn't always identical, hence why mapequal needs to be
+  # used.
   expect_mapequal(
     second_phase$data,
     uninterrupted_run[-(1:warmup_timesteps),])
@@ -159,3 +178,126 @@ test_that("Bednets intervention can be added when resuming", {
     parameters = base %>% set_default_bednets(c(25, 100)))
   expect_equal(data[diff(data$n_use_net) > 0, "timestep"], 100)
 })
+
+test_that("Correlations can be set when resuming with new interventions", {
+  set.seed(123)
+
+  # When adding a new intervention with a non-zero correlation, we cannot
+  # ensure that an uninterrupted run matches the stopped-and-resumed simulation
+  # exactly, as the correlation matrix ends up being randomly sampled in a
+  # different order. This stops us from using the `test_resume` used throughout
+  # the rest of this file. Instead we'll only do stopped-and-resumed simulations
+  #' and check its behaviour.
+  #
+  # We first do a warmup phase with only TBV enabled. We then resume that
+  # simulation three times, each time with MDA enabled. Each time we resume the
+  # simulation, we set a different correlation parameter between the TBV and
+  # MDA interventions, with values -1, 0 and 1.
+  #
+  # We look at the output data and confirm that the correlation worked as
+  # expected. For this we need not only how many people got each intervention,
+  # but also how many received both and how many received at least one. This is
+  # not normally exposed, so we add an extra process to render these values.
+  #
+  # For simplicity, for each intervention, we remove any selection process other
+  # than overall coverage, such as age range (set to 0-200years) and drug
+  # efficacy (set to 100%).
+  #
+  # We need a large population to make the statistical assertions succeed. We'll
+  # only simulate 3 timesteps to keep execution time down: one timestep for
+  # warmup during which TBV takes place, one in which MDA takes place and one
+  # final timestep to collect the updated variables.
+  population <- 10000
+  tbv_coverage <- 0.2
+  mda_coverage <- 0.4
+
+  warmup_parameters <- get_parameters(overrides=list(human_population=population)) %>%
+    set_tbv(
+      timesteps=1,
+      coverage=tbv_coverage,
+      ages=0:200)
+
+  drug <- SP_AQ_params
+  drug[1] <- 1. # Override the drug efficacy to 100%
+  parameters <- warmup_parameters %>%
+    set_drugs(list(drug)) %>%
+    set_mda(
+      drug = 1,
+      timesteps = 2,
+      coverages = mda_coverage,
+      min_ages = 0,
+      max_ages = 200*365)
+
+  create_processes_stub <- function(renderer, variables, events, parameters, ...) {
+    p <- function(t) {
+      pop <- parameters$human_population
+      tbv <- variables$tbv_vaccinated$get_index_of(a=-1, b=0)$not()
+      mda <- variables$drug_time$get_index_of(-1)$not()
+
+      renderer$render("total_tbv", tbv$size(), t)
+      renderer$render("total_mda", mda$size(), t)
+      renderer$render("total_tbv_and_mda", tbv$copy()$and(mda)$size(), t)
+      renderer$render("total_tbv_or_mda", tbv$copy()$or(mda)$size(), t)
+    }
+    c(create_processes(renderer, variables, events, parameters, ...), p)
+  }
+
+  mockery::stub(run_resumable_simulation, 'create_processes', create_processes_stub)
+
+  warmup_correlations <- get_correlation_parameters(warmup_parameters)
+  warmup_correlations$inter_round_rho('tbv', 1)
+
+  warmup <- run_resumable_simulation(1,
+    parameters=warmup_parameters,
+    correlations=warmup_correlations)
+
+  zero_correlation <- get_correlation_parameters(parameters)
+  zero_correlation$inter_round_rho('tbv', 1)
+  zero_correlation$inter_round_rho('mda', 1)
+
+  positive_correlation <- get_correlation_parameters(parameters)
+  positive_correlation$inter_round_rho('tbv', 1)
+  positive_correlation$inter_round_rho('mda', 1)
+  positive_correlation$inter_intervention_rho('tbv', 'mda', 1)
+
+  negative_correlation <- get_correlation_parameters(parameters)
+  negative_correlation$inter_round_rho('tbv', 1)
+  negative_correlation$inter_round_rho('mda', 1)
+  negative_correlation$inter_intervention_rho('tbv', 'mda', -1)
+
+  data <- run_resumable_simulation(
+    3,
+    initial_state=warmup$state,
+    parameters=parameters,
+    correlations=zero_correlation)$data %>% tail(1)
+  expect_equal(data$total_tbv, population * tbv_coverage, tolerance = 0.1)
+  expect_equal(data$total_mda, population * mda_coverage, tolerance = 0.1)
+  expect_equal(
+    data$total_tbv_and_mda,
+    population * (tbv_coverage * mda_coverage),
+    tolerance = 0.1)
+  expect_equal(
+    data$total_tbv_or_mda,
+    population * (tbv_coverage + mda_coverage - tbv_coverage * mda_coverage),
+    tolerance = 0.1)
+
+  data <- run_resumable_simulation(
+    3,
+    initial_state=warmup$state,
+    parameters=parameters,
+    correlations=positive_correlation)$data %>% tail(1)
+  expect_equal(data$total_tbv, population * tbv_coverage, tolerance = 0.1)
+  expect_equal(data$total_mda, population * mda_coverage, tolerance = 0.1)
+  expect_equal(data$total_tbv_and_mda, min(data$total_tbv, data$total_mda))
+  expect_equal(data$total_tbv_or_mda, max(data$total_tbv, data$total_mda))
+
+  data <- run_resumable_simulation(
+    3,
+    initial_state=warmup$state,
+    parameters=parameters,
+    correlations=negative_correlation)$data %>% tail(1)
+  expect_equal(data$total_tbv, population * tbv_coverage, tolerance = 0.1)
+  expect_equal(data$total_mda, population * mda_coverage, tolerance = 0.1)
+  expect_equal(data$total_tbv_and_mda, 0)
+  expect_equal(data$total_tbv_or_mda, data$total_tbv + data$total_mda)
+})