Merge pull request #118 from mrc-ide/mrc-6087-2

Tidy up rng functions, towards a real interface

DESCRIPTION

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

NAMESPACE

@@ -3,13 +3,15 @@
 S3method("+",monty_model)
 S3method(cnd_footer,monty_parse_error)
 S3method(coda::as.mcmc.list,monty_samples)
+S3method(length,monty_rng_state)
 S3method(posterior::as_draws_array,monty_samples)
 S3method(posterior::as_draws_df,monty_samples)
 S3method(print,monty_model)
 S3method(print,monty_model_properties)
 S3method(print,monty_observer)
 S3method(print,monty_packer)
 S3method(print,monty_packer_grouped)
+S3method(print,monty_rng_state)
 S3method(print,monty_runner)
 S3method(print,monty_sampler)
 S3method(print,monty_samples)
@@ -31,10 +33,21 @@ export(monty_model_split)
 export(monty_observer)
 export(monty_packer)
 export(monty_packer_grouped)
+export(monty_random_binomial)
+export(monty_random_exponential_rate)
+export(monty_random_n_binomial)
+export(monty_random_n_exponential_rate)
+export(monty_random_n_real)
+export(monty_random_real)
 export(monty_rng)
+export(monty_rng_create)
 export(monty_rng_distributed_pointer)
 export(monty_rng_distributed_state)
+export(monty_rng_jump)
+export(monty_rng_long_jump)
 export(monty_rng_pointer)
+export(monty_rng_set_state)
+export(monty_rng_state)
 export(monty_runner_callr)
 export(monty_runner_parallel)
 export(monty_runner_serial)

R/random.R

