Passing jacobian generated from ModelingToolkit to ODEFunction for EnsembleGPUArray

Specific Domains Modelling & Simulations
, ,
1

Hello all!

I want to do some tests with GPU Parallel Ensemble Simulation using DifferentialEquations.jl. As a starting point, I was looking into the example in the DiffEqGPU.jl page. Regarding Stiff ODEs, it’s said that:

Stiff ODEs require the analytical solution of every derivative function it requires. For example, Rosenbrock methods require the Jacobian and the gradient with respect to time, and so these two functions are required to be given. Note that they can be generated by the modelingtoolkitize approach.

Since my model is stiff (I use CVODE_BDF from Sundials.jl) and analytical Jacobians/gradients are impractical for this model, I wanted to use ModelingToolkit.jl to generate these for me.

However, when trying to generate the Jacobian and the time gradient for the Lorenz system from the example, I’m getting an error:

function lorenz(du,u,p,t)
    du[1] = p[1]*(u[2]-u[1])
    du[2] = u[1]*(p[2]-u[3]) - u[2]
    du[3] = u[1]*u[2] - p[3]*u[3]
end
u0 = Float32[1.0;0.0;0.0]
tspan = (0.0f0,100.0f0)
p = [10.0f0,28.0f0,8/3f0]
prob = ODEProblem(lorenz,u0,tspan,p)
sys = modelingtoolkitize(prob)
sys_jac = eval(generate_jacobian(sys)[2])
sys_tgrad = eval(generate_tgrad(sys)[2])
func = ODEFunction(sys,jac=sys_jac,tgrad=sys_tgrad)

ERROR: TypeError: non-boolean (var"#126#127") used in boolean context
Stacktrace:
 [1] (ODEFunction{true})(sys::ODESystem, dvs::Vector{Term{Real, Base.ImmutableDict{DataType, Any}}}, ps::Vector{Sym{Real, Base.ImmutableDict{DataType, Any}}}, u0::Nothing; version::Nothing, tgrad::Function, jac::Function, eval_expression::Bool, sparse::Bool, simplify::Bool, eval_module::Module, steady_state::Bool, checkbounds::Bool, sparsity::Bool, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})        
   @ ModelingToolkit C:\Users\Ricardo\.julia\packages\ModelingToolkit\tMgaW\src\systems\diffeqs\abstractodesystem.jl:272
 [2] #ODEFunction#387
   @ C:\Users\Ricardo\.julia\packages\ModelingToolkit\tMgaW\src\systems\diffeqs\abstractodesystem.jl:235 [inlined]

Am I using the wrong ModelingToolkit function? Maybe I’m not passing the correct part of the functions? I’m very unfamiliar with ModelingToolkit, so any help is appreciated!

2

ODEFunction on sys just needs Bools. I’m not sure why this was made so complicated.

function lorenz(du,u,p,t)
    du[1] = p[1]*(u[2]-u[1])
    du[2] = u[1]*(p[2]-u[3]) - u[2]
    du[3] = u[1]*u[2] - p[3]*u[3]
end
u0 = Float32[1.0;0.0;0.0]
tspan = (0.0f0,100.0f0)
p = [10.0f0,28.0f0,8/3f0]
prob = ODEProblem(lorenz,u0,tspan,p)
sys = modelingtoolkitize(prob)
func = ODEFunction(sys,jac=true,tgrad=true)
2 Likes
3

Thank you Chris! The example in the DiffEqGPU.jl page used actual functions as inputs, that’s what confused me. Cheers!

1 Like
4

Ahh yes. Symbolics systems want true/false for whether to generate these things. But the numerical ODEFunction constructor wants the functions.

1 Like
5

I am now running into a problem when trying to call solve:

function lorenz(du,u,p,t)
    du[1] = p[1]*(u[2]-u[1])
    du[2] = u[1]*(p[2]-u[3]) - u[2]
    du[3] = u[1]*u[2] - p[3]*u[3]
end
u0 = Float32[1.0;0.0;0.0]
tspan = (0.0f0,100.0f0)
p = [10.0f0,28.0f0,8/3f0]
prob = ODEProblem(lorenz,u0,tspan,p)
sys = modelingtoolkitize(prob)
func = ODEFunction(sys,jac=true,tgrad=true)
prob_func = (prob,i,repeat) -> remake(prob,p=rand(Float32,3).*p)
prob_jac = ODEProblem(func,u0,tspan,p)
monteprob_jac = EnsembleProblem(prob_jac, prob_func = prob_func)
@time solve(monteprob_jac,Rodas5(),EnsembleGPUArray(),dt=0.1,trajectories=10_000,saveat=1.0f0)

ERROR: GPU compilation of kernel #gpu_gpu_kernel(KernelAbstractions.CompilerMetadata{KernelAbstractions.NDIteration.DynamicSize, KernelAbstractions.NDIteration.DynamicCheck, Nothing, CartesianIndices{1, Tuple{Base.OneTo{Int64}}}, KernelAbstractions.NDIteration.NDRange{1, KernelAbstractions.NDIteration.DynamicSize, KernelAbstractions.NDIteration.DynamicSize, CartesianIndices{1, Tuple{Base.OneTo{Int64}}}, CartesianIndices{1, Tuple{Base.OneTo{Int64}}}}}, ModelingToolkit.var"#f#396"{RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", Expr}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋out, :ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", Expr}}, CUDA.CuDeviceMatrix{Float32, 1}, CUDA.CuDeviceMatrix{Float32, 1}, CUDA.CuDeviceMatrix{Float32, 1}, Float32) failed
KernelError: passing and using non-bitstype argument

Argument 3 to your kernel function is of type ModelingToolkit.var"#f#396"{RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", Expr}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋out, 
:ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", Expr}}, which is not isbits:
  .f_oop is of type RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", Expr} which is not isbits.
    .body is of type Expr which is not isbits.
      .head is of type Symbol which is not isbits.
      .args is of type Vector{Any} which is not isbits.
  .f_iip is of type RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:ˍ₋out, :ˍ₋arg1, :ˍ₋arg2, :t), ModelingToolkit.var"#_RGF_ModTag", ModelingToolkit.var"#_RGF_ModTag", Expr} which is not isbits.
    .body is of type Expr which is not isbits.
      .head is of type Symbol which is not isbits.
      .args is of type Vector{Any} which is not isbits.

Am I missing any steps?

6

Oh open an issue for that. Seems like the way we automatically generate the functions doesn’t play nicely with CUDA.jl. So then you’d have to do the function generation completely manually, i.e.:

sys_f = eval(generate_function(sys)[2])
sys_jac = eval(generate_jacobian(sys)[2])
sys_tgrad = eval(generate_tgrad(sys)[2])
func = ODEFunction(sys_f ,jac=sys_jac,tgrad=sys_tgrad)
1 Like
7

Thank you Chris! Shall I open it for ModelingToolkit.jl or DiffEqGPU.jl?

8

DiffEqGPU.jl

1 Like