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

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

@@ -39,12 +39,16 @@ EventBase <- R6Class(
 Event <- R6Class(
   'Event',
   inherit = EventBase,
+  private = list(
+    should_restore = FALSE
+  ),
   public = list(
     #' @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()
+      private$should_restore = restore
     },
 
     #' @description Schedule this event to occur in the future.
@@ -74,11 +78,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 (private$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()
+      }
     }
   )
 )