Simulations: managing model time-stepping

Simulation orchestrates model time-stepping loops that include stop conditions, writing output, and the execution of "callbacks" that can do everything from monitoring simulation progress to editing the model state.

Scripts for numerical experiments can be broken into four parts:

Definition of the "grid" and physical domain for the numerical experiment
Configuration of physics and model construction
Building a Simulation time-stepping loop with callbacks and output writers
Calling run! to integrate the model forward in time.

Simulation is the final boss object that users interact with in order in the process of performing a numerical experiment.

A simple simulation

A minimal example illustrates how Simulation works:

using Oceananigans

grid = RectilinearGrid(size=(4, 4, 4), extent=(1, 1, 1))
model = NonhydrostaticModel(; grid)
simulation = Simulation(model; Δt=7, stop_iteration=6)
run!(simulation)

simulation

Simulation of NonhydrostaticModel{CPU, RectilinearGrid}(time = 42 seconds, iteration = 6)
├── Next time step: 7 seconds
├── run_wall_time: 8.993 seconds
├── run_wall_time / iteration: 1.499 seconds
├── stop_time: Inf days
├── stop_iteration: 6.0
├── wall_time_limit: Inf
├── minimum_relative_step: 0.0
├── callbacks: OrderedDict with 4 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)
└── output_writers: OrderedDict with no entries

A more complicated setup might invoke multiple callbacks:

using Oceananigans

grid = RectilinearGrid(size=(4, 4, 4), extent=(1, 1, 1))
model = NonhydrostaticModel(; grid)
simulation = Simulation(model; Δt=7, stop_time=14)

print_progress(sim) = @info string("Iteration: ", iteration(sim), ", time: ", time(sim))
add_callback!(simulation, print_progress, IterationInterval(2))

declare_time(sim) = @info string("The simulation has been running for ", prettytime(sim.run_wall_time), "!")
add_callback!(simulation, declare_time, TimeInterval(10))

run!(simulation)

