First run of ModelingToolkit + DE solve() is extremely slow

[70Gage70]

The code in question takes some interpolants of wind/water vector fields (in the MWE I have replaced these with interpolations of a constant), then constructs a Clump, which is a system of two first order ODEs for the x, y coordinates of the Clump. Then, a Raft is a network of clumps connected by springs with the appropriate spring force. A RectangularRaft is a Raft initialized with Clumps in an evenly-spaced rectangular grid.

There are two files, parameters_mwe.jl which stores some physical constants:

struct BOMParameters{T<:Real}
    α::T
    τ::T
    R::T
    f::T
end

function BOMParameters(;
    δ::Real = 1.25,
    a::Real = 1.0e-4,
    ρ::Real = 1027.0e9,
    ρa::Real = 1.2e9,
    ν::Real = 8.64e-8,
    νa::Real = 1.296e-6,
    Ω::Real = 2*π,
    ϑ0::Real = 10.0)

    μ = ν * ρ
    μa = νa * ρa
    γ = μa/μ

    ψ = (im*sqrt(1 - (2/δ - 1)^2) + 2/δ - 1)^(1/3)
    Φ = real(im*sqrt(3)/2 * (1/ψ - ψ) - 1/2 * (1/ψ + ψ) + 1)
    Ψ = 1/π * acos(1 - Φ) - 1/π * (1 - Φ) * sqrt(1 - (1 - Φ)^2)

    α = γ*Ψ/(1 - (1 - γ)*Ψ)
    τ = (1 - Φ/6)/(1 - (1 - γ)*Ψ) * (a^2 * ρ / (3*μ*δ^4))
    R = (1 - Φ/2)/(1 - Φ/6)

    f = 2*Ω*sin(ϑ0*π/180)

    return BOMParameters(α, τ, R, f)
end

struct SpringParameters{F<:Function, T<:Real}
    k::F
    L::T
end

function Base.length(::SpringParameters)
    return 1
end

function Base.iterate(sp::SpringParameters)
    return (sp, nothing)
end

function Base.iterate(::SpringParameters, ::Nothing)
    return nothing
end

function Base.show(io::IO, x::SpringParameters)
    print(io, "SpringParameters[1->")
    show(io, x.k(1))
    print(io, ", ")
    show(io, length(x.L))
    print(io, "]")
end

function spring_force_x(x1::Real, x2::Real, y1::Real, y2::Real, parameters::SpringParameters)
    k, L = (parameters.k, parameters.L)
    d = sqrt((x1 - x2)^2 + (y1 - y2)^2)
    return k(d)*(x1 - x2)*(L/d - 1)
end

function spring_force_y(x1::Real, x2::Real, y1::Real, y2::Real, parameters::SpringParameters)
    k, L = (parameters.k, parameters.L)
    d = sqrt((x1 - x2)^2 + (y1 - y2)^2)
    return k(d)*(y1 - y2)*(L/d - 1)
end

And mwe.jl which constructs all the required functions via ModelingToolkit:

using Interpolations, DifferentialEquations, ModelingToolkit
using LinearAlgebra: ⋅

include("parameters_mwe.jl")

lon = -180:1:180
lat = -90:0.5:90
time = 0.0:1:365
water = fill(1.0, length(lon), length(lat), length(time))
wind = fill(1.0, length(lon), length(lat), length(time))

interpolant_type = BSpline(Cubic(Line(OnGrid()))) 
const water_itp = scale(interpolate(water, interpolant_type), lon, lat, time)
const wind_itp = scale(interpolate(water, interpolant_type), lon, lat, time)

MaterialDerivativeX(VF, x, y, t) = gradient(VF, x, y, t) ⋅ [VF(x, y, t), VF(x, y, t), 1.0]
MaterialDerivativeY(VF, x, y, t) = gradient(VF, x, y, t) ⋅ [VF(x, y, t), VF(x, y, t), 1.0]
Vorticity(VF, x, y, t) = gradient(VF, x, y, t)[1] - gradient(VF, x, y, t)[2]

@variables t            
ddt = Differential(t)

v_x(x, y, t) = water_itp(x, y, t)
v_y(x, y, t) =  water_itp(x, y, t)
Dv_xDt(x, y, t) = MaterialDerivativeX(water_itp, x, y, t)
Dv_yDt(x, y, t) = MaterialDerivativeY(water_itp, x, y, t)
u_x(x, y, t, α) = (1 - α) * water_itp(x, y, t) + α * wind_itp(x, y, t)
u_y(x, y, t, α) = (1 - α) * water_itp(x, y, t) + α * wind_itp(x, y, t)
Du_xDt(x, y, t, α) = (1 - α) * MaterialDerivativeX(water_itp, x, y, t) + α * MaterialDerivativeX(wind_itp, x, y, t) 
Du_yDt(x, y, t, α) = (1 - α) * MaterialDerivativeY(water_itp, x, y, t) + α * MaterialDerivativeY(wind_itp, x, y, t) 
ω(x, y, t) = Vorticity(water_itp, x, y, t)

