Type-instability with DifferentialEquations.jl when using nested structs

General Usage
, ,
1

Hi Julia community,

I have been using DifferentialEquations.jl recently, which is an excellent package. In the spirit of writing performant Julia code, I have been trying to eliminate unnecessary heap allocations.

I have come across a type-instability that I have been unable to fix. The type instability arises when I nest one struct within another. Here I include an MWE illustrating the instability with a struct Fit nested in another struct EOS:

using DifferentialEquations

const conversion_factor = 755400.0

struct Fit
    m₁::Float64
    c₁::Float64
    m₂::Float64
    c₂::Float64

    function Fit()
        m₁ = 1.595
        c₁ = 3.438
        m₂ = 1.075
        c₂ = 3.484

        new(m₁, c₁, m₂, c₂)
    end
end

nbfun(fit::Fit, pnum) = 10^((log10(pnum) - fit.c₁)/fit.m₁)
ε(fit::Fit, nb) = 10^(fit.m₂*log10(nb) + fit.c₂)

struct EOS
    fit::Fit

    function EOS()
        fit = Fit()

        new(fit)
    end
end

ε(eos::EOS, p) = ε(eos.fit, nbfun(eos.fit, conversion_factor*p))/conversion_factor

function structure!(dy_dr, y, eos, r)
    m, p = y

    εᵣ = ε(eos, p)
    dν_dr = 2*(m + 4*π*r^3*p)/(r*(r - 2*m))

    dy_dr[1] = 4*π*r^2*εᵣ
    dy_dr[2] = - (εᵣ + p)*dν_dr/2
    dy_dr
end

function main()
    eos = EOS()
    pc = 1.19e-4

    prob = ODEProblem{true}(structure!, [0, pc], (1e-5, 10), eos)
    sol = solve(prob, Tsit5())
end

When I examine the types, I find

julia> @code_warntype main()
MethodInstance for main()
  from main() @ Main REPL[9]:1
Arguments
  #self#::Core.Const(main)
Locals
  sol::Any
  prob::ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, EOS, ODEFunction{true, SciMLBase.AutoSpecialize, typeof(structure!), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardODEProblem}
  pc::Float64
  eos::EOS
Body::Any
1 ─       (eos = Main.EOS())
│         (pc = 0.000119)
│   %3  = Core.apply_type(Main.ODEProblem, true)::Core.Const(ODEProblem{true})
│   %4  = Main.structure!::Core.Const(structure!)
│   %5  = Base.vect(0, pc::Core.Const(0.000119))::Vector{Float64}
│   %6  = Core.tuple(1.0e-5, 10)::Core.Const((1.0e-5, 10))
│         (prob = (%3)(%4, %5, %6, eos::Core.Const(EOS(Fit(1.595, 3.438, 1.075, 3.484)))))
│   %8  = prob::ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, EOS, ODEFunction{true, SciMLBase.AutoSpecialize, typeof(structure!), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardODEProblem}
│   %9  = Main.Tsit5()::Core.Const(Tsit5(; stage_limiter! = trivial_limiter!, step_limiter! = trivial_limiter!, thread = static(false),))
│   %10 = Main.solve(%8, %9)::Any
│         (sol = %10)
└──       return %10

The type-unstable part is at the end with the variable sol as type Any.

Now, for this simple example, it is trivial to combine the two structs into one, which does eliminate the type-instability. However, for my purposes, I would like to keep the two structs separate.

How can I remove the type-instability? I would be grateful for any tips or suggestions on how I can address this problem.

2 Likes
2

Welcome to the discourse :slight_smile:
Unfortunately, I am not sure where this type instability is coming from. You took good care to not use global variables, no undertyped fields and the like. I suspect it could have to do with inlining, so you could try to just slap @inline in front of all functions and see whether that improves inference.

That being said, I think this particular type instability is not important for performance, since you can just use a function barrier when you continue using sol.

3

What if you do:

DiffEqBase.anyeltypedual(x:: EOS, counter = 0) = Any

?

1 Like
4

Thanks for the suggestions. I can confirm that

DiffEqBase.anyeltypedual(x::EOS, counter=0) = Any

fixes the type-instability.

Is there a way to understand how the type-instability arose in the first place?

5

It’s from the Dual detection, something about Julia’s inference is giving up early. Can you open an issue?

6

Sure. I can do that.

7

I fixed it in Check dualness by fields instead of properties by ChrisRackauckas · Pull Request #999 · SciML/DiffEqBase.jl · GitHub

1 Like