Allow events and variables to be added when restoring the simulation.

The simulation restoring feature required the number of events to be
identical to the original simulation's. This has turned out to be too
restrictive for our use cases in malariasimulation, where we want to
resume the simulation with new interventions enabled, which come with
new events.

I had originally hoped that this would be enough, and that we could run
the first phase of the simulation with all the events present, even
though they are unused. However there are interventions whose number of
events depends on the parametrization of it, hence we cannot create a
saved simulation state that fits all possible use cases when resuming.

The solution implemented here is to allow more events and variables to
be introduced when resuming the simulation, on the condition that the
objects are named when passed to `simulation_loop`. Without this
requirement, given more objects than were saved, it would be ambigous
which ones need to be restored and which ones are new.

When a new object is included in the simulation, its `restore_state`
method is called with a NULL argument. This is needed because events,
even new ones, still need to set their internal timestep variable.

Additionally, the list of events and variables can now be structured
using nested lists of objects. This has no impact on the way the
simulation is executed, but it allows for more complicated simulations
to be restored.

For example, this is a simplified example of what the list of events in
malariasimulation might look like:
```
list(
 mda_administer = Event$new(),
 mass_pev_doses = list(
  TargetedEvent$new(n),
  TargetedEvent$new(n),
  TargetedEvent$new(n)
 )
)
```

In this example, the top-level list is names, allowing the
`mass_pev_doses` events to be absent during the warmup simulation and
added only later when restoring. However, because the nested list of
targeted events is not named, more events cannot be added to it.

The names of the methods to save and restore state, together with the
helper functions `save_object_state` and `restore_object_state` are
tweaked and made public to allow this pattern to be reused in
applications that have their own state to save, as demonstrated by the
[stateful.R] file in malariasimulation.

[stateful.R]: https://github.com/mrc-ide/malariasimulation/blob/b3376d33cda29cc5e4d7217fefad4b367f9f9167/R/stateful.R

NAMESPACE

@@ -18,6 +18,8 @@ export(infection_age_process)
 export(multi_probability_bernoulli_process)
 export(multi_probability_multinomial_process)
 export(reschedule_listener)
+export(restore_object_state)
+export(save_object_state)
 export(simulation_loop)
 export(update_category_listener)
 importFrom(R6,R6Class)

R/RcppExports.R

@@ -165,8 +165,8 @@ double_variable_queue_shrink_bitset <- function(variable, index) {
     invisible(.Call(`_individual_double_variable_queue_shrink_bitset`, variable, index))
 }
 