@register_symbolic v_x(x, y, t)
@register_symbolic v_y(x, y, t)
@register_symbolic Dv_xDt(x, y, t)
@register_symbolic Dv_yDt(x, y, t)
@register_symbolic u_x(x, y, t, α)
@register_symbolic u_y(x, y, t, α)
@register_symbolic Du_xDt(x, y, t, α)
@register_symbolic Du_yDt(x, y, t, α)
@register_symbolic ω(x, y, t)

@register_symbolic spring_force_x(x1, x2, y1, y2, parameters)
@register_symbolic spring_force_y(x1, x2, y1, y2, parameters)

function Clump(
    xy0::Vector{<:Real},
    clump_parameters::BOMParameters;
    name::Symbol)

    ps = @parameters α = clump_parameters.α τ = clump_parameters.τ R = clump_parameters.R f = clump_parameters.f
    @variables x(t) = xy0[1] y(t) = xy0[2]
    @variables Fx(t) Fy(t)

    eqs = [
    ddt(x) ~ u_x(x, y, t, α) + τ * (
        R*Dv_xDt(x, y, t) - R*(f + ω(x, y, t)/3)*v_y(x, y, t) - Du_xDt(x, y, t, α) + (f + R*ω(x, y, t)/3)*u_y(x, y, t, α) + Fx
    ),
    ddt(y) ~ u_y(x, y, t, α) + τ * (
        R*Dv_yDt(x, y, t) + R*(f + ω(x, y, t)/3)*v_x(x, y, t) - Du_yDt(x, y, t, α) - (f + R*ω(x, y, t)/3)*u_x(x, y, t, α) + Fy
    )
    ]

    return ODESystem(eqs, t, [x, y, Fx, Fy], ps; name)
end

function Raft(
    xy0::Matrix{<:Real},
    clump_parameters::Vector{<:BOMParameters},
    spring_parameters::Matrix{<:SpringParameters};
    name::Symbol)

    N_clumps = size(xy0, 1)
    @named clump 1:N_clumps i -> Clump(xy0[i,:], clump_parameters[i])

    forces_x = [clump[i].Fx ~ sum([spring_force_x(clump[i].x, clump[j].x, clump[i].y, clump[j].y, spring_parameters[i, j])] for j = 1:N_clumps if j != i)[1]
        for i = 1:N_clumps
    ]

    forces_y = [clump[i].Fy ~ sum([spring_force_y(clump[i].x, clump[j].x, clump[i].y, clump[j].y, spring_parameters[i, j])] for j = 1:N_clumps if j != i)[1]
        for i = 1:N_clumps
    ]

    eqs = [forces_x ; forces_y]

    return compose(ODESystem(eqs, t; name = name), clump[1:N_clumps]...)
end

function RectangularRaft(
    x_range::AbstractRange{<:Real}, 
    y_range::AbstractRange{<:Real},
    clump_parameters::BOMParameters, 
    spring_parameters::SpringParameters; 
    name::Symbol)

    # a rectangular mesh
    network = reverse.(collect(Iterators.product(reverse(y_range), x_range))) # reverse so that the first row has the largest y
    n_col = length(x_range)
    n_row = length(y_range)
    N_clumps = length(network)

    xy0 = Matrix{Float64}(undef, N_clumps, 2)
    clump_parameters_raft = Vector{BOMParameters}(undef, N_clumps)
    spring_parameters_raft = Matrix{SpringParameters}(undef, N_clumps, N_clumps)

    n(i, j) = (i - 1) * n_col + j

    for i = 1:n_row
        for j = 1:n_col
            xy0[n(i, j), :] .= network[i, j] 
            clump_parameters_raft[n(i, j)] = clump_parameters
            connections = filter(idx -> (1 <= idx[1] <= n_row) && (1 <= idx[2] <= n_col), [(i-1, j), (i+1, j), (i, j-1), (i, j+1)])
            connections = map(x->n(x...), connections)

            spring_parameters_raft[n(i, j), connections] .= spring_parameters
            spring_parameters_raft[n(i, j), setdiff(1:N_clumps, connections)] .= SpringParameters(k->0, 0)
        end
    end

    return Raft(xy0, clump_parameters_raft, spring_parameters_raft, name = name)
end

function sol(n)
    @info "Generating model."

    x0, y0 = 0.0, 0.0
    x_range = range(start = x0 - 5, length = n, stop = x0 + 5)
    y_range = range(start = y0 - 5, length = n, stop = y0 + 5)
    clump_parameters = BOMParameters()
    spring_parameters = SpringParameters(k -> 20, 0.4)

    @named RRaft = RectangularRaft(x_range, y_range, clump_parameters, spring_parameters)
    @time RRaft = structural_simplify(RRaft)

    t_range = (0.0, 1.0)

    @info "Generating problem."

    @time prob = ODEProblem(RRaft, [], t_range, [], jac = true, sparse = true)

    @info "Solving problem."

    @time sol = solve(prob)

    return sol
end

The argument n of sol basically controls how many Clumps there are (n^2). In general, n could be very large (say ~100) but even with n = 10 I find that the first run of solve() takes ~ 125 s and the second takes ~ 0.008 s. See here:

Is there anything I can do about this?

Files
parameters_mwe.jl (1.6 KB)
mwe.jl (5.1 KB)

[ufechner7]

Create a system image with all required packages. I use this bash script:

#!/bin/bash -eu
update=true
if [[ $# -gt 0 ]]; then
    if [[ $1 != "--update" ]]; then
        echo "Invalid parameter! Use:"
        echo "./create_sys_image"
        exit 1
    fi
fi

if [[ $(basename $(pwd)) == "bin" ]]; then
    cd ..
fi

julia_version=$(julia --version | awk '{print($3)}')
julia_major=${julia_version:0:3} 
if test -f "kps-image-${julia_major}.so"; then
    mv bin/kps-image-${julia_major}.so bin/kps-image-${julia_major}.so.bak
fi

echo "Updating packages..."
if test -f "Manifest.toml"; then
   mv Manifest.toml Manifest.toml.bak
fi
julia --pkgimages=no --project -e "include(\"./test/update_packages.jl\");"

julia --pkgimages=no --project -e "using Pkg; Pkg.precompile()"
julia --pkgimages=no --project -e "include(\"./test/create_sys_image.jl\");"
mv kps-image_tmp.so bin/kps-image-${julia_major}.so

echo "Precompiling the plotting libraries, a few minutes left..."
julia -J  bin/kps-image-${julia_major}.so --project -e "using PythonPlot"
julia -J  bin/kps-image-${julia_major}.so --project -e "using Plots" > /dev/null 2>&1
cd ..

combined with these Julia scripts:

@info "Updating packages ..."
using Pkg
Pkg.instantiate()
Pkg.update()
Pkg.precompile()

and

using Pkg

@info "Loading packages ..."
using ModelingToolkit, OrdinaryDiffEq, DifferentialEquations, DataInterpolations, LaTeXStrings, Timers, 
      PackageCompiler, ControlSystemsBase, Roots, JLD2

@info "Creating sysimage ..."
push!(LOAD_PATH,joinpath(pwd(),"src"))

PackageCompiler.create_sysimage(
    [:ModelingToolkit, :OrdinaryDiffEq, :DifferentialEquations, :DataInterpolations, :LaTeXStrings, :Timers, 
     :ControlSystemsBase, :Roots, :JLD2];
    sysimage_path="kps-image_tmp.so",
    precompile_execution_file=joinpath("test", "test_for_precompile.jl")
)

and

let
    include("../src/init.jl")
    run_tests(basic)
    nothing   
end

using ControlSystemsBase, Roots
P = tf(-1.898e-8, [1, 0.01934])

function dc_gain(sys)
      sys.matrix[1].num[0] / sys.matrix[1].den[0]
 end
 
 # calculate the 3db bandwidth of a low-pass like system
 function w_3db(sys)
     gain_3db = abs(dc_gain(sys)) *  exp10(-3/20)
     function res(w)
         mag = bode(sys, [w])[1][1]
         mag - gain_3db
     end
     res = find_zero(res, (exp10(-4), exp10(1)))
     return res
 end

 println(w_3db(P))

 using JLD2
 dict=load("data/lin_turbine.jld2")

@info "Precompile script has completed execution."

This is probably more complex than needed, but works fine for me, and my first-time-to-model-simulated is about 10s.

[70Gage70]

Thanks for the suggestion. I must not understand the system images method because it didn’t work at all for me. I basically did this:

using PackageCompiler
create_sysimage(["ModelingToolkit, DifferentialEquations, Interpolations"], sysimage_path="mwe_test.so", precompile_execution_file="mwe.jl")

But this took longer to precompile (~ 4 min) than the code itself takes to run. More to the point, it didn’t speed up sol(10) at all when I reloaded julia with the --sysimage flag. It seems like from the docs, the file you make a system image of needs to have the exact code you want to run, but if I append sol(10) to mwe.jl and create the system image, it just runs sol(10) anyway. And what if I wanted sol(15) etc?

[ufechner7]

You do NOT need to have the exact code in your pre-compile file. Just make sure that a large part of the required code becomes part of the package image, that is usually sufficient.

[70Gage70]

But even if I make the entire code part of the package image, it has no effect on the speed:

[ufechner7]

OK, then your problem is different from what I am doing.

@ChrisRackauckas Can you comment on this? Is this to be expected, or is this problem particularly difficult to solve?

On the other hand, if you say this would be slow, what do you compare it with? I remember that I used to have Simulink models that needed a few minutes compilation time before the simulation started…

[ChrisRackauckas]

It’s a known issue with ModelingToolkit that large equations generated by for loops have a long compile time. This is fixed by what’s known as SymbolicIR.jl along with FastDifferentiation.jl, both of which are not currently integrated but should be rather soon. Wait a bit for the JuliaCon announcement.