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

@@ -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
+  }
+}