Introducing Callback for Simulations

[glwagner]

This PR implements Callback, designed to be used with Simulation.

Illustration:

using Oceananigans

model = NonhydrostaticModel(grid = RegularRectilinearGrid(size=(128, 128), extent=(2π, 2π), topology=(Periodic, Periodic, Flat)),
                            timestepper = :RungeKutta3,
                            advection = UpwindBiasedFifthOrder(),
                            buoyancy = nothing,
                            tracers = nothing)

set!(model, u = (x, y, z) -> randn(), v = (x, y, z) -> randn())

simulation = Simulation(model, Δt=0.01, stop_iteration=100)

print_progress(sim) = @info "Iteration: $(sim.model.clock.iteration), time: $(sim.model.clock.time)"
simulation.callbacks[:progress] = Callback(print_progress, schedule=IterationInterval(10))

run!(simulation)

we get

julia> run!(simulation)
[ Info: Iteration: 0, time: 0.0
[ Info: Iteration: 10, time: 0.09999999999999999
[ Info: Iteration: 20, time: 0.19999999999999984
[ Info: Iteration: 30, time: 0.30000000000000004
[ Info: Iteration: 40, time: 0.4000000000000007
[ Info: Iteration: 50, time: 0.5000000000000013
[ Info: Iteration: 60, time: 0.6000000000000003
[ Info: Iteration: 70, time: 0.6999999999999993
[ Info: Iteration: 80, time: 0.7999999999999983
[ Info: Iteration: 90, time: 0.8999999999999972
[ Info: Iteration: 100, time: 0.9999999999999962
[ Info: Simulation is stopping. Model iteration 100 has hit or exceeded simulation stop iteration 100.

We also support passing an iterable of callbacks to run!:

progress_callback = Callback(print_progress, schedule=IterationInterval(10))
run!(simulation, callbacks=[progress_callback])

just in case someone wants to do that...

I think we could also redesign TimeStepWizard to be a special kind of Callback, and nuke the progress property.

What do others think? Are they ok with this big change to the API?

cc @navidcy @ali-ramadhan @francispoulin

[tomchor]

Basically a "callback" is a way to run any function that takes a Simulation as an argument at a given schedule? Be it a function that prints things, updates model parameters (like TimeStepWizard) or whatever?

If so, I'm very much onboard with this feature 👍

[francispoulin]

Seems like a good idea but let me make sure I understand it.

If I take your example and modify it slightly, change 10 to 1, should it give output everytime step?

progress_callback = Callback(print_progress, schedule=IterationInterval(10))
run!(simulation, callbacks=[progress_callback])

[glwagner]

Basically a "callback" is a way to run any function that takes a Simulation as an argument at a given schedule? Be it a function that prints things, updates model parameters (like TimeStepWizard) or whatever?

If so, I'm very much onboard with this feature 👍

Yeah... "callback" is an annoyingly obtuse but somehow standard name for this kind of thing

https://en.wikipedia.org/wiki/Callback_(computer_programming)

https://trixi-framework.github.io/Trixi.jl/stable/callbacks/

We could be rebels and call this something else. Like Callafter, or something we come up with.

If TimeStepWizard is a Callback, it updates the time step in Simulation. But this feature can also be used to do things like update an array that holds boundary fluxes. This could be used to couple two domains together (this application is why I want to implement this now).

[glwagner]

Seems like a good idea but let me make sure I understand it.

If I take your example and modify it slightly, change 10 to 1, should it give output everytime step?

progress_callback = Callback(print_progress, schedule=IterationInterval(10))
run!(simulation, callbacks=[progress_callback])

That's right. Callback.func is called on schedule. schedule=IterationInterval(1) means it's called every time-step.

[tomchor]

Yeah... "callback" is an annoyingly obtuse but somehow standard name for this kind of thing

Yeah, definitely not the clearest name... but if it's standard I guess it's good to keep it.

This sounds like it's unambiguously positive though! Nice contribution

[glwagner]

I think I'd like to merge this now, and then make those additional changes later (since the other changes may take some time, and callbacks will be useful in the interim).

[tomchor]

Looks good!

[christophernhill]

@glwagner at some point it might good to think about whether exposing similar interfaces to DiffEq ( e.g. https://diffeq.sciml.ai/dev/features/callback_functions/#The-Callback-Types ) makes sense. Probably for the future.

Same could go for run! (e.g. https://diffeq.sciml.ai/dev/basics/integrator/#Initialization-and-Stepping ) too?

[glwagner]

@glwagner at some point it might good to think about whether exposing similar interfaces to DiffEq ( e.g. https://diffeq.sciml.ai/dev/features/callback_functions/#The-Callback-Types ) makes sense. Probably for the future.

Same could go for run! (e.g. https://diffeq.sciml.ai/dev/basics/integrator/#Initialization-and-Stepping ) too?

Ah wow, that is epic.

I think what's implemented here is more or less analogous to DiscreteCallback.

Reading over the features there makes me realize that we probably want to "align" the time-step (somehow, similar to the root finding feature that's provided for DiffEq) for callbacks. We align time-steps for output, but not for callbacks (yet). That would be nice to add.

I think we could also add an analog of the "continuous callback" --- to the models, not the simulations --- that's executed during update_state!, and therefore can be thought of as a function that "continuously" (eg valid at every moment in time) modifies the state of the model. Perhaps StateModifier or something like that? These objects would not have a schedule since they'd always be executed.

Initialization and finalization might be good wishlist features too for both.