Callbacks

A Callback can be used to execute an arbitrary user-defined function on the simulation at user-defined times.

For example, we can specify a callback which displays the run time every 2 iterations:

using Oceananigans

model = NonhydrostaticModel(grid=RectilinearGrid(size=(1, 1, 1), extent=(1, 1, 1)))

simulation = Simulation(model, Δt=1, stop_iteration=10)

show_time(sim) = @info "Time is $(prettytime(sim.model.clock.time))"

simulation.callbacks[:total_A] = Callback(show_time, IterationInterval(2))

simulation
Simulation of NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0)
├── Next time step: 1 second
├── Elapsed wall time: 0 seconds
├── Wall time per iteration: NaN days
├── Stop time: Inf days
├── Stop iteration : 10.0
├── Wall time limit: Inf
├── Callbacks: OrderedDict with 5 entries:
│   ├── stop_time_exceeded => Callback of stop_time_exceeded on IterationInterval(1)
│   ├── stop_iteration_exceeded => Callback of stop_iteration_exceeded on IterationInterval(1)
│   ├── wall_time_limit_exceeded => Callback of wall_time_limit_exceeded on IterationInterval(1)
│   ├── nan_checker => Callback of NaNChecker for u on IterationInterval(100)
│   └── total_A => Callback of show_time on IterationInterval(2)
├── Output writers: OrderedDict with no entries
└── Diagnostics: OrderedDict with no entries

Now when simulation runs the simulation the callback is called.

run!(simulation)
[ Info: Initializing simulation...
[ Info: Time is 0 seconds
[ Info:     ... simulation initialization complete (164.907 ms)
[ Info: Executing initial time step...
[ Info:     ... initial time step complete (6.706 seconds).
[ Info: Time is 2 seconds
[ Info: Time is 4 seconds
[ Info: Time is 6 seconds
[ Info: Time is 8 seconds
[ Info: Simulation is stopping after running for 6.925 seconds.
[ Info: Model iteration 10 equals or exceeds stop iteration 10.
[ Info: Time is 10 seconds

We can also use the convenience add_callback!:

add_callback!(simulation, show_time, name=:total_A_via_convenience, IterationInterval(2))

simulation
Simulation of NonhydrostaticModel{CPU, RectilinearGrid}(time = 10 seconds, iteration = 10)
├── Next time step: 1 second
├── Elapsed wall time: 6.925 seconds
├── Wall time per iteration: 692.546 ms
├── Stop time: Inf days
├── Stop iteration : 10.0
├── Wall time limit: Inf
├── Callbacks: OrderedDict with 6 entries:
│   ├── stop_time_exceeded => Callback of stop_time_exceeded on IterationInterval(1)
│   ├── stop_iteration_exceeded => Callback of stop_iteration_exceeded on IterationInterval(1)
│   ├── wall_time_limit_exceeded => Callback of wall_time_limit_exceeded on IterationInterval(1)
│   ├── nan_checker => Callback of NaNChecker for u on IterationInterval(100)
│   ├── total_A => Callback of show_time on IterationInterval(2)
│   └── total_A_via_convenience => Callback of show_time on IterationInterval(2)
├── Output writers: OrderedDict with no entries
└── Diagnostics: OrderedDict with no entries

The keyword argument callsite determines the moment at which the callback is executed. By default, callsite = TimeStepCallsite(), indicating execution after the completion of a timestep. The other options are callsite = TendencyCallsite() that executes the callback after the tendencies are computed but before taking a timestep and callsite = UpdateStateCallsite() that executes the callback within update_state!, after auxiliary variables have been computed (for multi-stage time-steppers, update_state! may be called multiple times per timestep).

As an example of a callback with callsite = TendencyCallsite() , we show below how we can manually add to the tendency field of one of the velocity components. Here we've chosen the :u field using parameters:

using Oceananigans

model = NonhydrostaticModel(grid=RectilinearGrid(size=(1, 1, 1), extent=(1, 1, 1)))

simulation = Simulation(model, Δt=1, stop_iteration=10)

function modify_tendecy!(model, params)
    model.timestepper.Gⁿ[params.c] .+= params.δ
    return nothing
end

simulation.callbacks[:modify_u] = Callback(modify_tendecy!, IterationInterval(1),
                                           callsite = TendencyCallsite(),
                                           parameters = (c = :u, δ = 1))

run!(simulation)
[ Info: Initializing simulation...
[ Info:     ... simulation initialization complete (755.427 μs)
[ Info: Executing initial time step...
[ Info:     ... initial time step complete (2.462 seconds).
[ Info: Simulation is stopping after running for 3.435 seconds.
[ Info: Model iteration 10 equals or exceeds stop iteration 10.

Above there is no forcing at all, but due to the callback the $u$-velocity is increased.

@info model.velocities.u
┌ Info: 1×1×1 Field{Face, Center, Center} on RectilinearGrid on CPU
├── grid: 1×1×1 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 1×1×1 halo
├── boundary conditions: FieldBoundaryConditions
│   └── west: Periodic, east: Periodic, south: Periodic, north: Periodic, bottom: ZeroFlux, top: ZeroFlux, immersed: ZeroFlux
└── data: 3×3×3 OffsetArray(::Array{Float64, 3}, 0:2, 0:2, 0:2) with eltype Float64 with indices 0:2×0:2×0:2
    └── max=10.0, min=10.0, mean=10.0
Example only for illustration purposes

The above is a redundant example since it could be implemented better with a simple forcing function. We include it here though for illustration purposes of how one can use callbacks.

Functions

Callback functions can only take one or two parameters sim - a simulation, or model for state callbacks, and optionally may also accept a NamedTuple of parameters.

Scheduling

The time that callbacks are called at are specified by schedule functions which can be: