(0.91.2) Remove instances of previous_Δt and fix a bug setting `las…

…t_Δt` in RK3 (#3595)

* Remove instances of previous_Δt and fix a bug in RK3

* Add the "stage" time step to Clock

* Update the other constructor

* Bugfix

* Also set the stage dt in AB2

* Fix constructor for AB2

* Bump patch version

* Fix simple bug

* Fix tests that used previous_Δt

* Fix shallow water model constructor

* Update clock.md

* Update Project.toml

* Bugfix

* Another bugbix

* Fix summary and doctests

Project.toml

@@ -1,7 +1,7 @@
 name = "Oceananigans"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 authors = ["Climate Modeling Alliance and contributors"]
-version = "0.91.1"
+version = "0.91.2"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

docs/src/model_setup/background_fields.md

@@ -59,7 +59,7 @@ model.background_fields.velocities.u
 FunctionField located at (Face, Center, Center)
 ├── func: U (generic function with 1 method)
 ├── grid: 1×1×1 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 1×1×1 halo
-├── clock: Clock(time=0 seconds, iteration=0, last_Δt=Inf days)
+├── clock: Clock{Float64, Float64}(time=0 seconds, iteration=0, last_Δt=Inf days)
 └── parameters: nothing
 ```
 
@@ -102,6 +102,6 @@ model.background_fields.tracers.b
 FunctionField located at (Center, Center, Center)
 ├── func: B (generic function with 1 method)
 ├── grid: 1×1×1 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 1×1×1 halo
-├── clock: Clock(time=0 seconds, iteration=0, last_Δt=Inf days)
+├── clock: Clock{Float64, Float64}(time=0 seconds, iteration=0, last_Δt=Inf days)
 └── parameters: (α = 3.14, N = 1.0, f = 0.1)
 ```

docs/src/model_setup/clock.md

@@ -14,15 +14,19 @@ end
 
 ```jldoctest
 julia> clock = Clock(time=0.0)
-Clock{Float64, Float64}: time = 0 seconds, last_Δt = Inf days, iteration = 0, stage = 1
+Clock{Float64, Float64}(time=0 seconds, iteration=0, last_Δt=Inf days)
+├── stage: 1
+└── last_stage_Δt: Inf days
 ```
 
 but can be modified to start the model clock at some other time.
 For example, passing
 
 ```jldoctest
 julia> clock = Clock(time=3600.0)
-Clock{Float64, Float64}: time = 1 hour, last_Δt = Inf days, iteration = 0, stage = 1
+Clock{Float64, Float64}(time=1 hour, iteration=0, last_Δt=Inf days)
+├── stage: 1
+└── last_stage_Δt: Inf days
 ```
 
 to the constructor for `NonhydrostaticModel` causes the simulation
@@ -37,7 +41,9 @@ for example, pass
 julia> using TimesDates
 
 julia> clock = Clock(time=TimeDate(2020))
-Clock{TimesDates.TimeDate, Float64}: time = 2020-01-01T00:00:00, last_Δt = Inf days, iteration = 0, stage = 1
+Clock{TimesDates.TimeDate, Float64}(time=2020-01-01T00:00:00, iteration=0, last_Δt=Inf days)
+├── stage: 1
+└── last_stage_Δt: Inf days
 ```
 
 to `NonhydrostaticModel`. `TimesDates.TimeDate` supports nanosecond resolution and is thus recommended

src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl

@@ -124,8 +124,8 @@ const OnlyParticleTrackingModel = HydrostaticFreeSurfaceModel{TS, E, A, S, G, T,
                  {TS, E, A, S, G, T, V, B, R, F, P<:AbstractLagrangianParticles, U<:PrescribedVelocityFields, C<:NamedTuple{(), Tuple{}}}
 
 function time_step!(model::OnlyParticleTrackingModel, Δt; callbacks = [], kwargs...)
-    model.timestepper.previous_Δt = Δt
     tick!(model.clock, Δt)
+    model.clock.last_Δt = Δt
     step_lagrangian_particles!(model, Δt)
     update_state!(model, callbacks)
 end

src/Models/ShallowWaterModels/shallow_water_model.jl

@@ -112,7 +112,7 @@ Keyword arguments
 function ShallowWaterModel(;
                            grid,
                            gravitational_acceleration,
-                               clock = Clock{eltype(grid), eltype(grid)}(0, Inf, 0, 1),
+                               clock = Clock{eltype(grid)}(time=0),
                   momentum_advection = UpwindBiasedFifthOrder(),
                     tracer_advection = WENO(),
                       mass_advection = WENO(),

src/OutputWriters/checkpointer.jl

@@ -230,6 +230,7 @@ function set!(model, filepath::AbstractString)
         # Update model clock
         model.clock.iteration = checkpointed_clock.iteration
         model.clock.time = checkpointed_clock.time
+        model.clock.last_Δt = checkpointed_clock.last_Δt
     end
 
     return nothing
@@ -260,8 +261,6 @@ end
 set_time_stepper!(timestepper::RungeKutta3TimeStepper, file, model_fields) =
     set_time_stepper_tendencies!(timestepper, file, model_fields)
 
-function set_time_stepper!(timestepper::QuasiAdamsBashforth2TimeStepper, file, model_fields)
+set_time_stepper!(timestepper::QuasiAdamsBashforth2TimeStepper, file, model_fields) =
     set_time_stepper_tendencies!(timestepper, file, model_fields)
-    timestepper.previous_Δt = file["timestepper/previous_Δt"]
-    return nothing
-end
+

src/OutputWriters/output_writer_utils.jl

@@ -160,7 +160,6 @@ end
 function serializeproperty!(file, address, ts::QuasiAdamsBashforth2TimeStepper)
     serializeproperty!(file, address * "/Gⁿ", ts.Gⁿ)
     serializeproperty!(file, address * "/G⁻", ts.G⁻)
-    serializeproperty!(file, address * "/previous_Δt", ts.previous_Δt)
     return nothing
 end
 

src/Simulations/time_step_wizard.jl

@@ -35,7 +35,7 @@ for advective and diffusive Courant-Friedrichs-Lewy (CFL) numbers (`cfl` and `di
 subject to the limits
 
 ```julia
-max(min_Δt, min_change * previous_Δt) ≤ new_Δt ≤ min(max_Δt, max_change * previous_Δt)
+max(min_Δt, min_change * last_Δt) ≤ new_Δt ≤ min(max_Δt, max_change * last_Δt)
 ```
 
 where `new_Δt` is the new time step calculated by the `TimeStepWizard`.

src/Solvers/heptadiagonal_iterative_solver.jl

@@ -21,7 +21,7 @@ mutable struct HeptadiagonalIterativeSolver{G, R, L, D, M, P, PM, PS, I, ST, T,
            iterative_solver :: I
                  state_vars :: ST
                   tolerance :: T
-                previous_Δt :: F
+                    last_Δt :: F
          maximum_iterations :: Int
                     verbose :: Bool
 end
@@ -70,7 +70,7 @@ The matrix constructors are calculated based on the pentadiagonal coeffients pas
 to `matrix_from_coefficients` function.
 
 To allow for variable time step, the diagonal term `- Az / (g * Δt²)` is only added later on
-and it is updated only when the previous time step changes (`previous_Δt != Δt`).
+and it is updated only when the previous time step changes (`last_Δt != Δt`).
 
 Preconditioning is done through the various methods implemented in `Solvers/sparse_preconditioners.jl`.
     
@@ -296,7 +296,7 @@ function solve!(x, solver::HeptadiagonalIterativeSolver, b, Δt)
     arch = architecture(solver.matrix)
 
     # update matrix and preconditioner if time step changes
-    if Δt != solver.previous_Δt
+    if Δt != solver.last_Δt
         constructors = deepcopy(solver.matrix_constructors)
         M = prod(solver.problem_size)
         update_diag!(constructors, arch, M, M, solver.diagonal, Δt, 0)
@@ -308,7 +308,7 @@ function solve!(x, solver::HeptadiagonalIterativeSolver, b, Δt)
                                                          solver.matrix,
                                                          solver.preconditioner_settings)
 
-        solver.previous_Δt = Δt
+        solver.last_Δt = Δt
     end
 
     solver.iterative_solver(x, solver.matrix, b, 

src/TimeSteppers/clock.jl

@@ -13,32 +13,61 @@ The `stage` is updated only for multi-stage time-stepping methods. The `time::T`
 either a number or a `DateTime` object.
 """
 mutable struct Clock{TT, DT}
-         time :: TT
-      last_Δt :: DT
+    time :: TT
+    last_Δt :: DT
+    last_stage_Δt :: DT
     iteration :: Int
-        stage :: Int
+    stage :: Int
 end
 
 """
-    Clock(; time, last_Δt = Inf, iteration=0, stage=1)
+    Clock(; time, last_Δt=Inf, last_stage_Δt=Inf, iteration=0, stage=1)
 
 Returns a `Clock` object. By default, `Clock` is initialized to the zeroth `iteration`
-and first time step `stage` with `last_Δt`.
+and first time step `stage` with `last_Δt=last_stage_Δt=Inf`.
 """
-Clock(; time::TT, last_Δt::DT=Inf, iteration=0, stage=1) where {TT, DT} = Clock{TT, DT}(time, last_Δt, iteration, stage)
+function Clock(; time,
+               last_Δt = Inf,
+               last_stage_Δt = Inf,
+               iteration = 0,
+               stage = 1)
+
+    TT = typeof(time)
+    DT = typeof(last_Δt)
+    last_stage_Δt = convert(DT, last_Δt)
+    return Clock{TT, DT}(time, last_Δt, last_stage_Δt, iteration, stage)
+end
+
 # TODO: when supporting DateTime, this function will have to be extended
 time_step_type(TT) = TT
 
-function Clock{TT}(; time, last_Δt=Inf, iteration=0, stage=1) where TT
+function Clock{TT}(; time,
+                   last_Δt = Inf,
+                   last_stage_Δt = Inf,
+                   iteration = 0,
+                   stage = 1) where TT
+
     DT = time_step_type(TT)
     last_Δt = convert(DT, last_Δt)
-    return Clock{TT, DT}(time, last_Δt, iteration, stage)
+    last_stage_Δt = convert(DT, last_stage_Δt)
+
+    return Clock{TT, DT}(time, last_Δt, last_stage_Δt, iteration, stage)
 end
 
-Base.summary(clock::Clock) = string("Clock(time=$(prettytime(clock.time)), iteration=$(clock.iteration), last_Δt=$(prettytime(clock.last_Δt)))")
+function Base.summary(clock::Clock)
+    TT = typeof(clock.time)
+    DT = typeof(clock.last_Δt)
+    return string("Clock{", TT, ", ", DT, "}",
+                  "(time=", prettytime(clock.time),
+                  ", iteration=", clock.iteration,
+                  ", last_Δt=", prettytime(clock.last_Δt), ")")
+end
 
-Base.show(io::IO, c::Clock{TT, DT}) where {TT, DT} =
-    println(io, "Clock{$TT, $DT}: time = $(prettytime(c.time)), last_Δt = $(prettytime(c.last_Δt)), iteration = $(c.iteration), stage = $(c.stage)")
+function Base.show(io::IO, clock::Clock)
+    return print(io, summary(clock), '\n',
+                 "├── stage: ", clock.stage, '\n',
+                 "└── last_stage_Δt: ", prettytime(clock.last_stage_Δt))
+end
 
 next_time(clock, Δt) = clock.time + Δt
 next_time(clock::Clock{<:AbstractTime}, Δt) = clock.time + Nanosecond(round(Int, 1e9 * Δt))
@@ -61,8 +90,6 @@ function tick!(clock, Δt; stage=false)
 
     tick_time!(clock, Δt)
 
-    clock.last_Δt = Δt
-
     if stage # tick a stage update
         clock.stage += 1
     else # tick an iteration and reset stage
@@ -73,7 +100,11 @@ function tick!(clock, Δt; stage=false)
     return nothing
 end
 
-"Adapt `Clock` to work on the GPU via CUDAnative and CUDAdrv."
-Adapt.adapt_structure(to, clock::Clock) =
-    (time=clock.time, last_Δt=clock.last_Δt, iteration=clock.iteration, stage=clock.stage)
+"""Adapt `Clock` for GPU."""
+Adapt.adapt_structure(to, clock::Clock) = (time          = clock.time,
+                                           last_Δt       = clock.last_Δt,
+                                           last_stage_Δt = clock.last_stage_Δt,
+                                           iteration     = clock.iteration,
+                                           stage         = clock.stage)
+
 

src/TimeSteppers/quasi_adams_bashforth_2.jl

@@ -3,7 +3,6 @@ using Oceananigans.Utils: @apply_regionally, apply_regionally!
 
 mutable struct QuasiAdamsBashforth2TimeStepper{FT, GT, IT} <: AbstractTimeStepper
                   χ :: FT
-        previous_Δt :: FT
                  Gⁿ :: GT
                  G⁻ :: GT
     implicit_solver :: IT
@@ -48,13 +47,10 @@ function QuasiAdamsBashforth2TimeStepper(grid, tracers,
     GT = typeof(Gⁿ)
     χ  = convert(FT, χ)
 
-    return QuasiAdamsBashforth2TimeStepper{FT, GT, IT}(χ, Inf, Gⁿ, G⁻, implicit_solver)
+    return QuasiAdamsBashforth2TimeStepper{FT, GT, IT}(χ, Gⁿ, G⁻, implicit_solver)
 end
 
-function reset!(timestepper::QuasiAdamsBashforth2TimeStepper)
-    timestepper.previous_Δt = Inf
-    return nothing
-end
+reset!(timestepper::QuasiAdamsBashforth2TimeStepper) = nothing
 
 #####
 ##### Time steppping
@@ -70,11 +66,11 @@ Setting `compute_tendencies=false` will not calculate new tendencies
 function time_step!(model::AbstractModel{<:QuasiAdamsBashforth2TimeStepper}, Δt;
                     callbacks=[], euler=false, compute_tendencies=true)
 
+    Δt == 0 && @warn "Δt == 0 may cause model blowup!"
+
     # Be paranoid and update state at iteration 0
     model.clock.iteration == 0 && update_state!(model, callbacks)
 
-    Δt == 0 && @warn "Δt == 0 may cause model blowup!"
-
     ab2_timestepper = model.timestepper
 
     # Change the default χ if necessary, which occurs if:
@@ -83,16 +79,15 @@ function time_step!(model::AbstractModel{<:QuasiAdamsBashforth2TimeStepper}, Δt
     #     need to take an euler step. Note that model.clock.last_Δt is
     #     initialized as Inf
     #   * The user has passed euler=true to time_step!
-    χ₀ = ab2_timestepper.χ
-    euler = euler || (Δt != ab2_timestepper.previous_Δt)
+    euler = euler || (Δt != model.clock.last_Δt)
 
     # If euler, then set χ = -0.5
     minus_point_five = convert(eltype(model.grid), -0.5)
     χ = ifelse(euler, minus_point_five, ab2_timestepper.χ)
 
     # Set time-stepper χ (this is used in ab2_step!, but may also be used elsewhere)
+    χ₀ = ab2_timestepper.χ # Save initial value
     ab2_timestepper.χ = χ
-    ab2_timestepper.previous_Δt = Δt
 
     # Ensure zeroing out all previous tendency fields to avoid errors in
     # case G⁻ includes NaNs. See https://github.com/CliMA/Oceananigans.jl/issues/2259
@@ -103,13 +98,13 @@ function time_step!(model::AbstractModel{<:QuasiAdamsBashforth2TimeStepper}, Δt
         end
     end
 
-    model.clock.iteration == 0 && update_state!(model, callbacks)
-
     ab2_step!(model, Δt) # full step for tracers, fractional step for velocities.
     calculate_pressure_correction!(model, Δt)
     @apply_regionally correct_velocities_and_store_tendencies!(model, Δt)
 
     tick!(model.clock, Δt)
+    model.clock.last_Δt = Δt
+    model.clock.last_stage_Δt = Δt # just one stage
     update_state!(model, callbacks; compute_tendencies)
     step_lagrangian_particles!(model, Δt)
 
@@ -176,3 +171,4 @@ Time step velocity fields via the 2nd-order quasi Adams-Bashforth method
 end
 
 @kernel ab2_step_field!(::FunctionField, Δt, χ, Gⁿ, G⁻) = nothing
+

src/TimeSteppers/runge_kutta_3.jl

@@ -108,6 +108,7 @@ function time_step!(model::AbstractModel{<:RungeKutta3TimeStepper}, Δt; callbac
     pressure_correct_velocities!(model, first_stage_Δt)
 
     tick!(model.clock, first_stage_Δt; stage=true)
+    model.clock.last_stage_Δt = first_stage_Δt
     store_tendencies!(model)
     update_state!(model, callbacks)
     step_lagrangian_particles!(model, first_stage_Δt)
@@ -122,6 +123,7 @@ function time_step!(model::AbstractModel{<:RungeKutta3TimeStepper}, Δt; callbac
     pressure_correct_velocities!(model, second_stage_Δt)
 
     tick!(model.clock, second_stage_Δt; stage=true)
+    model.clock.last_stage_Δt = second_stage_Δt
     store_tendencies!(model)
     update_state!(model, callbacks)
     step_lagrangian_particles!(model, second_stage_Δt)
@@ -135,10 +137,15 @@ function time_step!(model::AbstractModel{<:RungeKutta3TimeStepper}, Δt; callbac
     calculate_pressure_correction!(model, third_stage_Δt)
     pressure_correct_velocities!(model, third_stage_Δt)
 
-    # This formulation of the final time-step reduces the accumulation of round-off error when Δt
-    # is added to model.clock.time.
+    # This adjustment of the final time-step reduces the accumulation of
+    # round-off error when Δt is added to model.clock.time. Note that we still use 
+    # third_stage_Δt for the substep, pressure correction, and Lagrangian particles step.
     corrected_third_stage_Δt = tⁿ⁺¹ - model.clock.time
+
     tick!(model.clock, corrected_third_stage_Δt)
+    model.clock.last_stage_Δt = corrected_third_stage_Δt
+    model.clock.last_Δt = Δt
+
     update_state!(model, callbacks; compute_tendencies)
     step_lagrangian_particles!(model, third_stage_Δt)
 

test/regression_tests/ocean_large_eddy_simulation_regression_test.jl

@@ -102,7 +102,7 @@ function run_ocean_large_eddy_simulation_regression_test(arch, grid_type, closur
     model.clock.iteration = spinup_steps
 
     update_state!(model; compute_tendencies = true)
-    model.timestepper.previous_Δt = Δt
+    model.clock.last_Δt = Δt
 
     for n in 1:test_steps
         time_step!(model, Δt, euler=false)

test/regression_tests/rayleigh_benard_regression_test.jl

@@ -125,7 +125,7 @@ function run_rayleigh_benard_regression_test(arch, grid_type)
     # Step the model forward and perform the regression test
     update_state!(model)
 
-    model.timestepper.previous_Δt = Δt
+    model.clock.last_Δt = Δt
 
     for n in 1:test_steps
         time_step!(model, Δt, euler=false)

test/test_checkpointer.jl

@@ -102,7 +102,7 @@ function run_checkpointer_tests(true_model, test_model, Δt)
     test_model_equality(test_model, checkpointed_model)
 
     # This only applies to QuasiAdamsBashforthTimeStepper:
-    @test test_model.timestepper.previous_Δt == checkpointed_model.timestepper.previous_Δt
+    @test test_model.clock.last_Δt == checkpointed_model.clock.last_Δt
 
     #####
     ##### Test pickup from explicit checkpoint path