Arrays of arrays in parameters cannot JIT compile

[fkrauer]

Hi

I have extended my epidemic SIR model to two age groups, which have some contact rates among each other (captured with a mixing matrix), and some transition rates (ageing, demo matrix). Since diffeqr does not accept additional arguments, I have wrapped the two matrices in "p" with an R list. However, diffeqr seems to have a problem with the array or matrix type in R.

The following code:

library(diffeqr)

maxtime <- 3650.0

# Make up some demography data
lifespan <- 80
steps <- 40
alower <- seq(0,lifespan-steps,steps) # lower boundary of age groups
alength <- c(diff(alower), steps) # duration in years in age group
ngroups <- length(alower)
demo <- matrix(data=NA, nrow=ngroups, ncol=5)

demo[,1] <- alower
demo[,2] <- 100000/lifespan*alength # Popsizes
demo[,3] <- demo[,2]/sum(demo[,2]) # proportions in each age group
demo[,4] <- c(1/(alength*365)) # ageing rates out of compartment
demo[,5] <- c(demo[nrow(demo),4], demo[1:(nrow(demo)-1),4]) # ageing rates into compartment
demo

# Make up some social mixing data
mixmat <- matrix(data=c(10.9, 4.5, 4.5, 6.1), nrow=2)

# Theta
beta <- 0.02
gamma <- 1/5
mu <- 1/(70*365)
sigma <- 1/365
k <- 5
eta <- 0.2 
phi <- 10

init <- log(c(demo[,"N"], rep(1, ngroups), rep(1e-12, ngroups))) 

theta <- c(beta=beta, gamma=gamma, sigma=sigma, eta=eta, phi=phi, ngroups=ngroups, k=k)
theta <- list(theta, demo, mixmat)

enviro <- diffeqr::diffeq_setup()

model <- function(logstate, parms, times) {
  
  # this doesnt work
  theta = as.vector(parms[[1]]) # this should be a vector
  demo = as.array(parms[[2]]) # this should be matrix/array
  mixmat = as.array(parms[[3]]) # this should be matrix/array
  
  ageout = demo[,4]
  agein = demo[,5]
  
  beta = theta[1]
  gamma = theta[2]
  sigma = theta[3]
  eta = theta[4]
  phi = theta[5]
  ngroups = theta[6]
  
  # states
  state = exp(logstate)
  S = state[(1:ngroups)+0*ngroups]
  I = state[(1:ngroups)+1*ngroups]
  R = state[(1:ngroups)+2*ngroups] 
  N = S + I + R
  
  Slower <- c(N[ngroups], S[-ngroups])
  Ilower <- c(0, I[-ngroups])
  Rlower <- c(0, R[-ngroups])
  
  # calculate FOI
  beta_eff = beta * (1+eta*sin(pi*(times-phi)/365.0)^2)
  FOI = beta_eff*colSums(mixmat*I/N)
  
  dSdt = -FOI*S + sigma*R - ageout*S  + agein*Slower
  dIdt = FOI*S - gamma*I - ageout*I + agein*Ilower
  dRdt = gamma*I - sigma*R - ageout*R + agein*Rlower
  
  return(list(cbind(dSdt/S, dIdt/I, dRdt/R)))
  
}

problem <- enviro$ODEProblem(model, init, c(0.0, maxtime), theta)
problemacc <- diffeqr::jitoptimize_ode(enviro,problem) 
solution <- enviro$solve(problemacc, enviro$AutoVern7(enviro$KenCarp4()), 
                         saveat=1.0, abstol=1e-8, reltol=1e-8) # default tolerance leads to weird behaviour with this model (generation of individuals)

throws this error:

Error: Error happens in Julia.
MethodError: no method matching reshape(::Operation, ::Int64)
Closest candidates are:
  reshape(!Matched::FillArrays.AbstractFill, ::Union{Colon, Int64}...) at C:\Users\fabienne\.julia\packages\FillArrays\RM6r2\src\FillArrays.jl:204
  reshape(!Matched::OffsetArrays.OffsetArray, ::Union{Colon, Int64}...) at C:\Users\fabienne\.julia\packages\OffsetArrays\7j7P7\src\OffsetArrays.jl:240
  reshape(!Matched::AbstractMultiScaleArray, ::Int64...) at C:\Users\fabienne\.julia\packages\MultiScaleArrays\c9WNi\src\diffeq.jl:49
  ...
Stacktrace:
 [1] simple_call(::String, ::Operation, ::Vararg{Any,N} where N; kwargs::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{,Tuple{}}}) at C:\Users\fabienne\Documents\R\win-library\4.0\JuliaCall\julia\interface1.jl:11
 [2] simple_call(::String, ::Operation, ::Vararg{Any,N} where N) at C:\Users\fabienne\Documents\R\win-library\4.0\JuliaCall\julia\interface1.jl:2
 [3] julia_extptr_callback(::Ptr{ListSxp}) at C:\Users\fabienne\.julia\packages\ Error: Error happens in Julia.
REvalError:

Is there a way around this? If not it might be necessary to write the function as JuliaCall::julia_eval. Is there an example of an ODE function that uses array/matrix calculation in Julia lang (doesn't have to be an SIR, any compartmental model will do)? Thanks

[ChrisRackauckas]

Arrays and matrices are fine. The issue here seems to be that the parameters theta are an array of arrays. The JIT compilation portion cannot handle that yet. That's on the list but it'll take some time.

[fkrauer]

ok, thanks for the info

[mbkoltai]

Hi,
I have a problem when passing a matrix as input too:
p <- c(10,28,8/3); K_lor=matrix(c(-p[1],p[1],0,p[2],-1,0,0,0,-p[3]),3,3)
f_ode_julia_vect <- function(u,x,t) {du=x%*%u + c(0,-u[1]*u[3],u[1]*u[2]); return(du) }
prob<-de_jul$ODEProblem(f_ode_julia_vect, u0,tspan,K_lor); fastprob=diffeqr::jitoptimize_ode(de_jul,prob)

gives me:
Error in x %*% u : requires numeric/complex matrix/vector arguments

(if I write 'x*u' in the equation then it's not a vector product I get.)

Does 'jitoptimize_ode' use some other notation for vector multiplication?
Thanks for any info.

[ChrisRackauckas]

It tries to scalarize it. I assume the issue might be the fact that R is using a BLAS for the matmul which isn't getting properly traced and replaced. If it was using an algorithm with indexing under the hood it would.