How to reinit! an integrator in a non-allocating and consistent way?

I am trying to create a control function from a ModelingToolkit.jl model:
f(x, u, p) → x_next
I want to do that by re-initializing the integrator with the current state x and input u, and then stepping the integrator:

function f!(xnext, x, u, _, p) # inplace version of f(x, u, p)
    (integrator, set_x, set_u, _) = p
    set_x(integrator, x)
    set_u(integrator, u)
    step!(integrator)
    xnext .= integrator.u # integrator.u is the integrator state, called x in the function
    nothing
end

How can I reinitialize the integrator in a non-allocating way? I am using setu and setp now, but there seems to be some caches in the integrator that are not being reset, as xnext is not consistent when calling the control function multiple times with the same inputs:

xnext = [1.000003514961648, 0.9999867038583831]
xnext = [1.0000016858995648, 0.9999936227203033]
xnext = [1.0000168585118323, 0.9999362274452563]
xnext = [1.0001685367381081, 0.999362298719702]
xnext = [1.0016805307279277, 0.9936254586572106]

And using reinit!(integrator) works, but that uses over 1000 allocations, and performance is critical for this function.

Here is the MWE I created:

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)
set_x = setu(mtk_model, unknowns(mtk_model))
set_u = setp(mtk_model, [mtk_model.τ])
get_h = getu(mtk_model, [mtk_model.y])
p = (integrator, set_x, set_u, get_h)

function f!(xnext, x, u, _, p)
    (integrator, set_x, set_u, _) = p
    set_x(integrator, x)
    set_u(integrator, u)
    step!(integrator)
    xnext .= integrator.u # integrator.u is the integrator state, called x in the function
    nothing
end

for i in 1:10
    local xnext = zeros(2)
    f!(xnext, ones(2), 1.0, nothing, p)
    @show xnext
end
for i in 1:10
    local xnext = zeros(2)
    f!(xnext, zeros(2), 1.0, nothing, p)
    @show xnext
end
1 Like

If you add reset_dt = true you get what you want?

julia> @time reinit!(integrator; reset_dt=true)
  0.000716 seconds (1.05 k allocations: 30.953 KiB)

That still results in 1.05k allocations. reinit!(integrator) works, but has a lot of allocations. While setu is inconsistent, but isn’t allocating.

Open an issue to look into those allocations. My guess is that it has to do with the initial mapping, so it’s related specifically to the interaction of MTK and reinit!. I presume non-MTK models are non-allocating on this?

This works:

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

But reinit_dae is not documented as far as I know. I don’t really know what it does, and if it could be problematic to set it to false.

It is problematic to set it to false because it won’t run the parameter inits and such. But yeah, we can make an issue to make the MTK dae reinit leaner. That would be an MTK issue.

1 Like

Submitted an issue here:

1 Like

As I am still having problems with inconsistent results after resetting the integrator (probably because of some internal states being reset inconsistently), I wanted to try using solve(prob) instead. But now I am getting a new problem: when using setu to update the problem state, this new state is simply ignored by solve. Using setp to update the parameters works fine. What could be the problem here, and is there a workaround? @ChrisRackauckas

MWE:

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

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))
set_x = setu(mtk_model, unknowns(mtk_model))
get_x = getu(mtk_model, unknowns(mtk_model))
set_u = setp(mtk_model, [mtk_model.τ])
get_u = getp(mtk_model, [mtk_model.τ])
get_h = getu(mtk_model, [mtk_model.y])
p = (prob, set_x, set_u, get_h)

function f!(xnext, x, u, _, p)
    (prob, set_x, set_u, _) = p
    set_x(prob, x)
    set_u(prob, u)
    sol = solve(prob, Tsit5(); save_on=false, save_start=false)
    xnext .= sol.u[end] # sol.u is the state, called x in the function
    nothing
end

xnext = zeros(2)

println("Initial")
f!(xnext, zeros(2), 0.0, nothing, p)
@show xnext

println("Just changing state x - doesn't work, xnext stays at zero")
f!(xnext, ones(2), 0.0, nothing, p)
@show xnext

println("Just changing input u - does work, xnext is non-zero")
f!(xnext, zeros(2), 1.0, nothing, p)
@show xnext
nothing

The inconsistent results with the same inputs were due to the adaptive dt in the solver. This issue is solved by resetting the integrator time and dt before stepping:

function f!(xnext, x, u, _, p)
    (integrator, set_x, set_u, _) = p
    set_t!(integrator, 0.0)
    set_proposed_dt!(integrator, Ts)
    set_x(integrator, x)
    set_u(integrator, u)
    step!(integrator, Ts)
    xnext .= integrator.u # sol.u is the state, called x in the function
    nothing
end
1 Like

Open an issue with that in ModelingToolkit.jl. setu needs a similar treatment to what we did to reinit!, in that for MTK-built models it needs a symbolic hook.

1 Like