Periodic Callbacks with general affect function (PoC)

[dodoplus]

This is a proof-of-concept of running general Julia code in event callbacks (in this case, periodic callbacks which are now supported alongside continuous callbacks).

Using equations to define an affect! is great when it works, but it can be limiting.

Defining a new event

function affect!(u,p, t, c)
    update!(c, p.R, p.T)
    p.R = some_complex_logic(p.R, t)
end

ctrl = MyControl()
ps = @parameters R, T
ODESystem(eqs, t, [], ps; name=name, periodic_events=(period, affect!, ctrl))

# ctrl was modified

[YingboMa]

These are defined in the example folder. You can just include the file.

[dodoplus]

Here's a more simplified example/test.

The affect! function performs an update based in the state of the controller, which can perform a complex logic (here simulated by pseudo random number generator) as well as keep (and use) a history. This can be useful, e.g. when simulating hysteresis.

using ModelingToolkit, OrdinaryDiffEq, Test
using Random

struct MyController
    rng
    h
    MyController() = new(MersenneTwister(1234), Float32[])
end

function newdirection!(c::MyController, pars, t)
    # random update
    pars.a, = Random.rand(c.rng, 1)

    # use history for update
    pars.b = t + sum(c.h)
end

function updatehist!(c::MyController, sts)
    push!(c.h, sts.x)
end

function affect!(u, p, t, ctrl::MyController)
    updatehist!(ctrl, u)
    newdirection!(ctrl, p, t)
end

period = 1.0

@variables t, x(t)
@parameters a, b
D = Differential(t)

eqs = [ D(x) ~ a * x + b]

controller = MyController()

@named model = ODESystem(eqs, t, [x], [a, b], periodic_events=(1.0, affect!, controller))

sys = structural_simplify(model)

prob = ODAEProblem(sys, [x => 1.0], (0, 2.0), [a => 1.0, b => 3.0])
sol = solve(prob, Tsit5())

@test length(controller.h) == 1
@test controller.h[1] ≈ 7.873127

[YingboMa]

I am not sure about this design. Cannot we detect periodic callbacks via discrete operators?

[ChrisRackauckas]

That would only allow update equations, not arbitrary affect! functions which is what I think is wanted here. I think we should really just make an arbitrary affect! support for the symbolic condition discrete callbacks though.

[dodoplus]

That would only allow update equations, not arbitrary affect! functions which is what I think is wanted here. I think we should really just make an arbitrary affect! support for the symbolic condition discrete callbacks though.

Absolutely. The point is less about periodic (and discrete in general) events, and more about a more general-purpose affect!.

I've tried to define a reasonable (I hope) interface for such affect!

[YingboMa]

To have better scaling, you can use Dict(reverse(en) for en in enumerate(all_dvs))

[YingboMa]

Do you need Ref? You can always resize! and copyto!.

[YingboMa]

Also, should it be just Vector{Int}?

[dodoplus]

Do you need Ref? You can always resize! and copyto!.

Not sure. I assume that each affect! only needs to access a small part of the overall variable array, so it may be cheaper to use references to those instead of copying the whole thing.

[YingboMa]

Each callback requires at least one dictionary look up. Can we just generate affect functions such that it only does integer indexing? The symbol => index mapping is fixed after the model is compiled.

[dodoplus]

Each callback requires at least one dictionary look up. Can we just generate affect functions such that it only does integer indexing? The symbol => index mapping is fixed after the model is compiled.

I'm not sure what you mean by that: setvec! which is called for each call of the callback function is constant-time. It just assigns a Ref:

function setvec!(vd::VarDict, v)
    vd.₊v[] = v
end

Within the affect! function, accessing a specific variable/parameter would require a dictionary lookup & a de-Ref:

function Base.getproperty(vd::VarDict, name::Symbol)
  ...
  ₊v[][₊d[name]]
end

As you say, the mapping is fixed (and is only computed once), so we could compile the affect! function, replacing e.g. pars.a with integ.p[2] (although I'm not sure how exactly. Rewriting the AST?). However, I think it's not worth it in this case.
What do you think?

[dodoplus]

Hmm... maybe you mean this by @ChrisRackauckas?
https://www.stochasticlifestyle.com/zero-cost-abstractions-in-julia-indexing-vectors-by-name-with-labelledarrays/

[ChrisRackauckas]

LArrays won't scale because there type structure gets bigger, so if you have an ODE of more than 100 variables this is going to be slow to compile. What @YingboMa means is that it should be lowered to a callback function which is just using normal Julia indexing, so the symbolics exists until the ODEProblem is being constructed, and then all of the states/parameters orderings are fixed so you can do codegen that has no symbolic naming anymore.

[dodoplus]

LArrays won't scale because there type structure gets bigger, so if you have an ODE of more than 100 variables this is going to be slow to compile.

Is this a problem in this case, though?
PeriodicEventCallback has states and ps for the list of state variables and parameters the user want to be "exposed" in the callback. By default (if these parameters are not set), only the states/parameters for the sub-system callback is part of are exposed (currently including all sub-systems - which may not be a good behavior by default).

I think it would be quite rare for a single component to have even tens of parameters/states.

What @YingboMa means is that it should be lowered to a callback function which is just using normal Julia indexing, so the symbolics exists until the ODEProblem is being constructed, and then all of the states/parameters orderings are fixed so you can do codegen that has no symbolic naming anymore

I'm not sure what you mean by that. The dictionary in the original proposal was only calculated once (as part of the ODEProblem setup). At runtime, the code should only include a single lookup (which I think is what you're suggesting) and a de-Ref. The Ref can be avoided by using @views.

[dodoplus]

@ChrisRackauckas @YingboMa I've switched back to an improved VarDict (no more Ref).
However, there's a Dict{Symbol, Int} which adds some overhead compared to LArray.

[dodoplus]

will be replaced with a new pull request