4  Arrays

4.1 When scalars are just not enough

The aim of this section is to show you how you use odin’s array syntax, via some motivating examples. The examples here are necessarily a bit longer than in the previous sections because we will generally need a few more moving parts.

library(odin2)
library(dust2)

4.2 Example: a two-group age structured SIR model

Let’s take the simple stochastic SIR model from Section 3.2 and add age structure to it by having two groups (children and adults), with heterogeneous mixing between the groups.

sir <- odin({
  
  # Equations for transitions between compartments by age group
  update(S[]) <- S[i] - n_SI[i]
  update(I[]) <- I[i] + n_SI[i] - n_IR[i]
  update(R[]) <- R[i] + n_IR[i]
  update(incidence) <- incidence + n_SI[1] + n_SI[2]
  
  # Individual probabilities of transition:
  
  p_SI[] <- 1 - exp(-lambda[i] * dt) # S to I
  p_IR <- 1 - exp(-gamma * dt) # I to R
  
  # Calculate force of infection
  
  # age-structured contact matrix: m[i, j] is mean number of contacts an
  # individual in group i has with an individual in group j per time unit
  
  m <- parameter()
  
  # here s_ij[i, j] gives the mean number of contacts and individual in group
  # i will have with the currently infectious individuals of group j
  s_ij[, ] <- m[i, j] * I[j]
  
  # lambda[i] is the total force of infection on an individual in group i 
  lambda[] <- beta * (s_ij[i, 1] + s_ij[i, 2])
  
  # Draws from binomial distributions for numbers changing between
  # compartments:
  
  n_SI[] <- Binomial(S[i], p_SI[i])
  n_IR[] <- Binomial(I[i], p_IR)
  
  initial(S[]) <- S0[i]
  initial(I[]) <- I0[i]
  initial(R[]) <- 0
  initial(incidence, zero_every = 1) <- 0
  
  # User defined parameters - default in parentheses:
  S0 <- parameter()
  I0 <- parameter()
  beta <- parameter(0.0165)
  gamma <- parameter(0.1)
  
  # Dimensions of arrays
  dim(S0) <- 2
  dim(I0) <- 2
  dim(S) <- 2
  dim(I) <- 2
  dim(R) <- 2
  dim(n_SI) <- 2
  dim(n_IR) <- 2
  dim(p_SI) <- 2
  dim(m) <- c(2, 2)
  dim(s_ij) <- c(2, 2)
  dim(lambda) <- 2
})

In the odin code above

update(S[]) <- S[i] - n_SI[i]

becomes (approximately)

for (int i = 0; i < S_length; ++i) {
  update_S[i] = S[i] + n_SI[i];
}

so there is an implicit indexing by i on the LHS in that equation of the odin code, and the generated code will then produce a for loop over all values of i.

4.3 Declaring dimensions

We can see in our first example that it is necessary to declare the dimensions of all arrays within your odin code using a dim equation.

Dimensions can be hard-coded:

dim(S) <- 2

or can be generalised so that the dimensions themselves become parameters:

dim(S) <- n_age
n_age <- parameter()

You may have array parameters, these can have dimensions explicitly declared

dim(m) <- c(n_age, n_age)

or they can also have their dimensions detected when the parameters are passed into the system - this still needs a dim equation where you must explicitly state the rank:

dim(m) <- parameter(rank = 2)

You can also declare the dimensions of one array to be the same as the dimensions of another:

dim(m) <- c(n_age, n_age)
dim(s_ij) <- dim(m)

or collectively declare dimensions of arrays that have the same dimensions:

dim(m, s_ij) <- c(n_age, n_age)

4.4 Summing over arrays

Frequently, you will want to take a sum over an array, or part of an array, using sum. To sum over all elements of an array, use sum() with the name of the array you would like to sum over:

dim(x) <- 10
x_tot <- sum(x)

If m is a matrix you can compute the sums over the second column by writing:

m1_tot <- sum(m[, 2])

This partial sum approach is frequently used within implicit loops:

m_col_totals[] <- sum(m[, i])

In our example we calculated the force of infection by summing over the two age groups

lambda[] <- beta * (s_ij[i, 1] + s_ij[i, 2])

so we could have used sum to write this as

lambda[] <- beta * sum(s_ij[i, ])

and this equation has been generalised to any number of age groups now!

Note

The same syntax applies to other functions that reduce arrays: min, max and prod.

