Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Expand arrays #8

Merged
merged 6 commits into from
Nov 8, 2024
Merged

Expand arrays #8

Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
fb80566
add vaccination example for arrays
edknock Nov 7, 2024
57ef95e
more explanation of array syntax
edknock Nov 7, 2024
ed156d2
some formatting
edknock Nov 7, 2024
f4a3412
Apply Rich's suggestions
edknock Nov 7, 2024
65e89e8
Merge branch 'main' into arrays
edknock Nov 7, 2024
8f8e041
add 2 group example
edknock Nov 8, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
350 changes: 283 additions & 67 deletions arrays.qmd
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,105 +5,321 @@
source("common.R")
```

::: {.callout-warning}
These examples come from the 2023 teaching slides, and have not been turned into actual content yet.
:::
## When scalars are just not enough

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.

The aim of this section is to show you how you use odin's array syntax, via some motivating examples. We may need to break this up into several sections! 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)
```

* Need a good motivating example for an array - previously we used an age structured model
* show example code

```
update(S) <- S - n_SI
update(I) <- I + n_SI - n_IR
update(R) <- R + n_IR
```
## Example: a two-group age structured SIR model

```
update(S[]) <- S[i] - n_SI[i]
update(I[]) <- I[i] + n_SI[i] - n_IR[i]
update(R[]) <- R[i] + n_IR[i]
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_2 <- 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]

# 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

# 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
})
```

The odin code
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];
}
```

* Don't use index variables (`i`, `j`, `k`, etc) on the left hand side
* Use multiple lines to account for boundary conditions
* Can crash the program if out of bounds (we have plans for this)
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

## An age structured SIR model
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 hardcoded:
```r
# 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]
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)
```

# Individual probabilities of transition:
p_SI[] <- 1 - exp(-lambda[i] * dt) # S to I
p_IR <- 1 - exp(-gamma * dt) # I to R
## Summing over arrays

# Force of infection
m <- parameter() # age-structured contact matrix
s_ij[, ] <- m[i, j] * I[j]
lambda[] <- beta * sum(s_ij[i, ])
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:

# Draws from binomial distributions for numbers changing between
# compartments:
n_SI[] <- Binomial(S[i], p_SI[i])
n_IR[] <- Binomial(I[i], p_IR)
```r
m1_tot <- sum(m[, 2])
```

initial(S[]) <- S_ini[i]
initial(I[]) <- I_ini[i]
initial(R[]) <- 0
This partial sum approach is frequently used within implicit loops:

# User defined parameters - default in parentheses:
S_ini <- parameter()
I_ini <- parameter()
beta <- parameter(0.0165)
gamma <- parameter(0.1)
```r
m_col_totals[] <- sum(m[, i])
```

# Dimensions of arrays
N_age <- parameter()
dim(S_ini) <- N_age
dim(I_ini) <- N_age
dim(S) <- N_age
dim(I) <- N_age
dim(R) <- N_age
dim(n_SI) <- N_age
dim(n_IR) <- N_age
dim(p_SI) <- N_age
dim(m) <- c(N_age, N_age)
dim(s_ij) <- c(N_age, N_age)
dim(lambda) <- N_age
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`.
:::

Relevant changes from the above:

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

# 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

# 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(S0) <- N_age
dim(I0) <- N_age
dim(S) <- N_age
dim(I) <- N_age
dim(R) <- N_age
dim(n_SI) <- N_age
dim(n_IR) <- N_age
dim(p_SI) <- N_age
dim(m) <- c(N_age, N_age)
dim(s_ij) <- c(N_age, N_age)
dim(lambda) <- N_age
Comment on lines +188 to +199
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Wes' fixes here will make this nicer, we'll want to go back and add it

})
```
m <- parameter() # age-structured contact matrix

::: {.callout-note}
With the equations
```r
s_ij[, ] <- m[i, j] * I[j]
lambda[] <- beta * sum(s_ij[i, ])
p_SI[] <- 1 - exp(-lambda[i] * dt)
update(S[]) <- S[i] - n_SI[i]
N_age <- user()
dim(S) <- N_age
dim(m) <- c(N_age, N_age)
```
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 1d and 2d 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 1d arrays and `i` and `j` for 2d 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

Let's take the above model and additionally add some vaccination to it.
edknock marked this conversation as resolved.
Show resolved Hide resolved

```{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]

# 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, ])
Comment on lines +241 to +244
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should spell out in the text that this is a matrix multiplication, and give the appropriate R code (lambda <- beta * (m %*% I) I think) and say that we're looking to simplify this in future?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added a note about this!


# 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

# 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(S0) <- c(N_age, N_vax)
dim(I0) <- c(N_age, N_vax)
dim(S) <- c(N_age, N_vax)
dim(I) <- c(N_age, N_vax)
dim(R) <- c(N_age, N_vax)
dim(n_SI) <- c(N_age, N_vax)
dim(n_IR) <- c(N_age, N_vax)
dim(p_SI) <- c(N_age, N_vax)
dim(m) <- c(N_age, N_age)
dim(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) <- c(N_age, N_vax)
dim(n_S_vax) <- c(N_age, N_vax)
dim(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
```

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