In mrc-ide/odin.dust: Compile Odin to Dust

knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>"
)
lang_output <- function(x, lang) {
  cat(c(sprintf("```%s", lang), x, "```"), sep = "\n")
}
cc_output <- function(x) lang_output(x, "cc")
r_output <- function(x) lang_output(x, "r")
plain_output <- function(x) lang_output(x, "plain")

Stochastic SIR model definition

A simple definition of the SIR model, as given in the odin documentation is: $$\begin{align} \frac{dS}{dt} &= -\beta \frac{SI}{N} \ \frac{dI}{dt} &= \beta \frac{SI}{N} - \gamma I \ \frac{dR}{dt} &= \gamma I \ \end{align}$$ $S$ is the number of susceptibles, $I$ is the number of infected and $R$ is the number recovered; the total population size $N = S + I + R$ is constant. $\beta$ is the infection rate, $\gamma$ is the recovery rate.

Discretising this model in time steps of width $dt$ gives the following update equations for each time step:

$$\begin{align} S_{t+1} &= S_t - n_{SI} \ I_{t+1} &= I_t + n_{SI} - n_{IR} \ R_{t+1} &= R_t + n_{IR} \end{align}$$

where $$\begin{align} n_{SI} &\sim B(S, 1 - e^{-\beta \frac{I}{N} \cdot dt}) \ n_{IR} &\sim B(I, 1 - e^{-\gamma \cdot dt}) \end{align}$$

Implementing the SIR model using odin.dust

The above equations can straightforwardly be written out using the odin DSL:

r_output(readLines(file.path("sir.R")))

This is converted to a C++ dust model, and compiled into a library in a single step, using odin.dust. Save the above code as a file named sir.R. File names must not contain special characters.

library(odin.dust)
gen_sir <- odin.dust::odin_dust("sir.R")

Saving a model into a package

If you want to distribute your model in an R package, rather than building it locally, you will want to use the odin_dust_package() function instead. This will write the transpiled C++ dust code into src and its R interface in R/dust.R. Package users are not required to regenerate the dust code, and their compiler will build the library when they install the package. You may also wish to add a file R/zzz.R to your package myrpackage with the following lines:

##' @useDynLib myrpackage, .registration = TRUE
NULL

to help automate the compilation of the model.

Running the SIR model with dust

Now we can use the new method on the generator to make dust objects. dust can be driven directly from R, and also interfaces with the mcstate package to allow parameter inference and forecasting.

new() takes the data needed to run the model i.e. a list with any parameters defined as user in the odin code above, the value of the initial time step $t_0$, and the number of particles, each for now can simply be thought of as an independent stochastic realisation of the model, but in the next time step will be used when inferring model parameters.

Additional arguments include the number of threads to parallelise the particles over, and the seed for the random number generator. The seed must be an integer, and using the same seed will ensure reproducible results for all particles. To use this to directly create a new dust object with 10 particles, run using 4 threads:

sir_model <- gen_sir$new(pars = list(dt = 1,
                                     S_ini = 1000,
                                     I_ini = 10,
                                     beta = 0.2,
                                     gamma = 0.1),
                         time = 1,
                         n_particles = 10L,
                         n_threads = 4L,
                         seed = 1L)

The initial state is ten particles wide, four states deep ($t$, $S$, $I$, $R$):

sir_model$state()

Run the particles (repeats) forward 10 time steps of length $dt$, followed by another 10 time steps:

sir_model$run(10)
sir_model$run(20)

We can change the parameters, say by increasing the infection rate and the population size, by reinitalising the model with reset(). We will also use a smaller time step to calculate multiple transitions per unit time:

dt <- 0.25
n_particles <- 10L
p_new <- list(dt = dt, S_ini = 2000, I_ini = 10, beta = 0.4, gamma = 0.1)
sir_model$update_state(pars = p_new, time = 0)
sir_model$state()

Let's run this epidemic forward, and plot the trajectories:

n_times <- 200
x <- array(NA, dim = c(sir_model$info()$len, n_particles, n_times))

for (t in seq_len(n_times)) {
  x[ , , t] <- sir_model$run(t)
}
time <- x[1, 1, ]
x <- x[-1, , ]

par(mar = c(4.1, 5.1, 0.5, 0.5), las = 1)
cols <- c(S = "#8c8cd9", I = "#cc0044", R = "#999966")
matplot(time, t(x[1, , ]), type = "l",
         xlab = "Time", ylab = "Number of individuals",
         col = cols[["S"]], lty = 1, ylim = range(x))
matlines(time, t(x[2, , ]), col = cols[["I"]], lty = 1)
matlines(time, t(x[3, , ]), col = cols[["R"]], lty = 1)
legend("left", lwd = 1, col = cols, legend = names(cols), bty = "n")

Adding age structure to the model