@@ -1,22 +1,263 @@
-monty_random_alloc <- function(n_streams = 1L, seed = NULL,
-                               deterministic = FALSE) {
+##' Create a monty random number generator.  This allows you to sample
+##' random numbers from the same random number algorithms as monty
+##' provides via C++ to [dust2](https://mrc-ide.github.io/dust2), and
+##' which it uses within its samplers and filters.  This function
+##' creates the internal state, and will be passed to actual
+##' generation functions `monty_random_*`, such as
+##' [monty_random_real()].
+##'
+##' Monty's random number generation is very different to that in base
+##' R.  We have the concept of "streams" of random numbers, with a
+##' generator having 1 or many streams.  Each stream is statistically
+##' independent, and can be sampled from simultaneously.  When you use
+##' any of the random number functions from R, you will draw one
+##' number *per stream*.
+##'
+##' Because the random number state can have multiple streams, and we
+##' return a vector, we have a separate set of functions where
+##' multiple numbers are requested *per stream*; these are all
+##' prefixed `monty_random_n_` (e.g., [monty_random_n_real()]).  These
+##' will return a matrix where you have used multiple streams, with
+##' each column representing a stream.  If you have a single stream
+##' and you set `preserve_stream_dimension` to `FALSE` then we will
+##' drop this dimension and return a matrix.
+##'
+##' @title Create a monty random number genreator
+##'
+##' @param n_streams The number of streams to create (see Details)
+##'
+##' @param seed The initial seed of the random number generator.  This
+##'   can be `NULL`, in which case we seed the generator from R's
+##'   random number state (meaning that we respond to `set.seed` as
+##'   one would expect).  Alternatively, you can provide an integer
+##'   here, but this should be used sparingly and primarily for
+##'   testing.
+##'
+##' @param n_threads The number of threads to use, if OpenMP is
+##'   enabled.
+##'
+##' @param deterministic Logical, indicating if we should use
+##'   "deterministic" mode where distributions return their
+##'   expectations and the state is never changed.
+##'
+##' @param preserve_stream_dimension Logical, indicating if the stream
+##'   dimension should be preserved in the case where `n_streams` is 1
+##'   and the multiple-sample functions are used.  Set this to `TRUE`
+##'   to ensure that the rank of the result does not change with the
+##'   number of streams (see Details).
+##'
+##' @return An object of class `monty_rng_state`, which can be
+##'   passed as the `state` argument to random-number producing
+##'   functions, such as [monty_random_real]
+##'
+##' @export
+##' @examples
+##' state <- monty_rng_create()
+##' state
+##'
+##' monty_random_real(state)
+monty_rng_create <- function(n_streams = 1L, seed = NULL,
+                                n_threads = 1L, deterministic = FALSE,
+                                preserve_stream_dimension = FALSE) {
+  assert_scalar_logical(deterministic)
+  assert_scalar_logical(preserve_stream_dimension)
+  preserve_stream_dimension <- preserve_stream_dimension || n_streams > 1
   ptr <- monty_rng_alloc(seed, n_streams, deterministic)
-  ret <- list(ptr = ptr, n_streams = n_streams, deterministic = deterministic)
-  class(ret) <- "monty_random_state"
-  ret
+  attr(ptr, "n_streams") <- n_streams
+  attr(ptr, "n_threads") <- n_threads
+  attr(ptr, "n_deterministic") <- deterministic
+  attr(ptr, "preserve_stream_dimension") <- preserve_stream_dimension
+  class(ptr) <- "monty_rng_state"
+  ptr
 }
 
 
+##' Get and set internal random number state
+##'
+##' @title Get and set random number state
+##'
+##' @param state The random number state, from [monty_rng_create]
+##'
+##' @return A vector of raws
+##' @export
+##' @examples
+##' s1 <- monty_rng_create()
+##' r1 <- monty_rng_state(s1)
+##'
+##' s2 <- monty_rng_create(seed = r1)
+##' identical(r1, monty_rng_state(s2))
+##' monty_random_real(s1)
+##' monty_random_real(s2)
+##'
+##' monty_rng_set_state(r1, s1)
+##' monty_random_real(s1)
+##' monty_random_real(s1)
+##' monty_random_real(s2)
+monty_rng_state <- function(state) {
+  cpp_monty_rng_state(state)
+}
+
+
+##' @param value A vector of raw values, typically the result of
+##'   exporting a random state with `monty_rng_state()`
+##'
+##' @export
+##' @rdname monty_rng_state
+monty_rng_set_state <- function(value, state) {
+  cpp_monty_rng_set_state(state, value)
+}
+
+
+##' Jump random number state.  There are two "lengths" of jumps; a
+##' normal jump and a long jump.  The normal jump is the distance
+##' between streams within a random number state, so if you have a
+##' multi-stream rng this shifts states left.  The long jump is used
+##' to create distributed states.  We will properly explain all this
+##' once the interface stabilises.
+##'
+##' @title Jump random number state
+##'
+##' @param state Either a `monty_rng_state` object (created via
+##'   [monty_rng_create]) or a raw vector suitable for creating
+##'   one.
+##'
+##' @param n The number of jumps to take (integer, 1 or more)
+##'
+##' @return The `monty_rng_state` object (invisibly, modified in
+##'   place) or a raw vector, matching the input argument `state`
+##'   (visibly).
+##'
+##' @export
+monty_rng_jump <- function(state, n = 1) {
+  if (is.raw(state)) {
+    rng <- monty_rng_create(seed = state, n_streams = length(state) %/% 32)
+    monty_rng_state(monty_rng_jump(rng, n))
+  } else {
+    assert_is(state, "monty_rng_state")
+    assert_scalar_size(n, allow_zero = FALSE)
+    cpp_monty_rng_jump(state, n)
+    invisible(state)
+  }
+}
+
+
+##' @export
+##' @rdname monty_rng_jump
+monty_rng_long_jump <- function(state, n = 1) {
+  if (is.raw(state)) {
+    rng <- monty_rng_create(seed = state, n_streams = length(state) %/% 32)
+    monty_rng_state(monty_rng_long_jump(rng, n))
+  } else {
+    assert_is(state, "monty_rng_state")
+    assert_scalar_size(n, allow_zero = FALSE)
+    cpp_monty_rng_long_jump(state, n)
+    invisible(state)
+  }
+}
+
+
+
+##' @export
+print.monty_rng_state <- function(x, ...) {
+  cli::cli_h1("<monty_rng_state>")
+  cli::cli_li("{length(x)} random number stream{?s}")
+  cli::cli_li("{attr(x, 'n_threads')} execution thread{?s}")
+  invisible(x)
+}
+
+
+##' @export
+length.monty_rng_state <- function(x) {
+  attr(x, "n_streams")
+}
+
+
+##' Generate a random number uniformly sampled on the range 0 to 1;
+##' this is the most basic of all random number functions in monty and
+##' all other algorithms are composed from this.  Quite often, you
+##' will want a number on $\[0, 1\]$ (e.g., for a Bernoulli trial), and
+##' this function is the most efficient way of generating one.
+##'
+##' @title Sample from Uniform(0, 1)
+##'
+##' @param state The random number state, from [monty_rng_create]
+##'
+##' @return A vector of random numbers, the same length as the number
+##'   of streams in `state`.
+##'
+##' @export
+##' @examples
+##' state <- monty_rng_create()
+##' monty_random_real(state)
+##' monty_random_n_real(5, state)
 monty_random_real <- function(state) {
-  cpp_monty_random_real(state$ptr)
+  cpp_monty_random_real(state)
+}
+
+
+##' @param n_samples The number of samples to take, **per stream**.
+##'   When using the multiple-sample interface, all other parameters
+##'   are held constant (per stream).
+##'
+##' @export
+##' @rdname monty_random_real
+monty_random_n_real <- function(n_samples, state) {
+  cpp_monty_random_n_real(n_samples, state)
 }
 
 
+##' Sample from the binomial distribution
+##'
+##' @title Sample from binomial distribution
+##'
+##' @param size The number of trials
+##'
+##' @param prob The probability of success on each trial
+##'
+##' @inheritParams monty_random_real
+##' @inherit monty_random_real return
+##'
+##' @export
+##' @examples
+##' state <- monty_rng_create()
+##' monty_random_binomial(10, 0.3, state)
+##' table(monty_random_n_binomial(2000, 10, 0.3, state))
 monty_random_binomial <- function(size, prob, state) {
-  cpp_monty_random_binomial(size, prob, state$ptr)
+  cpp_monty_random_binomial(size, prob, state)
+}
+
+
+##' @export
+##' @rdname monty_random_binomial
+monty_random_n_binomial <- function(n_samples, size, prob, state) {
+  cpp_monty_random_n_binomial(n_samples, size, prob, state)
 }
 
 
+##' Sample from an exponential distribution.  There are two
+##' parameterisations here, one in terms of the rate of the
+##' exponential, and one in terms of the mean (or scale).
+##'
+##' @title Sample from exponential distribution
+##'
+##' @param rate The rate of the exponential
+##'
+##' @inheritParams monty_random_real
+##' @inherit monty_random_real return
+##'
+##' @export
+##' @rdname monty_random_exponential
+##' @examples
+##' state <- monty_rng_create()
+##' monty_random_exponential_rate(0.2, state)
+##' summary(monty_random_n_exponential_rate(2000, 0.2, state))
 monty_random_exponential_rate <- function(rate, state) {
-  cpp_monty_random_exponential_rate(rate, state$ptr)
+  cpp_monty_random_exponential_rate(rate, state)
+}
+
+
+##' @export
+##' @rdname monty_random_exponential
+monty_random_n_exponential_rate <- function(n_samples, rate, state) {
+  cpp_monty_random_n_exponential_rate(n_samples, rate, state)
 }