Support Matrix Multiplication #291
Description
Benchmarking an R+deSolve code against the equivalent odin code yielded a surprising result:
# A tibble: 2 × 13
expression min median `itr/sec` mem_alloc `gc/sec` n_itr n_gc total_time result memory time
<bch:expr> <bch:tm> <bch:tm> <dbl> <bch:byt> <dbl> <int> <dbl> <bch:tm> <list> <list> <list>
1 odin_c 303ms 305ms 3.28 8.59MB 0 2 0 610ms <NULL> <Rprofmem> <bench_tm>
2 deSolve_r 121ms 121ms 8.23 38.04MB 24.7 1 3 121ms <NULL> <Rprofmem> <bench_tm>
# ℹ 1 more variable: gc <list>
I suspect the culprit is the lack of matrix multiplication in odin (or maybe I don't know how to invoke it).
In the deSolve part:
between_sites <- transmission_rate * S * (foi_matrix %*% I)However, in the odin part I do:
foi[, ] <- transmission_rate * foi_mat[i, j] * I[j]
delta_transmission[] <- transmission_rate * S[i] * I[i] + transmission_rate * S[i] * sum(foi[i,])
deriv(S[]) <- -delta_transmission[i]
deriv(I[]) <- +delta_transmission[i] - delta_recovery[i]Here I've omitted the parts I don't think are necessary.
- Is matrix multiplication supported?
Unfortunately, R's C-facilities have a weird way of doing matrix multiplication, so it might not be supported yet.
I'll work on a minimal testcase to check if this is indeed the problem.
Under details, I have more complete excerpts of my code:
Details
odin::odin({
foi[, ] <- foi_mat[i, j] * I[j]
delta_transmission[] <- transmission_rate * S[i] * I[i] + transmission_rate * S[i] * sum(foi[i,])
delta_recovery[] <- recovery_rate * I[i]
deriv(S[]) <- -delta_transmission[i]
deriv(I[]) <- +delta_transmission[i] - delta_recovery[i]
deriv(R[]) <- delta_recovery[i]
source_id <- user()
target_id <- user()
Total[] <- S[i] + I[i] + R[i]
output(source_prevalence) <- I[as.integer(source_id)] / Total[as.integer(source_id)]
output(target_prevalence) <- I[as.integer(target_id)] / Total[as.integer(target_id)]
transmission_rate <- user(0.05)
recovery_rate <- user(0.01)
S0[] <- user()
I0[] <- user()
foi_mat[,] <- user()
initial(S[]) <- S0[i]
initial(I[]) <- I0[i]
initial(R[]) <- Total[i] - S0[i] - I0[i]
dim(S0) <- user()
dim(S) <- N
dim(I) <- N
dim(R) <- N
dim(I0) <- N
dim(foi_mat) <- c(N, N)
dim(foi) <- c(N, N)
dim(Total) <- N
dim(delta_transmission) <- N
dim(delta_recovery) <- N
N <- length(S0)
},
verbose = TRUE, validate = TRUE, target = "c", pretty = TRUE,
skip_cache = FALSE) ->
model_generator
model <-
model_generator$new(S0 = site_S,
I0 = site_I,
foi_mat = foi_matrix,
source_id = as.integer(source_site_id),
target_id = as.integer(target_site_id))
model$set_user(transmission_rate = 0.05, recovery_rate = 0.01)Benchmarking:
bench::mark(
odin_c = model$run(0:100),
deSolve_r = {
site_R <- site_S
site_R[] <- 0
deSolve::ode(
y = c(S = site_S, I = site_I, R = site_R),
times = 0:100,
func = function(time, state, parms) {
with(parms, {
S <- state[1:N]
I <- state[(N + 1):(2 * N)]
R <- state[(2 * N + 1):(3 * N)]
Total <- S + I + R
between_sites <- transmission_rate * S * (foi_matrix %*% I)
source_prevalence <- I[[source_id]] / Total[[source_id]]
target_prevalence <- I[[target_id]] / Total[[target_id]]
list(c(
dS = -transmission_rate * S * I - between_sites,
dI = +transmission_rate * S * I + between_sites - recovery_rate * I,
dR = recovery_rate * I
),
source_prevalence = source_prevalence,
target_prevalence = target_prevalence)
})
},
parms = list(
transmission_rate = 0.05,
recovery_rate = 0.01,
source_id = as.integer(source_site_id),
target_id = as.integer(target_site_id),
N = length(site_S)
)
)
},
check = FALSE
) %>%
print()