arrays.qmd

# Arrays {#sec-arrays}

```{r}
#| include: false
source("common.R")
```

## 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.

```{r}
library(odin2)
library(dust2)
```


## Example: a two-group age structured SIR model

Let's take the simple stochastic SIR model from @sec-stochastic-sir and add age structure to it by having two groups (children and adults), with heterogeneous mixing between the groups.

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

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

becomes (approximately)

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

## 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:
```r
dim(S) <- 2
```
or can be generalised so that the dimensions themselves become parameters:
```r
dim(S) <- n_age
n_age <- parameter()
```
You may have array parameters, these can have dimensions explicitly declared
```r
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:
```r
dim(m) <- parameter(rank = 2)
```
You can also declare the dimensions of one array to be the same as the dimensions of another:
```r
dim(m) <- c(n_age, n_age)
dim(s_ij) <- dim(m)
```
or collectively declare dimensions of arrays that have the same dimensions:
```r
dim(m, s_ij) <- c(n_age, n_age)
```

## 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:

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

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

```r
m1_tot <- sum(m[, 2])
```

This partial sum approach is frequently used within implicit loops:

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

In our example we calculated the force of infection by summing over the two age groups
```r
lambda[] <- beta * (s_ij[i, 1] + s_ij[i, 2])
```
so we could have used `sum` to write this as
```r
lambda[] <- beta * sum(s_ij[i, ])
```
and this equation has been generalised to any number of age groups now!

::: {.callout-note}
The same syntax applies to other functions that reduce arrays: `min`, `max` and `prod`.
:::


## Example: a generalised age structured SIR model {#sec-stochastic-age}

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

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

::: {.callout-note}
With the equations
```r
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
```r
lambda <- beta * (m %*% I)
```
We are aiming to support matrix multiplication in future to help simplify this code.
:::


## 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.

## Example: an age-structured SIR model with vaccination {#sec-age-vax}

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

::: {.callout-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!
:::


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


## Semantics of random number draws

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

```r
x[] <- Normal(0, 1)
```

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

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


## 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 @sec-age-vax we see this in the equations governing movement from `I` to `R`:

```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](https://en.wikipedia.org/wiki/Erlang_distribution)  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`:

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

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