Support for multinomial draws

[annecori]

At the moment we are using nested binomials to model multinomial draws, e.g.:

# Compute the new infections with multiple strains using nested binomials
n_inf[1] <- rbinom(S, p_inf[1])
n_inf[2:n_strains] <- rbinom(S - sum(n_inf[1:(i - 1)]), p_inf[i])

It would be great if we could write something simpler e.g.:

# Compute the new infections with multiple strains using nested binomials
n_inf[1:n_strains] <- rmultinom(S, p_inf[i])

[richfitz]

## Vector form
n_inf[1] <- rbinom(S, p_inf[1])
n_inf[2:n_strains] <- rbinom(S - sum(n_inf[1:(i - 1)]), p_inf[i])

## Matrix form, assigning into columns
n_inf[1, ] <- rbinom(S[j], p_inf[1, j])
n_inf[2:n_strains, ] <- rbinom(S[j] - sum(n_inf[1:(i - 1)], j), p_inf[i, j])

## Matrix form, assigning into rows
n_inf[, 1] <- rbinom(S[i], p_inf[i, 1]) # could be p_inf[1, i]
n_inf[, 2:n_strains] <- rbinom(S[i] - sum(n_inf[i, 1:(i - 1)]), p_inf[i, j])

[cwhittaker1000]

Whilst playing around with multinomial draws using something based on https://github.com/mrc-ide/odin/issues/213#issue-773660376, I think I came across an issue, centred around the need to update the probability used in each binomial draw once you've assigned to an element of n_inf[]. I think this is what the correct version of the code in https://github.com/mrc-ide/odin/issues/213#issue-773660376 should look like!

total <- rbinom(S, sum(p_inf))
n_inf[1] <- rbinom(total, p_inf[1]/sum(p_inf)) 
n_inf[2:n_strains] <- rbinom(total - sum(n_inf[1:(i - 1)]), p_inf[i]/sum(p_inf[i:n_strains]))

When doing the nested binomials as an approximation for the multinomial, each time you do one of the binomial draws (and assign to an element of n_inf[]) I think you need to update the probability used in the next binomial draw, so that it's relative to the sum of all the remaining probabilities you have left (i.e. p_inf[i:n_strains]).

Note that there's some nuance around whether you're trying to assign the entirety of S into n_inf[]'s elements (which is not what https://github.com/mrc-ide/odin/issues/213#issue-773660376 is doing; or whether you're aiming to assig all of S into n_inf[]'s elements (which more closely resembles a classical multinomial draw). But the above code should be able to handle both instances.

[richfitz]

Thoughts on this with Ed

# simplest cast
x <- Multinomal(n, p)
dim(x) <- 10
dim(p) <- 10
# n is scalar

# Loop over rows of x.  We replace each _column_ with a vector draws from a
# multinomial with parameters n[i] (scalar) and p[i, .] (vector)
x[, .] <- Multinomial(n[i], p[i, .])
dim(x) <- c(nr, nc)
dim(p) <- c(nr, nc)
dim(n) <- nr

# for (i in 1:nrow(x)) {
#   x[i, ] <- Multinomaial(n[i], p[i, ])
# }

x[, ., ] <- Multinomial(n[i, k], p[i, ., k])
dim(x) <- c(n1, n2, n3)
dim(p) <- c(n1, n2, n3)
dim(n) <- c(n1, n3)

## Matrix multiplication
x <- a %*% b

## allowed if dim 2 of a is same as dim 1 of b
## special cases for vectors

## dim of x is dim1 of a dim2 of b, allowed to drop dimension down one rank if
## either is knowably 1

x[???] <- a[???] %*% b[???]

x[, ., ] <- Multinomial(n[i, k], p[i, ., k])
dim(x) <- c(n1, n2, n3)
dim(p) <- c(n1, n2, n3)
dim(n) <- c(n1, n3)

for (c(i, k) in dim(x)) {
  x[i, , k] <- Multinomial(n[i, k], p[i, , k])
}

x <- Multinomial(n[2], p[5, ])

for (c(i, l) in dim(x)) {
  x[i, , , l] <- a %*% b[i, , , l]
}

x <- a %*% b

x[, J, K, ] <- a %*% b[i, J, K, l]
x[, J, K, ] <- a %*% b[K, J, i, l]

x[, J, ] <- Multinomial(n[i, k], p[i, J, k])