Adding age structure to the model consists of the following steps, which turn variables into arrays:

• Define the number of age categories as a user parameter N_age.
• Add age structure to each compartment, by adding square brackets to the lvalues.
• Modify the rvalues to use quantities from the appropriate compartment, by adding an i index to the rvalues. These will automatically be looped over.
• Where an age compartment needs to be reduced into a single compartment/variable, use sum (or another array function from https://mrc-ide.github.io/odin/articles/functions.html as appropriate)
• Define the dimensions of all arrays, for example by setting dim(S) <- N_age.

This would simply give N_age independent processes equivalent to the first model, scaled by the size of the population in each age category. To actually make this useful, you likely want to add some interaction or transitions between the compartments. An example of this would be to add an age-specific contact matrix, demonstrated below, which defines a different force of infection $\lambda$ for each age group. This is calculated by $$\lambda_i = \frac{\beta}{N} \cdot \sum_{j=1}^{N_{\mathrm{age}}} I_j m_{ij}$$ In the odin code:

m[, ] <- user() # age-structured contact matrix
s_ij[, ] <- m[i, j] * I[i]
lambda[] <- beta / N * sum(s_ij[i, ])

The probability of infection of a susceptible is then indexed by this force of infection:

p_SI[] <- 1 - exp(-lambda[i] * dt)

Putting this all together, the age structured SIR model is as follows:

r_output(readLines(file.path("sirage.R")))

As before, save the file, and use odin.dust to compile the model:

gen_age <- odin.dust::odin_dust("sirage.R")

We can generate an age-structured contact matrix based on the POLYMOD survey by using the socialmixr package:

data(polymod, package = "socialmixr")

## In this example we use 8 age bands of 10 years from 0-70 years old.
age.limits = seq(0, 70, 10)

## Get the contact matrix from socialmixr
contact <- socialmixr::contact_matrix(
  survey = polymod,
  countries = "United Kingdom",
  age.limits = age.limits,
  symmetric = TRUE)

## Transform the matrix to the (symetrical) transmission matrix
## rather than the contact matrix. This transmission matrix is
## weighted by the population in each age band.
transmission <- contact$matrix /
  rep(contact$demography$population, each = ncol(contact$matrix))
transmission

This can be given as a parameter to the model by adding the argument pars = list(m = transmission) when using the new() method on the gen_age generator.

N_age <- length(age.limits)
n_particles <- 5L
dt <- 0.25
model <- gen_age$new(pars = list(dt = dt,
                                 S_ini = contact$demography$population,
                                 I_ini = c(0, 10, 0, 0, 0, 0, 0, 0),
                                 beta = 0.2 / 12.11,
                                 gamma = 0.1,
                                 m = transmission,
                                 N_age = N_age),
                     time = 1,
                     n_particles = n_particles,
                     n_threads = 1L,
                     seed = 1L)

# Define how long the model runs for, number of time steps
n_times <- 1000

# Create an array to contain outputs after looping the model.
# Array contains 35 rows = Total S, I, R (3), and
# in each age compartment (24) as well as the cumulative incidence (8)
x <- array(NA, dim = c(model$info()$len, n_particles, n_times))

# For loop to run the model iteratively
for (t in seq_len(n_times)) {
  x[ , , t] <- model$run(t)
}
time <- x[1, 1, ]
x <- x[-1, , ]
# Plotting the trajectories
par(mfrow = c(2,4), oma=c(2,3,0,0))
for (i in 1:N_age) {
  par(mar = c(3, 4, 2, 0.5))
  cols <- c(S = "#8c8cd9", I = "#cc0044", R = "#999966")
  matplot(time, t(x[i + 3,, ]), type = "l", # Offset to access numbers in age compartment
          xlab = "", ylab = "", yaxt="none", main = paste0("Age ", contact$demography$age.group[i]),
          col = cols[["S"]], lty = 1, ylim=range(x[-1:-3,,]))
  matlines(time, t(x[i + 3 + N_age, , ]), col = cols[["I"]], lty = 1)
  matlines(time, t(x[i + 3 + 2*N_age, , ]), col = cols[["R"]], lty = 1)
  legend("right", lwd = 1, col = cols, legend = names(cols), bty = "n")
  axis(2, las =2)
}
mtext("Number of individuals", side=2,line=1, outer=T)
mtext("Time", side = 1, line = 0, outer =T)

Note that a scale factor is introduced for beta so that $R_0$ remains at 2, offsetting the scaling of the contact matrix m. This is calculated using the next-generation matrix. See Driessche and Towers and Feng.

Many other methods are available on the dust object to support inference, but typically you may want to use the mcstate package instead of calling these directly. However, if you wish to simulate state space models with set parameters, this can be done entirely using the commands above from dust.

mrc-ide/odin.dust documentation built on July 3, 2024, 1:33 p.m.