---
title: "SIR models"
output: rmarkdown::html_vignette
vignette: >
  %\VignetteIndexEntry{sir_models}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

```{r, include = FALSE}
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](https://mrc-ide.github.io/odin/articles/discrete.html) 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`](https://mrc-ide.github.io/odin.dust/)

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

```{r odin_sir, echo = FALSE, results = "asis"}
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`](https://mrc-ide.github.io/odin.dust/). Save the above code as a file named `sir.R`. File names must not contain special characters.

```{r}
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:

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

to help automate the compilation of the model.

# Running the SIR model with [`dust`](https://mrc-ide.github.io/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`](https://mrc-ide.github.io/mcstate/) package to allow parameter inference and forecasting.

[`new()`](https://mrc-ide.github.io/dust/reference/dust.html#method-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:

```{r}
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$):

```{r}
sir_model$state()
```

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

```{r}
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:

```{r}
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:

```{r fig.height=5, fig.width=7}
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:

```r
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:

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

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


```{r odin_age, echo = FALSE, results = "asis"}
r_output(readLines(file.path("sirage.R")))
```

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

```{r}
gen_age <- odin.dust::odin_dust("sirage.R")
```

We can generate an age-structured contact matrix based on the [POLYMOD](https://journals.plos.org/plosmedicine/article?id=10.1371/journal.pmed.0050074) survey by using the [socialmixr](https://github.com/sbfnk/socialmixr) package:

```{r}
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.

```{r, fig.height=3.5, fig.width=7}
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](https://doi.org/10.1016/j.idm.2017.06.002) and [Towers and Feng](https://dx.doi.org/10.1016/j.mbs.2012.07.007).

Many other methods are available on the dust object to support inference, but typically you may want to use the [`mcstate`](https://mrc-ide.github.io/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`](https://mrc-ide.github.io/dust/).