Autodifferencing and type stability in DifferentialEquations.jl

I’m working to optimize a solver for a large, stiff, system of differential equations using DifferentialEquations.jl. I’ve been following through the guide on Code Optimization in the documentation, and have been trying to minimize my allocations. I’m using QNDF() as my solver and I’ve found that I can’t just preallocate all of my arrays because u and du change type from Vector{Float64} to something of type Float64(::ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEq.OrdinaryDiffEqTag, Float64}, Float64, 12}}

Currently, in my ODEFunction, I test for the type of du, and allocate a new array if it is not of type Vector{Float64}. I was wondering if there is a more efficient approach to this. I have tried turning autodifferencing off, by setting QNDF(autodiff=false), but this doesn’t improve my overall solver speed.

My ODEFunction roughly looks like this - I can post a better example if needed. I’m mostly just wondering if there’s any documentation I’m missing that helps handle this change of type in a memory efficient way.

function odefunc!(du, u, p, t)
    #load parameters 
    L,B,Lu,Bu= p

    #check type
    if typeof(du) != Vector{Float64}
        Lu = similar(du)
        Bu = similar(Lu)
    end

    mul!(Lu, L, u)
    mul!(Bu, B, u)
    @. du = Lu + Bu
    nothing
end

Any thoughts appreciated!

Hi @brynn, welcome!

I don’t believe I can check your MWE, but at least I’m missing how BF is used?

Edit: ah, OK! But there are still some parts missing?

Sorry - I removed most of the code to just try and show what my real question is. I can come up with code that actually can be run to show what my problem is (it will just take me a little while to figure that out!)
Not sure if this clarification helps at all but: in my main script, Lu and Bu are instantiated as type Vector{Float64} and then passed in as a parameter, so that each time the solver calls my ODEFunction, it should reuse those variables and not allocate anything new. However, for some timesteps, u and du are not simple vectors and when multiplied by u, the result can’t be allocated into Lu and Bu because it’s the wrong type. I wasn’t sure if this was a common problem and if there is a cleaner solution than what I’m doing.

Thank you!

If you try to come up with a MWE you either will find a problem yourself or help us to reproduce and (hopefully;) improve. As far as I know type instability should not be a problem in your case.

This sounds unusual however. Maybe convert them before in this case?

Check out PreallocationTools.jl. See the FAQ on this.

https://diffeq.sciml.ai/stable/basics/faq/#I-get-Dual-number-errors-when-I-solve-my-ODE-with-Rosenbrock-or-SDIRK-methods

I should probably add an example on this to that tutorial.

1 Like

Thanks for pointing that out, it worked really well!