4.5 Example: a generalised age structured SIR model

Let’s now take the two group SIR model and generalise it to n_age age groups

sir_age <- odin({
  # Equations for transitions between compartments by age group
  update(S[]) <- S[i] - n_SI[i]
  update(I[]) <- I[i] + n_SI[i] - n_IR[i]
  update(R[]) <- R[i] + n_IR[i]
  update(incidence) <- incidence + sum(n_SI)
  
  # Individual probabilities of transition:
  p_SI[] <- 1 - exp(-lambda[i] * dt) # S to I
  p_IR <- 1 - exp(-gamma * dt) # I to R
  
  # Calculate force of infection
  
  # age-structured contact matrix: m[i, j] is mean number of contacts an
  # individual in group i has with an individual in group j per time unit
  
  m <- parameter()
  
  # here s_ij[i, j] gives the mean number of contacts and individual in group
  # i will have with the currently infectious individuals of group j
  s_ij[, ] <- m[i, j] * I[j]
  
  # lambda[i] is the total force of infection on an individual in group i 
  lambda[] <- beta * sum(s_ij[i, ])
  
  # Draws from binomial distributions for numbers changing between
  # compartments:
  n_SI[] <- Binomial(S[i], p_SI[i])
  n_IR[] <- Binomial(I[i], p_IR)
  
  initial(S[]) <- S0[i]
  initial(I[]) <- I0[i]
  initial(R[]) <- 0
  initial(incidence, zero_every = 1) <- 0
  
  # User defined parameters - default in parentheses:
  S0 <- parameter()
  I0 <- parameter()
  beta <- parameter(0.0165)
  gamma <- parameter(0.1)
  
  # Dimensions of arrays
  n_age <- parameter()
  dim(S, S0, n_SI, p_SI) <- n_age
  dim(I, I0, n_IR) <- n_age
  dim(R) <- n_age
  dim(m, s_ij) <- c(n_age, n_age)
  dim(lambda) <- n_age
})

Note

With the equations

s_ij[, ] <- m[i, j] * I[j]
lambda[] <- beta * sum(s_ij[i, ])

what we are really doing is matrix multiplication, and in R this would be

lambda <- beta * (m %*% I)

We are aiming to support matrix multiplication in future to help simplify this code.

4.6 Indexing

We have seen in the above examples uses of 1-dimensions and 2-dimensional arrays, making use of index variables i and j. Note that while we never use these index variables on the LHS, it is implicit that on the LHS we are indexing by i for 1-dimensional arrays and i and j for 2-dimensional arrays! Use of these index variables on the RHS corresponds to the indexing on the LHS.

Of course you may want to use higher-dimensional arrays! We currently supporting up to 8 dimensions, with index variables i, j, k, l, i5, i6, i7 and i8.

Be careful with array equations! It’s possible that in an array equation you end up going out of bounds with an array used on the RHS. This can crash the program when running.

4.7 Example: an age-structured SIR model with vaccination

Let’s take the above model and additionally add some vaccination to it.

Note

This should not be considered a guide on how to model vaccination with odin, as it merely presents one model for the purposes of illustrating how arrays work in odin - there are many other ways one could model vaccination!

