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
├── 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)
│ └── 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 (1.957 seconds)
[ Info: Executing initial time step...
[ Info: ... initial time step complete (6.910 seconds).
[ Info: Time is 2.000 seconds
[ Info: Time is 4 seconds
[ Info: Time is 6 seconds
[ Info: Time is 8 seconds
[ Info: Simulation is stopping after running for 8.958 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: 8.961 seconds
├── Wall time per iteration: 896.093 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)
│ ├── 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
using Oceananigans: TendencyCallsite
model = NonhydrostaticModel(grid=RectilinearGrid(size=(1, 1, 1), extent=(1, 1, 1)))
simulation = Simulation(model, Δt=1, stop_iteration=10)
function modify_tendency!(model, params)
model.timestepper.Gⁿ[params.c] .+= params.δ
return nothing
end
simulation.callbacks[:modify_u] = Callback(modify_tendency!, IterationInterval(1),
callsite = TendencyCallsite(),
parameters = (c = :u, δ = 1))
run!(simulation)
[ Info: Initializing simulation...
[ Info: ... simulation initialization complete (823.020 μs)
[ Info: Executing initial time step...
[ Info: ... initial time step complete (2.817 seconds).
[ Info: Simulation is stopping after running for 3.216 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
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:
IterationInterval
: runs everyn
iterationsTimeInterval
: runs everyn
s of model run timeSpecifiedTimes
: runs at the specified timesWallTimeInterval
: runs everyn
s of wall time