Confusing memory allocations when using the integrator of DifferentialEquations.jl

General Usage Performance
, ,
1

Here is a prototype of my problem:

using OrdinaryDiffEq
n = 2^18;
u0 = rand(ComplexF64, n);

function f(du, u, p, t)
    for i in 1:length(u)
        du[i] = -1.0im * u[i]
    end
    return
end

mutable struct A
    u0::Vector{ComplexF64}
    integrator

    function A(u0)
        a = new()
        a.u0 = u0
        prob = ODEProblem(f, u0, (0, 1), nothing)
        a.integrator = init(prob, Tsit5(), save_on=false, save_start=false, save_end=false, calck=false, alias_u0=true, abstol=1e-8, reltol=1e-8);
        return a
    end
end

function g(u1, u2)
    x = 0.0im
    for i in 1:length(u1)
        x += u1[i] * conj(u2[i])
    end
    return x
end

function g2(a::A)
    u1 = a.u0
    u2 = a.integrator.u
    x = 0.0im
    for i in 1:length(u1)
        x += u1[i] * conj(u2[i])
    end
    return x
end

a = A(u0);
step!(a.integrator, 1, true)

Basically, g() and g2() calculate the dot product of the initial and final states. (Yes, I could have used dot(), but this is a prototype.) Naively, I expect they behave similarly. It turns out that g2() allocates extra memory in each call. Can someone please give me some insight? Thanks!

julia> @time g(a.u0, a.integrator.u)
  0.000627 seconds (1 allocation: 32 bytes)
174863.6176038208 + 0.0im

julia> @time g2(a)
  0.042788 seconds (1.57 M allocations: 43.992 MiB, 36.44% gc time)
174863.6176038208 + 0.0im
2

Your struct isn’t type stable. Specifically, the compiler can’t figure out what the type of a.integrator is.

3 
mutable struct A{I}
    u0::Vector{ComplexF64}
    integrator::I
4

Indeed. Since I know the type of a.integrator.u, the problem is fixed by redefining g2().

function g2(a::A)
    u1 = a.u0
    u2 = a.integrator.u::Vector{ComplexF64}
    x = 0.0im
    for i in 1:length(u1)
        x += u1[i] * conj(u2[i])
    end
    return x
end
5

It would work except that the concrete type of an ODE integrator is very complicated :sweat:

6

That’s why you just parameterize it as I.

7

In my problem I don’t know the concrete type of the integrator until during the construction of the structure …

8

You don’t need to. Use parametric typing.

struct A{T}
  x::T
end

A(2.0) # A{Float64}(2.0)
A(2) # A{Int}(2)
9

In the code example I posted, the type of a.integrator is known inside the constructor. It is given by init()depending on the parameters I chose.

10 
mutable struct A{I}
    u0::Vector{ComplexF64}
    integrator::I

    function A(u0)
        u0 = u0
        prob = ODEProblem(f, u0, (0, 1), nothing)
        integrator = init(prob, Tsit5(), save_on=false, save_start=false, save_end=false, calck=false, alias_u0=true, abstol=1e-8, reltol=1e-8);
        return new(u0,integrator)
    end
end
11

Thanks for the suggestions @ChrisRackauckas. I misremembered. I actually passed a as the last parameter to ODEProblem() (replacing nothing) in the real code. Because Julia doesn’t support mutual reference, I defined a.integrator not in the constructor but in a separate function.