Merge pull request #15 from mrc-ide/refactor-dim

explain new dim functionality, reorder usage slightly

arrays.qmd

@@ -65,9 +65,16 @@ sir <- odin({
   gamma <- parameter(0.1)
   
   # Dimensions of arrays
-  dim(S0, I0, S, I, R) <- 2
-  dim(n_SI, n_IR, p_SI) <- 2
-  dim(m, s_ij) <- c(2, 2)
+  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
 })
 ```
@@ -109,6 +116,15 @@ or they can also have their dimensions detected when the parameters are passed i
 ```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
 
@@ -194,8 +210,9 @@ sir_age <- odin({
   
   # Dimensions of arrays
   n_age <- parameter()
-  dim(S0, I0, S, I, R) <- n_age
-  dim(n_SI, n_IR, p_SI) <- n_age
+  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
 })
@@ -277,8 +294,9 @@ sir_age_vax <- odin({
   # Dimensions of arrays
   n_age <- parameter()
   n_vax <- parameter()
-  dim(S0, I0, S, I, R) <- c(n_age, n_vax)
-  dim(n_SI, n_IR, p_SI) <- c(n_age, n_vax)
+  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)
@@ -335,9 +353,7 @@ 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) <- c(n_age, n_vax)
-dim(n_IR) <- c(n_age, n_vax)
-
+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`:
@@ -350,14 +366,12 @@ 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) <- c(n_age, n_vax, k_I)
-dim(n_I_progress) <- c(n_age, n_vax, 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]
-dim(I0) <- c(n_age, n_vax, k_I)
 ```