sir_age_vax <- odin({
  # Equations for transitions between compartments by age group
  update(S[, ]) <- new_S[i, j]
  update(I[, ]) <- I[i, j] + n_SI[i, j] - n_IR[i, j]
  update(R[, ]) <- R[i, j] + n_IR[i, j]
  update(incidence) <- incidence + sum(n_SI)
  
  # Individual probabilities of transition:
  p_SI[, ] <- 1 - exp(-rel_susceptibility[j] * lambda[i] * dt) # S to I
  p_IR <- 1 - exp(-gamma * dt) # I to R
  p_vax[, ] <- 1 - exp(-eta[i, j] * dt)
  
  # Force of infection
  m <- parameter() # age-structured contact matrix
  s_ij[, ] <- m[i, j] * sum(I[j, ])
  lambda[] <- beta * sum(s_ij[i, ])
  
  # Draws from binomial distributions for numbers changing between
  # compartments:
  n_SI[, ] <- Binomial(S[i, j], p_SI[i, j])
  n_IR[, ] <- Binomial(I[i, j], p_IR)
  
  # Nested binomial draw for vaccination in S
  # Assume you cannot move vaccine class and get infected in same step
  n_S_vax[, ] <- Binomial(S[i, j] - n_SI[i, j], p_vax[i, j])
  new_S[, 1] <- S[i, j] - n_SI[i, j] - n_S_vax[i, j] + n_S_vax[i, n_vax]
  new_S[, 2:n_vax] <- S[i, j] - n_SI[i, j] - n_S_vax[i, j] + n_S_vax[i, j - 1]
  
  initial(S[, ]) <- S0[i, j]
  initial(I[, ]) <- I0[i, j]
  initial(R[, ]) <- 0
  initial(incidence, zero_every = 1) <- 0
  
  # User defined parameters - default in parentheses:
  S0 <- parameter()
  I0 <- parameter()
  beta <- parameter(0.0165)
  gamma <- parameter(0.1)
  eta <- parameter()
  rel_susceptibility <- parameter()
  
  # Dimensions of arrays
  n_age <- parameter()
  n_vax <- parameter()
  dim(S, S0, n_SI, p_SI) <- c(n_age, n_vax)
  dim(I, I0, n_IR) <- c(n_age, n_vax)
  dim(R) <- c(n_age, n_vax)
  dim(m, s_ij) <- c(n_age, n_age)
  dim(lambda) <- n_age
  dim(eta) <- c(n_age, n_vax)
  dim(rel_susceptibility) <- c(n_vax)
  dim(p_vax, n_S_vax, new_S) <- c(n_age, n_vax)
})

We see we can use multiple lines to deal with boundary conditions:

new_S[, 1] <- S[i, j] - n_SI[i, j] - n_S_vax[i, j] + n_S_vax[i, n_vax]
new_S[, 2:n_vax] <- S[i, j] - n_SI[i, j] - n_S_vax[i, j] + n_S_vax[i, j - 1]

which we could also write as

new_S[, ] <- S[i, j] - n_SI[i, j] - n_S_vax[i, j]
new_S[, 1] <- new_S[i, j] + n_S_vax[i, n_vax]
new_S[, 2:n_vax] <- new_S[i, j] + n_S_vax[i, j - 1]

or another alternative way of writing this would be to use if else

new_S[, ] <- S[i, j] - n_SI[i, j] - n_S_vax[i, j] +
    (if (j == 1) n_S_vax[i, n_vax] else n_S_vax[i, j - 1])

Note that in odin, an if always requires an else!

4.8 Semantics of random number draws

Stochastic functions are called for each element in an array they are assigned to, at each time. So here:

x[] <- Normal(0, 1)

x will be filled with each element having a different draw from a standard normal. In contrast, in:

a <- Normal(0, 1)
x[] <- a

x will be a vector where every element is the same, the result of a single draw from a standard normal.

4.9 Using arrays for Erlang durations

In our examples we have assumed that infectious periods follow a discretised exponential distribution (rounded up to the nearest dt), for instance in Section 4.7 we see this in the equations governing movement from I to R:

p_IR <- 1 - exp(-gamma * dt) # I to R
n_IR[, ] <- Binomial(I[i, j], p_IR)

update(I[, ]) <- I[i, j] + n_SI[i, j] - n_IR[i, j]
update(R[, ]) <- R[i, j] + n_IR[i, j]

dim(I, I0, n_IR) <- c(n_age, n_vax)

However, exponential distributions often do not capture the true distribution of infectious periods or other such delays you may be interested in modelling. Arrays in odin can be used to implement (discretised) Erlang distributions by breaking them down into stages of iid exponential distributions by adding a dimension to an array corresponding to these stages. For instance we could generalise the above to an Erlang with shape parameter k_I (the number of stages) and rate gamma:

k_I <- parameter()
p_I_progress <- 1 - exp(-gamma * dt)
n_I_progress[, , ] <- Binomial(I[i, j, k], p_I_progress)

update(I[, , ]) <- I[i, j, k] - n_I_progress[i, j, k] +
  (if (k == 1) n_SI[i, j] else n_I_progress[i, j, k - 1])
update(R[, ]) <- R[i, j] + n_I_progress[i, j, k_I]
dim(I, I0, n_I_progress) <- c(n_age, n_vax, k_I)

and there would be some other bits we’d need to change to deal with the increased dimensionality of I:

s_ij[, ] <- m[i, j] * sum(I[j, ,])
initial(I[, , ]) <- I0[i, j, k]