[ Info: Initializing simulation...
[ Info: Iteration: 0, time: 0.0
[ Info: The simulation has been running for 0 seconds!
[ Info:     ... simulation initialization complete (35.992 ms)
[ Info: Executing initial time step...
[ Info:     ... initial time step complete (4.774 ms).
[ Info: Iteration: 2, time: 10.0
[ Info: The simulation has been running for 41.027 ms!
[ Info: Simulation is stopping after running for 41.692 ms.
[ Info: Simulation time 14 seconds equals or exceeds stop time 14 seconds.

Simulation book-keeps the total iterations performed, the next Δt, and the lists of callbacks and output_writers. Simulations can be continued, which is helpful for interactive work:

simulation.stop_time = 42
run!(simulation)

[ Info: Initializing simulation...
[ Info:     ... simulation initialization complete (108.582 μs)
[ Info: Executing initial time step...
[ Info: Iteration: 4, time: 21.0
[ Info:     ... initial time step complete (599.849 μs).
[ Info: The simulation has been running for 849.925 μs!
[ Info: Iteration: 6, time: 31.0
[ Info: The simulation has been running for 1.886 ms!
[ Info: Iteration: 8, time: 41.0
[ Info: Simulation is stopping after running for 2.904 ms.
[ Info: Simulation time 42 seconds equals or exceeds stop time 42 seconds.

Stop criteria and time-step control

The Simulation constructor accepts three stopping conditions:

stop_iteration: maximum number of steps
stop_time: maximum model clock time (same units as the model's clock)
wall_time_limit: maximum wall-clock seconds before the run aborts

simulation = Simulation(model; Δt=0.1, stop_time=10, stop_iteration=10000, wall_time_limit=30)

Simulation of NonhydrostaticModel{CPU, RectilinearGrid}(time = 42 seconds, iteration = 9)
├── Next time step: 100 ms
├── run_wall_time: 0 seconds
├── run_wall_time / iteration: 0 seconds
├── stop_time: 10 seconds
├── stop_iteration: 10000.0
├── wall_time_limit: 30.0
├── minimum_relative_step: 0.0
├── callbacks: OrderedDict with 4 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)
└── output_writers: OrderedDict with no entries

Callback basics

Callback executes arbitrary code on a schedule. Callbacks can be used to monitor simulation progress, compute diagnostics, and adjust the course of a simulation. To illustrate a hierarchy of callbacks we use a simulation with a forced passive tracer:

using Oceananigans

grid = RectilinearGrid(size=(4, 4, 4), extent=(1, 1, 1))

c_source(x, y, z, t, c) = 0.1 * c
c_forcing = Forcing(c_source, field_dependencies=:c)
model = NonhydrostaticModel(; grid, tracers=:c, forcing=(; c=c_forcing))
simulation = Simulation(model; Δt=0.1, stop_time=10)

# Add a callback that prints progress
print_progress(sim) = @info "Iter $(iteration(sim)): $(prettytime(sim))"
add_callback!(simulation, print_progress, IterationInterval(25), name=:progress)
run!(simulation)

[ Info: Initializing simulation...
[ Info: Iter 0: 0 seconds
[ Info:     ... simulation initialization complete (1.380 seconds)
[ Info: Executing initial time step...
[ Info:     ... initial time step complete (4.831 seconds).
[ Info: Iter 25: 2.500 seconds
[ Info: Iter 50: 5.000 seconds
[ Info: Iter 75: 7.500 seconds
[ Info: Iter 100: 10.000 seconds
[ Info: Simulation is stopping after running for 6.395 seconds.
[ Info: Simulation time 10 seconds equals or exceeds stop time 10 seconds.

Naming callbacks

Callbacks can optionally be assigned a name via add_callback! or by adding the callback manually to the callbacks dictionary: simulation.callbacks[:mine] = my_callback. Names can be used to identify, modify, or delete callbacks from Simulation.

Callbacks: for stopping a simulation

To spark your imagination, consider that callbacks can be used to implement arbitrary stopping criteria. As an example we consider a stopping criteria based on the magnitude of a tracer:

using Printf

set!(model, c=1)

function stop_simulation(sim)
    if maximum(sim.model.tracers.c) >= 2
        sim.running = false
        @info "The simulation is stopping because max(c) >= 2."
    end
    return nothing
end

# The default schedule is IterationInterval(1)
add_callback!(simulation, stop_simulation)

function print_progress(sim)
    max_c = maximum(sim.model.tracers.c)
    @info @sprintf("Iter %d: t = %s, max(c): %.2f",
                   iteration(sim), prettytime(sim), max_c)
    return nothing
end

add_callback!(simulation, stop_simulation, IterationInterval(10), name=:progress)
simulation.stop_time += 10
run!(simulation)

[ Info: Initializing simulation...
[ Info:     ... simulation initialization complete (5.772 ms)
[ Info: Executing initial time step...
[ Info:     ... initial time step complete (289.591 ms).
[ Info: The simulation is stopping because max(c) >= 2.

Callback execution order

Callbacks are executed in the order they appear within the callbacks OrderedDict, which in turn corresponds to the order they are added.

Adaptive time-stepping with `TimeStepWizard`

The time-step can be changed by modifying simulation.Δt. To decrease the computational cost of simulations of flows that significantly grow or decay in time, users may invoke a special callback called TimeStepWizard (which is associated with a special helper function conjure_time_step_wizard!). TimeStepWizard monitors the advective and diffusive Courant numbers, increasing or decreasing simulation.Δt to keep the time-step near its maximum stable value while respecting bounds such as max_change or max_Δt.

ϵ(x, y, z) = 1e-2
set!(model, c=1, u=ϵ, v=ϵ, w=ϵ)

conjure_time_step_wizard!(simulation, cfl=0.7)
simulation

Simulation of NonhydrostaticModel{CPU, RectilinearGrid}(time = 17.000 seconds, iteration = 171)
├── Next time step: 100 ms
├── run_wall_time: 420.305 ms
├── run_wall_time / iteration: 2.458 ms
├── stop_time: 20 seconds
├── stop_iteration: Inf
├── wall_time_limit: Inf
├── minimum_relative_step: 0.0
├── callbacks: OrderedDict with 7 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)
│   ├── progress => Callback of stop_simulation on IterationInterval(10)
│   ├── callback1 => Callback of stop_simulation on IterationInterval(1)
│   └── time_step_wizard => Callback of TimeStepWizard(cfl=0.7, max_Δt=Inf, min_Δt=0.0) on IterationInterval(10)
└── output_writers: OrderedDict with no entries

A TimeStepWizard has been added to the list of simulation.callbacks. TimeStepWizard is authorized to modify the time-step as the simulation progresses (every 10 iterations by default, but this can be modified):

print_progress(sim) = @info string("Iter: ", iteration(sim), ": Δt = ", prettytime(sim.Δt))
add_callback!(simulation, print_progress, IterationInterval(10), name=:progress)
simulation.stop_time += 1
run!(simulation)

[ Info: Initializing simulation...
[ Info:     ... simulation initialization complete (411.197 μs)
[ Info: Executing initial time step...
[ Info:     ... initial time step complete (1.067 ms).
[ Info: Iter: 180: Δt = 100 ms
[ Info: Iter: 190: Δt = 110.000 ms
[ Info: Iter: 200: Δt = 121.000 ms
[ Info: Simulation is stopping after running for 632.696 ms.
[ Info: Simulation time 21 seconds equals or exceeds stop time 21 seconds.

Time-step alignment with scheduled events

By default, Simulation "aligns" Δt so that events which are scheduled by time (such as TimeInterval and SpecifiedTimes) occur exactly on schedule. Time-step alignment may be disabled by setting align_time_step = false in the Simulation constructor.

print_progress(sim) = @info "At t = $(time(sim)), iter = $(iteration(sim))"
add_callback!(simulation, print_progress, TimeInterval(0.2), name=:progress)
simulation.stop_time += 0.6
run!(simulation)

[ Info: Initializing simulation...
[ Info:     ... simulation initialization complete (8.535 ms)
[ Info: Executing initial time step...
[ Info:     ... initial time step complete (5.565 ms).
[ Info: At t = 21.2, iter = 208
[ Info: At t = 21.4, iter = 210
[ Info: Simulation is stopping after running for 32.616 ms.
[ Info: Simulation time 21.600 seconds equals or exceeds stop time 21.600 seconds.
[ Info: At t = 21.6, iter = 212

Minimum relative step

Due to round-off error, time-step alignment can produce very small Δt close to machine epsilon –- for example, when TimeInterval is a constant multiple of a fixed Δt. In some cases, Δt close to machine epsilon is undesirable: for example, with NonhydrostaticModel a machine epsilon Δt will produce a pressure field that is polluted by numerical error due to its pressure correction algorithm. To mitigate issues associated with very small time-steps, the Simulation constructor accepts a minimum_relative_step argument (a typical choice is 1e-9). When minimum_relative_step > 0, a time-step will be skipped (instead, the clock is simply reset) if an aligned time-step is less than minimum_relative_step * previous_Δt.

Callback "callsites"

The callsite keyword lets callbacks hook into different parts of the time-stepping cycle. For example, use callsite = TendencyCallsite() to modify tendencies before a step, or UpdateStateCallsite() to react after auxiliary variables update. See the Callbacks page for more customization patterns, including parameterized callbacks and state callbacks.

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 (6.081 ms)
[ Info: Executing initial time step...
[ Info:     ... initial time step complete (594.339 ms).
[ Info: Simulation is stopping after running for 608.046 ms.
[ Info: Model iteration 10 equals or exceeds stop iteration 10.

Introduction to writing output

Output writers live in the ordered dictionary simulation.output_writers. Each writer pairs outputs with a schedule describing when to export them and a backend (NetCDF, JLD2, or Checkpointer) describing how to serialize them. Time-step alignment ensures the writer's schedule is honored without user intervention.

using NCDatasets

fields = Dict("u" => simulation.model.velocities.u)

simulation.output_writers[:surface_slice] = NetCDFWriter(simulation.model, fields;
                                                         filename = "demo.nc",
                                                         schedule = TimeInterval(30),
                                                         indices = (:, :, grid.Nz))
simulation.output_writers

OrderedCollections.OrderedDict{Symbol, Oceananigans.AbstractOutputWriter} with 1 entry:
  :surface_slice => NetCDFWriter scheduled on TimeInterval(30 seconds):…

During run!, Oceananigans calls each writer whenever its schedule actuates. More comprehensive information may be found on the output writers page.

Putting it together

Combining callbacks, output writers, and adaptive time-stepping turns a few lines of model code into a robust workflow:

using Oceananigans.Units: hours, minutes

model = NonhydrostaticModel(; grid, tracers=:T)
simulation = Simulation(model, Δt=20, stop_time=2hours)

progress(sim) = @info "t = $(prettytime(sim)), Δt = $(prettytime(sim.Δt))"
add_callback!(simulation, progress, IterationInterval(10))
conjure_time_step_wizard!(simulation, cfl=0.8)

simulation.output_writers[:snapshots] = JLD2Writer(simulation.model, simulation.model.velocities;
                                                   filename = "snapshots.jld2",
                                                   schedule = TimeInterval(30minutes))

run!(simulation)

[ Info: Initializing simulation...
[ Info: t = 0 seconds, Δt = 20 seconds
[ Info:     ... simulation initialization complete (2.049 seconds)
[ Info: Executing initial time step...
[ Info:     ... initial time step complete (4.329 seconds).
[ Info: t = 3.667 minutes, Δt = 22 seconds
[ Info: t = 7.700 minutes, Δt = 24.200 seconds
[ Info: t = 12.137 minutes, Δt = 26.620 seconds
[ Info: t = 17.017 minutes, Δt = 29.282 seconds
[ Info: t = 22.385 minutes, Δt = 32.210 seconds
[ Info: t = 28.291 minutes, Δt = 35.431 seconds
[ Info: t = 34.547 minutes, Δt = 38.974 seconds
[ Info: t = 41.692 minutes, Δt = 42.872 seconds
[ Info: t = 49.552 minutes, Δt = 47.159 seconds
[ Info: t = 58.198 minutes, Δt = 51.875 seconds
[ Info: t = 1.127 hours, Δt = 57.062 seconds
[ Info: t = 1.301 hours, Δt = 1.046 minutes
[ Info: t = 1.493 hours, Δt = 1.151 minutes
[ Info: t = 1.690 hours, Δt = 1.266 minutes
[ Info: t = 1.922 hours, Δt = 1.392 minutes
[ Info: Simulation is stopping after running for 6.573 seconds.
[ Info: Simulation time 2 hours equals or exceeds stop time 2 hours.

For more recipes continue with the page on schedules.