-create_event <- function(restoreable) {
-    .Call(`_individual_create_event`, restoreable)
+create_event <- function() {
+    .Call(`_individual_create_event`)
 }
 
 create_targeted_event <- function(size) {
@@ -181,6 +181,10 @@ event_base_get_timestep <- function(event) {
     .Call(`_individual_event_base_get_timestep`, event)
 }
 
+event_base_set_timestep <- function(event, time) {
+    invisible(.Call(`_individual_event_base_set_timestep`, event, time))
+}
+
 event_base_should_trigger <- function(event) {
     .Call(`_individual_event_base_should_trigger`, event)
 }
@@ -197,8 +201,8 @@ event_checkpoint <- function(event) {
     .Call(`_individual_event_checkpoint`, event)
 }
 
-event_restore <- function(event, time, schedule) {
-    invisible(.Call(`_individual_event_restore`, event, time, schedule))
+event_restore <- function(event, schedule) {
+    invisible(.Call(`_individual_event_restore`, event, schedule))
 }
 
 targeted_event_clear_schedule_vector <- function(event, target) {
@@ -257,8 +261,8 @@ targeted_event_checkpoint <- function(event) {
     .Call(`_individual_targeted_event_checkpoint`, event)
 }
 
-targeted_event_restore <- function(event, time, state) {
-    invisible(.Call(`_individual_targeted_event_restore`, event, time, state))
+targeted_event_restore <- function(event, state) {
+    invisible(.Call(`_individual_targeted_event_restore`, event, state))
 }
 
 process_listener <- function(event, listener) {

R/categorical_variable.R

@@ -96,21 +96,31 @@ CategoricalVariable <- R6Class(
     .update = function() variable_update(self$.variable),
     .resize = function() variable_resize(self$.variable),
 
-    .checkpoint = function() {
+    #' @description save the state of the variable
+    save_state = function() {
       categories <- self$get_categories()
       values <- lapply(categories, function(c) self$get_index_of(c)$to_vector())
       names(values) <- categories
       values
     },
 
-    .restore = function(values) {
-      stopifnot(names(values) == self$get_categories())
-      stopifnot(sum(sapply(values, length)) == categorical_variable_get_size(self$.variable))
+    #' @description restore the variable from a previously saved state.
+    #' @param timestep the timestep at which simulation is resumed. This
+    #' parameter's value is ignored, it only exists to conform to a uniform
+    #' interface with events.
+    #' @param state the previously saved state, as returned by the
+    #' \code{save_state} method. NULL is passed when restoring from a saved
+    #' simulation in which this variable did not exist.
+    restore_state = function(timestep, state) {
+      if (!is.null(state)) {
+        stopifnot(names(state) == self$get_categories())
+        stopifnot(sum(sapply(state, length)) == categorical_variable_get_size(self$.variable))
 
-      for (c in names(values)) {
-        self$queue_update(c, values[[c]])
+        for (c in names(state)) {
+          self$queue_update(c, state[[c]])
+        }
+        self$.update()
       }
-      self$.update()
     }
   )
 )

R/double_variable.R

@@ -147,11 +147,22 @@ DoubleVariable <- R6Class(
     .update = function() variable_update(self$.variable),
     .resize = function() variable_resize(self$.variable),
 
-    .checkpoint = function() self$get_values(),
-    .restore = function(values) {
-      stopifnot(length(values) == variable_get_size(self$.variable))
-      self$queue_update(values)
-      self$.update()
+    #' @description save the state of the variable
+    save_state = function() self$get_values(),
+
+    #' @description restore the variable from a previously saved state.
+    #' @param timestep the timestep at which simulation is resumed. This
+    #' parameter's value is ignored, it only exists to conform to a uniform
+    #' interface with events.
+    #' @param state the previously saved state, as returned by the
+    #' \code{save_state} method. NULL is passed when restoring from a saved
+    #' simulation in which this variable did not exist.
+    restore_state = function(timestep, state) {
+      if (!is.null(state)) {
+        stopifnot(length(state) == variable_get_size(self$.variable))
+        self$queue_update(state)
+        self$.update()
+      }
     }
   )
 )

R/event.R

@@ -40,11 +40,14 @@ Event <- R6Class(
   'Event',
   inherit = EventBase,
   public = list(
+    .should_restore = FALSE,
+
     #' @description Initialise an Event.
-    #' @param restore if true, the schedule of this event is restored when restoring from a saved
-    #' simulation.
+    #' @param restore if true, the schedule of this event is restored when
+    #' restoring from a saved simulation.
     initialize = function(restore = TRUE) {
-      self$.event <- create_event(restore)
+      self$.event <- create_event()
+      self$.should_restore = restore
     },
 
     #' @description Schedule this event to occur in the future.
@@ -74,11 +77,23 @@ Event <- R6Class(
     # NOTE: intentionally empty
     .resize = function() {},
 
-    .checkpoint = function() {
+    #' @description save the state of the event
+    save_state = function() {
       event_checkpoint(self$.event)
     },
-    .restore = function(time, schedule) {
-      event_restore(self$.event, time, schedule)
+
+    #' @description restore the event from a previously saved state.
+    #' If the event was constructed with \code{restore = FALSE}, the state
+    #' argument is ignored.
+    #' @param timestep the timestep at which simulation is resumed.
+    #' @param state the previously saved state, as returned by the
+    #' \code{save_state} method. NULL is passed when restoring from a saved
+    #' simulation in which this variable did not exist.
+    restore_state = function(timestep, state) {
+      event_base_set_timestep(self$.event, timestep)
+      if (self$.should_restore && !is.null(state)) {
+        event_restore(self$.event, state)
+      }
     }
   )
 )

R/integer_variable.R

@@ -190,11 +190,22 @@ IntegerVariable <- R6Class(
     .update = function() variable_update(self$.variable),
     .resize = function() variable_resize(self$.variable),
 
-    .checkpoint = function() self$get_values(),
-    .restore = function(values) {
-      stopifnot(length(values) == variable_get_size(self$.variable))
-      self$queue_update(values)
-      self$.update()
+    #' @description save the state of the variable
+    save_state = function() self$get_values(),
+
+    #' @description restore the variable from a previously saved state.
+    #' @param timestep the timestep at which simulation is resumed. This
+    #' parameter's value is ignored, it only exists to conform to a uniform
+    #' interface with events.
+    #' @param state the previously saved state, as returned by the
+    #' \code{save_state} method. NULL is passed when restoring from a saved
+    #' simulation in which this variable did not exist.
+    restore_state = function(timestep, state) {
+      if (!is.null(state)) {
+        stopifnot(length(state) == variable_get_size(self$.variable))
+        self$queue_update(state)
+        self$.update()
+      }
     }
   )
 )

R/ragged_double.R

@@ -152,11 +152,22 @@ RaggedDouble <- R6Class(
     .update = function() variable_update(self$.variable),
     .resize = function() variable_resize(self$.variable),
 
-    .checkpoint = function() self$get_values(),
-    .restore = function(values) {
-      stopifnot(length(values) == variable_get_size(self$.variable))
-      self$queue_update(values)
-      self$.update()
+    #' @description save the state of the variable
+    save_state = function() self$get_values(),
+
+    #' @description restore the variable from a previously saved state.
+    #' @param timestep the timestep at which simulation is resumed. This
+    #' parameter's value is ignored, it only exists to conform to a uniform
+    #' interface with events.
+    #' @param state the previously saved state, as returned by the
+    #' \code{save_state} method. NULL is passed when restoring from a saved
+    #' simulation in which this variable did not exist.
+    restore_state = function(timestep, state) {
+      if (!is.null(state)) {
+        stopifnot(length(state) == variable_get_size(self$.variable))
+        self$queue_update(state)
+        self$.update()
+      }
     }
   )
 )

R/ragged_integer.R

@@ -152,11 +152,22 @@ RaggedInteger <- R6Class(
     .update = function() variable_update(self$.variable),
     .resize = function() variable_resize(self$.variable),
 
-    .checkpoint = function() self$get_values(),
-    .restore = function(values) {
-      stopifnot(length(values) == variable_get_size(self$.variable))
-      self$queue_update(values)
-      self$.update()
+    #' @description save the state of the variable
+    save_state = function() self$get_values(),
+
+    #' @description restore the variable from a previously saved state.
+    #' @param timestep the timestep at which simulation is resumed. This
+    #' parameter's value is ignored, it only exists to conform to a uniform
+    #' interface with events.
+    #' @param state the previously saved state, as returned by the
+    #' \code{save_state} method. NULL is passed when restoring from a saved
+    #' simulation in which this variable did not exist.
+    restore_state = function(timestep, state) {
+      if (!is.null(state)) {
+        stopifnot(length(state) == variable_get_size(self$.variable))
+        self$queue_update(state)
+        self$.update()
+      }
     }
   )
 )