# Type-instability with DifferentialEquations.jl when using nested structs

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

Welcome to the discourse
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`.

What if you do:

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

?

1 Like

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?

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

Sure. I can do that.

1 Like