Ensure type stability when using OrdinaryDiffEq

New to Julia
,
1

Hi,

I am trying to solve large numbers of ODEs and my approach is to put solving the equations part into a function and it should spit out the relevant numbers/arrays. Using some a simple equation, it is akin to the following code

using OrdinaryDiffEq                                      
                                                          
function test(u₀::Float64, tspan::Tuple{Float64, Float64})
    #Half-life of Carbon-14 is 5,730 years.
    C₁ = 5.730

    #Define the problem
    radioactivedecay(u, p, t) = -C₁ * u

    #Pass to solver
    prob = ODEProblem(radioactivedecay, u₀, tspan)
    sol = solve(prob, Tsit5())
    return sol.t, sol[:]
end

Now I am trying to improve the performance and functions shows type instability via @code_warntype (and JET.jl also).

julia> @code_warntype test(1.0, (0.0, 1.0))
MethodInstance for test(::Float64, ::Tuple{Float64, Float64})
  from test(u₀::Float64, tspan::Tuple{Float64, Float64}) @ Main ~/Cloud/GPP/GPP_SUGRA.jl/test-diffeq-type.jl:3
Arguments
  #self#::Core.Const(test)
  u₀::Float64
  tspan::Tuple{Float64, Float64}
Locals
  sol::Any
  prob::Any
  radioactivedecay::var"#radioactivedecay#1"{Float64}
  C₁::Float64
Body::Tuple{Any, Any}
1 ─       (C₁ = 5.73)
│   %2  = Main.:(var"#radioactivedecay#1")::Core.Const(var"#radioactivedecay#1")
│   %3  = Core.typeof(C₁::Core.Const(5.73))::Core.Const(Float64)
│   %4  = Core.apply_type(%2, %3)::Core.Const(var"#radioactivedecay#1"{Float64})
│         (radioactivedecay = %new(%4, C₁::Core.Const(5.73)))
│         (prob = Main.ODEProblem(radioactivedecay::Core.Const(var"#radioactivedecay#1"{Float64}(5.73)), u₀, tspan))
│   %7  = prob::Any
│   %8  = Main.Tsit5()::Core.Const(Tsit5(; stage_limiter! = trivial_limiter!, step_limiter! = trivial_limiter!, thread = static(false),))
│         (sol = Main.solve(%7, %8))
│   %10 = Base.getproperty(sol, :t)::Any
│   %11 = Base.getindex(sol, Main.:(:))::Any
│   %12 = Core.tuple(%10, %11)::Tuple{Any, Any}
└──       return %12

My guess is that the compiler cannot determine the type of prob and sol. My question now is: is this kind of type instability harmful to performance? If so, how can I improve this? Thanks.

2

Creating an ODEProblem might take time. It is better to create it once and then reuse it. In your case, you could do instead

#Define the problem
radioactivedecay(u, p, t) = -p.C₁ * u

p = (C₁ = 5.730,)
u₀ = 1.0
tspan = (0.0, 1.0)
prob = ODEProblem(radioactivedecay, u₀, tspan)

# use this link to change tspan and initial data + solving
solve(remake(prob, p = p, tspan = tspan, u0 = u₀), Tsit5())

See also Parallel Ensemble Simulations · DifferentialEquations.jl which provides implementations exactly for your setting. (WIth added features such as parallelisation…)

3

Since you didn’t use the in-place form of ode, try prob = ODEProblem{false}(radioactivedecay, u₀, tspan).

4

Thanks for the suggestion. I just tried using Parallel Ensemble simulation, but it somehow results in worse performance in my actual code. (Creating ODEProblem takes tiny fraction of time of my whole simulation anyway… and my code was perfectly parallelized also)

5

Thanks for the answer, it seems to be the solution. The radioactive decay example runs now ~5 times faster and @code_warntype doesn’t give red lines :))
It would be interesting to know what causes the difference exactly

6

In your previous codes, I think the compiler can not detect whether the ODEProblem is in place or not, i.e., the form defines the vector field:

function f1(du,u,p,t)
    du[1]=-u[1]
end

or

function f2(u,p,t)
    -u[1]
end

Code like prob=ODEProblem{false}(...) indicates that ODEProblem is not in place, so there is no type-instability now.