Oceananigans.jl: Redesigning `Simulation` and `run!`

[glwagner]

@ali-ramadhan and I have discussed a potential redesign of Simulation and run! that has several key features:

1. Coalesce all of the "callbacks" (arbitrary functions that are executed during a time-stepping loop) other than OutputWriters into a single list. Current objects that we can classify / redesign as callbacks are: stop criteria, TimeStepWizard, and diagnostics. All callbacks are required to possess callback.schedule, and we can provide convenience objects for coupling simple callback functions to AbstractSchedule.

2. Within the time-stepping loop, execute simulation.callbacks prior to writing output. This ensures that data calculated in a callback can be output during the same time-step, such as WindowedTimeAverage and other non-local-in-time output.

3. Wrap the time-stepping loop inside a try / catch block and throw exceptions to stop a simulation. This generalizes the concept of stopping a simulation and also means that a simulation can be stopped inside any callback. Further, when an AbstractStopException is called we will loop over the OutputWriter callbacks a final time, passing the exception into the OutputWriter callback functions. This allows output behavior specialized on the type of exception. For example:

• If NaNsDetected is thrown, no output will be written.
• If WallTimeExceeded is thrown, the checkpointer may write output.

1. simulation.Δt becomes a number corresponding to the next time-step, always (rather than sometimes being a TimeStepWizard). The TimeStepWizard callback changes this number on its schedule. Otherwise, the time-step is held constant. This changes the API, since the initial time-step must now be provided to Simulation.

2. (Somewhat unrelated, but enabled by the new pattern) Use a new function align!(simulation.Δt, writer.schedule, simulation.model) to adjust a subsequent time-step if output writing is scheduled. This ensures output writing on TimeIntervals will always be on schedule rather than chronically late as it is now. Since output is called after all the other callbacks, the output writers get the final say as to the next time-step.

These changes will break the existing API. Notably, we'll use

push!(simulation.callbacks, TimeStepWizard(cfl=1))

rather than setting simulation.Δt to be a TimeStepWizard. This is probably an improvement. We can still provide the keywords stop_time and stop_iteration in the Simulation constructor as a convenience; however rather than being properties of Simulation we will create callback objects that get scheduled every iteration and add them to simulation.callbacks.

I think we should also provide a few other features, like to ability to pass Δt to run!, perhaps along with a few other run!-specific objects.

[glwagner]

NaNChecker will also have to be moved from diagnostics to callbacks, see discussion on #1198.

We also may want to make the list of callbacks an OrderedDict, so that we can write something like

simulation.callbacks[:NaNChecker].schedule = IterationInterval(1)

If can overload getproperty for these custom OrderedDict we can also support the syntax

simulation.callbacks.NaNChecker.schedule = IterationInterval(1)

[ali-ramadhan]

I'll add a sixth feature (from #1251):

1. Time step should be updated before simulation callback/progress function so that we print the actual next time step in the simulation callback/progress function. Even better would be to define a new function next_time_step(simulation) since the next time step might be shortened due to alignment.

[ali-ramadhan]

I'll also add a seventh feature (from #1250):

1. Run simulation callback/progress function at iteration 0. Right now the callback/progress function is only called for the first time when iteration = iteration_interval but I think we actually want to run the progress function at iteration 0 as it helps provide more feedback to the user at a time of heavy compilation.

@glwagner also suggested that all the callbacks should be initialized at the beginning of run.

[glwagner]

To clarify, there's two possibilities:

1. When run! begins, execute [initialize!(callback) for callback in simulation.callbacks]
2. When run! begins, execute the callbacks themselves.

or both.

Which do we want?

There won't be a progress function if this issue is resolved; we'll have a generic list of callbacks that includes default stop criteria as well as user-defined callbacks and stop criteria.

[ali-ramadhan]

I was envisioning option 2 (or both if we feel that initialize!(callback) is an important option).

An important practical reason to have option 2 is that some callbacks will only be called very infrequently so it's better to quickly find out that there's a typo or mistake in your callbacks before the simulation runs for 1,000,000 iterations if the callback's schedule is IterationInterval(1_000_000).

[glwagner]

That makes sense. We can also let users assume that "iteration = 0" means "initialization". This doesn't cover cases where simulations are run from iterations other than 0 though.