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

grid = RectilinearGrid(size=(1, 1, 1), extent=(1, 1, 1))
model = NonhydrostaticModel(grid)
simulation = Simulation(model, Δt=1, stop_iteration=10)

show_time(sim) = @info "Time is $(prettytime(sim.model.clock.time))"
simulation.callbacks[:pretty_basic] = Callback(show_time, IterationInterval(2))

simulation

Simulation of NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0)
├── Next time step: 1 second
├── run_wall_time: 0 seconds
├── run_wall_time / iteration: NaN days
├── stop_time: Inf days
├── stop_iteration: 10.0
├── wall_time_limit: Inf
├── minimum_relative_step: 0.0
├── 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)
│   └── pretty_basic => Callback of show_time on IterationInterval(2)
└── output_writers: 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 (107.573 ms)
[ Info: Executing initial time step...
[ Info:     ... initial time step complete (476.112 ms).
[ 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 584.118 ms.
[ 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, IterationInterval(2))

simulation

Simulation of NonhydrostaticModel{CPU, RectilinearGrid}(time = 10 seconds, iteration = 10)
├── Next time step: 1 second
├── run_wall_time: 584.176 ms
├── run_wall_time / iteration: 58.418 ms
├── stop_time: Inf days
├── stop_iteration: 10.0
├── wall_time_limit: Inf
├── minimum_relative_step: 0.0
├── 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)
│   ├── pretty_basic => Callback of show_time on IterationInterval(2)
│   └── callback1 => Callback of show_time on IterationInterval(2)
└── output_writers: 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. This is the only callsite that is owned by Simulation.

Other callsite options are callsite = TendencyCallsite() and callsite = UpdateStateCallsite(). When these callsites are used, the callback is a function of the model in question rather than Simulation. These callsites must be implemented by models: for NonhydrostaticModel and HydrostaticFreeSurfaceModel, TendencyCallsite callbacks are executed after computing tendencies but before taking a timestep. UpdateStateCallsite callbacks are executed within update_state!, after auxiliary variables have been computed (for multi-stage time-steppers, update_state! may be called multiple times per timestep).

TendencyCallsite() callbacks act on the momentum tendencies. HydrostaticFreeSurfaceModel computes its tracer tendencies later in each substep, after the velocities and the free surface are stepped, so a callback that modifies them needs callsite = TracerTendencyCallsite() (imported with using Oceananigans: TracerTendencyCallsite). NonhydrostaticModel and ShallowWaterModel compute all their tendencies together and execute both kinds of callbacks after them. On multi-region grids, HydrostaticFreeSurfaceModel executes TendencyCallsite callbacks once per region, with that region's model, and TracerTendencyCallsite callbacks once, with the whole model.

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
using Oceananigans: TendencyCallsite

grid = RectilinearGrid(size=(1, 1, 1), extent=(1, 1, 1))
model = NonhydrostaticModel(grid)
simulation = Simulation(model, Δt=1, stop_iteration=10)

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

add_callback!(simulation,
              modify_tendency!,
              callsite = TendencyCallsite(),
              parameters = (c = :u, δ = 1))


run!(simulation)

[ Info: Initializing simulation...
[ Info:     ... simulation initialization complete (4.773 ms)
[ Info: Executing initial time step...
[ Info:     ... initial time step complete (217.941 ms).
[ Info: Simulation is stopping after running for 223.104 ms.
[ Info: Model iteration 10 equals or exceeds stop iteration 10.

Above there is no forcing at all, but due to the callback the -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: Nothing
│ └── 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

Callbacks rely on the schedules described in Scheduling callbacks and output writers. The most common choices are:

• IterationInterval: runs every n iterations

• TimeInterval: runs every n seconds of model time

• SpecifiedTimes: runs at the specified times

• WallTimeInterval: runs every n seconds of wall time