Less allocating reinit_dae

[1-Bart-1]

reinit!(integrator, u0) currently allocates a lot due to reinit_dae. It would be great if these allocations could be reduced.

In this MWE, reinit_dae causes around 1000 allocations:

using ModelingToolkit, OrdinaryDiffEq
using ModelingToolkit: D_nounits as D, t_nounits as t, setu, setp, getu

Ts = 0.1
@mtkmodel Pendulum begin
    @parameters begin
        g = 9.8
        L = 0.4
        K = 1.2
        m = 0.3
        τ = 0.0 # input
    end
    @variables begin
        θ(t) = 0.0 # state
        ω(t) = 0.0 # state
        y(t) # output
    end
    @equations begin
        D(θ)    ~ ω
        D(ω)    ~ -g/L*sin(θ) - K/m*ω + τ/m/L^2
        y       ~ θ * 180 / π
    end
end

@mtkbuild mtk_model = Pendulum()
prob = ODEProblem(mtk_model, nothing, (0.0, Ts))
integrator = OrdinaryDiffEq.init(prob, Tsit5(); dt=Ts, abstol=1e-8, reltol=1e-8, save_on=false, save_everystep=false)

@time reinit!(integrator; reinit_dae=true) # 1.05k allocs
@time reinit!(integrator; reinit_dae=false) # 6 allocs

From this discussion on discourse.

[AayushSabharwal]

The latest numbers are

julia> @time reinit!(integrator; reinit_dae=true)
  0.000070 seconds (52 allocations: 1.656 KiB)

julia> @time reinit!(integrator; reinit_dae=false)
  0.000095 seconds (7 allocations: 240 bytes)

While the first case is not ideal, it is better. Almost all of these allocations would be solved by inplace getter functions in SII. I'll try and get around to it, but it'll take a while.

[1-Bart-1]

Huge improvement, thanks! Is this on the main branch?

[1-Bart-1]

I tried to test this with my kite model, but I am still getting 13k allocations with reinit_dae=true using ModelingToolkit#master.

julia> @time OrdinaryDiffEqCore.reinit!(s.integrator; reinit_dae=true)
  0.017841 seconds (12.98 k allocations: 811.258 KiB)

[AayushSabharwal]

What does @time solve(prob.f.initialization_data.initializeprob) look like for your system?

[1-Bart-1]

I am getting an error:

julia> @time NonlinearSolve.solve(s.prob.f.initialization_data.initializeprob)
ERROR: type SCCNonlinearProblem has no field kwargs
Stacktrace:
 [1] getproperty(prob::SCCNonlinearProblem{Vector{…}, false, Vector{…}, Vector{…}, NonlinearFunction{…}, ModelingToolkit.MTKParameters{…}}, name::Symbol)
   @ SciMLBase ~/.julia/packages/SciMLBase/cwnDi/src/problems/nonlinear_problems.jl:519
 [2] solve(::SCCNonlinearProblem{…}; sensealg::Nothing, u0::Nothing, p::Nothing, wrap::Val{…}, kwargs::@Kwargs{})
   @ DiffEqBase ~/.julia/packages/DiffEqBase/VCM6X/src/solve.jl:1081
 [3] solve(::SCCNonlinearProblem{Vector{…}, false, Vector{…}, Vector{…}, NonlinearFunction{…}, ModelingToolkit.MTKParameters{…}})
   @ DiffEqBase ~/.julia/packages/DiffEqBase/VCM6X/src/solve.jl:1079
 [4] macro expansion
   @ ./timing.jl:581 [inlined]
 [5] top-level scope
   @ ./REPL[13]:1
Some type information was truncated. Use `show(err)` to see complete types.

This might be because I am using NonlinearSolve v.4.5.0.

Edit: I found the workaround. And also with the latest version of NonlinearSolve v.4.8.0 I am getting a lot of allocations.

julia> @time NonlinearSolve.solve(s.prob.f.initialization_data.initializeprob, FastShortcutNonlinearPolyalg());
  0.086019 seconds (41.68 k allocations: 7.753 MiB)

[1-Bart-1]

I need this fast reinit for my application. @AayushSabharwal can I help with implementing the inplace getter functions in SII?

[AayushSabharwal]

So it looks like the majority of your allocations are in solving initialization. Passing use_scc = false to ODEProblem should be a decent stopgap measure if your initialization isn't too stiff. The solution then is to optimize SCCNonlinearSolve.jl.

[ChrisRackauckas]

Yes SCCNonlinearSolve.jl still needs a bit of work.

[1-Bart-1]

Some updates seem to have reduced the allocations a lot already. use_scc=false removes half of the allocations, but this is not enough to actually speed up my stiff problem.

use_scc=false:
@time OrdinaryDiffEqCore.reinit!(s.integrator)
0.026057 seconds (2.04 k allocations: 712.352 KiB)
use_ssc=true:
@time OrdinaryDiffEqCore.reinit!(s.integrator)
0.017217 seconds (4.08 k allocations: 574.000 KiB)

[ChrisRackauckas]

Yeah the SCC is just a better way to solve

[hersle]

But here it's the other way around, though. EDIT: stupid me, scc is faster.

[AayushSabharwal]

Maybe the ODE integrator can store the NonlinearSolve cache, so reinit! can reuse that instead of trying to solve

[AayushSabharwal]

But here it's the other way around, though.

The SCC version takes less time and allocates less total memory, despite having a greater number of allocations.

[ChrisRackauckas]

I'll be improving the SCC solvers over the last month but even for now it should be faster