Merge pull request #115 from mrc-ide/mrc-6055

Support for restoring models after deserialisation

DESCRIPTION

@@ -1,6 +1,6 @@
 Package: monty
 Title: Monte Carlo Models
-Version: 0.3.4
+Version: 0.3.5
 Authors@R: c(person("Rich", "FitzJohn", role = c("aut", "cre"),
                     email = "rich.fitzjohn@gmail.com"),
              person("Wes", "Hinsley", role = "aut"),

R/combine.R

@@ -119,6 +119,8 @@ monty_model_combine <- function(a, b, properties = NULL,
     parts, properties)
   observer <- model_combine_observer(
     parts, parameters, properties, call)
+  restore <- model_combine_restore(
+    parts)
 
   data <- list(
     parts = unname(parts),
@@ -133,6 +135,7 @@ monty_model_combine <- function(a, b, properties = NULL,
          gradient = gradient,
          get_rng_state = stochastic$get_rng_state,
          set_rng_state = stochastic$set_rng_state,
+         restore = restore,
          observer = observer,
          direct_sample = direct_sample),
     properties)
@@ -455,3 +458,13 @@ model_combine_allow_multiple_parameters <- function(parts, properties,
           "not supported by both of your models"),
     call = call)
 }
+
+
+model_combine_restore <- function(parts) {
+  a <- parts[[1]]
+  b <- parts[[2]]
+  function() {
+    a$restore()
+    b$restore()
+  }
+}

R/model.R

@@ -199,6 +199,7 @@ monty_model <- function(model, properties = NULL) {
   observer <- validate_model_observer(model, properties, call)
   rng_state <- validate_model_rng_state(model, properties, call)
   parameter_groups <- validate_model_parameter_groups(model, properties, call)
+  restore <- validate_model_restore(model, properties, call)
 
   ## Update properties based on what we found:
   properties$has_gradient <- !is.null(gradient)
@@ -217,6 +218,7 @@ monty_model <- function(model, properties = NULL) {
               gradient = gradient,
               direct_sample = direct_sample,
               observer = observer,
+              restore = restore,
               rng_state = rng_state,
               properties = properties)
   class(ret) <- "monty_model"
@@ -539,6 +541,16 @@ validate_model_parameter_groups <- function(model, properties, call) {
 }
 
 
+validate_model_restore <- function(model, propertioes, call) {
+  restore <- model$restore
+  if (is.function(restore)) {
+    restore
+  } else {
+    function() {}
+  }
+}
+
+
 require_monty_model <- function(model, arg = deparse(substitute(model)),
                                 call = parent.frame()) {
   if (!inherits(model, "monty_model")) {

R/rng.R

@@ -464,10 +464,15 @@ monty_rng <- R6::R6Class(
 
     ##' @description
     ##' Returns the state of the random number stream. This returns a
-    ##' raw vector of length 32 * n_streams. It is primarily intended for
-    ##' debugging as one cannot (yet) initialise a monty_rng object with this
-    ##' state.
+    ##' raw vector of length 32 * `n_streams`.
     state = function() {
       monty_rng_state(private$ptr)
+    },
+
+    ##' @description
+    ##' Sets the state of the random number stream.
+    ##' @param state Raw vector of state, with length 32 * `n_streams`.
+    set_state = function(state) {
+      monty_rng_set_state(private$ptr, state)
     }
   ))

R/sample-manual.R

@@ -126,6 +126,12 @@ monty_sample_manual_run <- function(chain_id, path, progress = NULL) {
   restart <- inputs$restart
   is_continue <- is.list(restart)
 
+  if (is_continue) {
+    restart$model$restore()
+  } else {
+    inputs$model$restore()
+  }
+
   pb <- progress_bar(n_chains, steps$total, progress,
                      show_overall = FALSE, single_chain = TRUE)
 

inst/include/monty/random/prng.hpp

@@ -81,6 +81,11 @@ class prng {
     return state_[i];
   }
 
+  /// Return total number of elements in state
+  size_t state_size() const {
+    return size() * rng_state::size();
+  }
+
   /// Convert the random number state of all generators into a single
   /// vector. This can be used to save the state to restore using
   /// `import_state()`