Checkpointer updates

src/Models/Models.jl

@@ -21,6 +21,7 @@ using Oceananigans.Utils: Time
 
 import Oceananigans: initialize!
 import Oceananigans.Architectures: architecture
+import Oceananigans.TimeSteppers: reset!, set_clock!
 import Oceananigans.Solvers: iteration
 import Oceananigans.Simulations: timestepper
 import Oceananigans.TimeSteppers: reset!, set_clock!

src/OutputWriters/checkpointer.jl

@@ -1,24 +1,34 @@
 using Glob
 
 using Oceananigans
-using Oceananigans: fields, prognostic_fields
+using Oceananigans: AbstractModel, fields, prognostic_fields
 using Oceananigans.Fields: offset_data
 using Oceananigans.TimeSteppers: QuasiAdamsBashforth2TimeStepper
 
 import Oceananigans.Fields: set!
 
-mutable struct Checkpointer{T, P} <: AbstractOutputWriter
+mutable struct Checkpointer{T} <: AbstractOutputWriter
     schedule :: T
     dir :: String
     prefix :: String
-    properties :: P
     overwrite_existing :: Bool
     verbose :: Bool
     cleanup :: Bool
 end
 
 required_checkpoint_properties(model) = [:grid, :clock]
 
+# Certain properties are required for `set!` to pickup from a checkpoint.
+function required_checkpointed_properties(model)
+    properties = [:grid, :clock]
+
+    if has_ab2_timestepper(model)
+        push!(properties, :timestepper)
+    end
+
+    return properties
+end
+
 """
     Checkpointer(model;
                  schedule,
@@ -30,16 +40,18 @@ required_checkpoint_properties(model) = [:grid, :clock]
                  properties = required_checkpoint_properties(model))
 
 Construct a `Checkpointer` that checkpoints the model to a JLD2 file on `schedule.`
-The `model.clock.iteration` is included in the filename to distinguish between multiple checkpoint files.
+The `model.clock.iteration` is included in the filename to distinguish between multiple
+checkpoint files.
 
-To restart or "pickup" a model from a checkpoint, specify `pickup = true` when calling `run!`, ensuring
-that the checkpoint file is in directory `dir`. See [`run!`](@ref) for more details.
+To restart or "pickup" a model from a checkpoint, specify `pickup = true` when
+calling `run!`, ensuring that the checkpoint file is in directory `dir`.
+See [`run!`](@ref) for more details.
 
 Note that extra model `properties` can be specified, but removing crucial properties
-such as `:timestepper` will render restoring from the checkpoint impossible.
+such as `:timestepper` might render restoring from the checkpoint impossible.
 
 The checkpointer attempts to serialize as much of the model to disk as possible,
-but functions or objects containing functions cannot be serialized at this time.
+but note that functions or objects containing functions cannot be serialized.
 
 Keyword arguments
 =================
@@ -93,7 +105,7 @@ function Checkpointer(model; schedule,
 
     mkpath(dir)
 
-    return Checkpointer(schedule, dir, prefix, properties, overwrite_existing, verbose, cleanup)
+    return Checkpointer(schedule, dir, prefix, overwrite_existing, verbose, cleanup)
 end
 
 #####
@@ -158,32 +170,43 @@ end
 ##### Writing checkpoints
 #####
 
-function write_output!(c::Checkpointer, model)
+function write_output!(c::Checkpointer, model, addr=checkpointer_address(model))
     filepath = checkpoint_path(model.clock.iteration, c)
     c.verbose && @info "Checkpointing to file $filepath..."
-    addr = checkpointer_address(model)
 
     t1 = time_ns()
 
-    jldopen(filepath, "w") do file
-        file["$addr/checkpointed_properties"] = c.properties
-        serializeproperties!(file, model, c.properties, addr)
-        model_fields = prognostic_fields(model)
-        field_names = keys(model_fields)
-        for name in field_names
-            full_address = "$addr/$name"
-            serializeproperty!(file, full_address, model_fields[name])
-        end
-    end
+    write_output!(c, model, filepath, "w")
 
     t2, sz = time_ns(), filesize(filepath)
+
     c.verbose && @info "Checkpointing done: time=$(prettytime((t2 - t1) * 1e-9)), size=$(pretty_filesize(sz))"
 
     c.cleanup && cleanup_checkpoints(c)
 
     return nothing
 end
 
+function write_output!(c, model, filepath::AbstractString, mode::AbstractString;
+                       properties = default_checkpointed_properties(model))
+    @show properties
+    @show model
+
+    properties = validate_checkpointed_properties(model, properties)
+    addr = checkpointer_address(model)
+
+    jldopen(filepath, mode) do file
+        file["$addr/checkpointed_properties"] =
+        serializeproperties!(file, model, properties, addr)
+        model_fields = prognostic_fields(model)
+        field_names = keys(model_fields)
+        for name in field_names
+            full_address = "$addr/$name"
+            serializeproperty!(file, full_address, model_fields[name])
+        end
+    end
+end
+
 function cleanup_checkpoints(checkpointer)
     filepaths = glob(checkpoint_superprefix(checkpointer.prefix) * "*.jld2", checkpointer.dir)
     latest_checkpoint_filepath = latest_checkpoint(checkpointer, filepaths)
@@ -197,12 +220,13 @@ end
 
 # Should this go in Models?
 """
-    set!(model, filepath::AbstractString)
+    set!(model::AbstractModel, filepath::AbstractString)
 
 Set data in `model.velocities`, `model.tracers`, `model.timestepper.Gⁿ`, and
 `model.timestepper.G⁻` to checkpointed data stored at `filepath`.
 """
-function set!(model, filepath::AbstractString)
+function set!(model::AbstractModel, filepath::AbstractString)
+
     addr = checkpointer_address(model)
 
     jldopen(filepath, "r") do file
@@ -225,7 +249,8 @@ function set!(model, filepath::AbstractString)
             end
         end
 
-        set_time_stepper!(model.timestepper, model.architecture, file, model_fields, addr)
+        set_time_stepper!(model.timestepper, file, model_fields, addr)
+        @show model.timestepper.Gⁿ[:u]
 
         if !isnothing(model.particles)
             copyto!(model.particles.properties, file["$addr/particles"])
@@ -251,7 +276,11 @@ function set_time_stepper_tendencies!(timestepper, arch, file, model_fields, add
             parent_data = on_architecture(arch, file["$addr/timestepper/Gⁿ/$name/data"])
 
             tendencyⁿ_field = timestepper.Gⁿ[name]
+
             @apply_regionally copyto!(parent(tendencyⁿ_field), parent_data)
+            if name==:u
+                @show tendencyⁿ_field
+            end
 
             # Tendency "n-1"
             parent_data = on_architecture(arch, file["$addr/timestepper/G⁻/$name/data"])

src/Simulations/run.jl

@@ -132,7 +132,7 @@ function time_step!(sim::Simulation)
         sim.Δt
     end
 
-    initial_time_step = !(sim.initialized)
+    @show initial_time_step = !(sim.initialized)
     initial_time_step && initialize!(sim)
 
     if initial_time_step && sim.verbose
@@ -191,13 +191,13 @@ we_want_to_pickup(pickup::String) = true
 we_want_to_pickup(pickup) = throw(ArgumentError("Cannot run! with pickup=$pickup"))
 
 """
-    initialize!(sim::Simulation, pickup=false)
+    initialize!(sim::Simulation)
 
 Initialize a simulation:
 
-- Update the auxiliary state of the simulation (filling halo regions, computing auxiliary fields)
-- Evaluate all diagnostics, callbacks, and output writers if sim.model.clock.iteration == 0
-- Add diagnostics that "depend" on output writers
+- Update the auxiliary state of the simulation (filling halo regions, computing auxiliary fields).
+- Evaluate all diagnostics, callbacks, and output writers if `sim.model.clock.iteration == 0`.
+- Add diagnostics that "depend" on output writers.
 """
 function initialize!(sim::Simulation)
     if sim.verbose
@@ -207,7 +207,7 @@ function initialize!(sim::Simulation)
 
     model = sim.model
     initialize!(model)
-    update_state!(model)
+    update_state!(model, compute_tendencies=true)
 
     # Output and diagnostics initialization
     [add_dependencies!(sim.diagnostics, writer) for writer in values(sim.output_writers)]
@@ -253,4 +253,3 @@ function initialize!(sim::Simulation)
 
     return nothing
 end
-

src/Simulations/simulation.jl

@@ -4,7 +4,7 @@ using Oceananigans.DistributedComputations: Distributed, all_reduce
 using Oceananigans.OutputWriters: JLD2Writer, NetCDFWriter
 
 import Oceananigans.Utils: prettytime
-import Oceananigans.TimeSteppers: reset!
+import Oceananigans.TimeSteppers: reset!, set_clock!
 import Oceananigans.OutputWriters: write_output!
 import Oceananigans.Solvers: iteration
 
@@ -34,6 +34,7 @@ end
                stop_iteration = Inf,
                stop_time = Inf,
                wall_time_limit = Inf,
+               align_time_step = true,
                minimum_relative_step = 0)
 
 Construct a `Simulation` for a `model` with time step `Δt`.
@@ -48,6 +49,14 @@ Keyword arguments
 
 - `stop_time`: Stop the simulation once this much model clock time has passed. Default: `Inf`.
 
+- `align_time_step`: When `true` it implies that the simulation will automatically adjust the
+                     time-step to meet a constraint imposed by various schedules like `ScheduledTimes`,
+                     `TimeInterval`, `AveragedTimeInterval`, as well as a `stop_time` criterion.
+                     If `false`, i.e., no time-step alignment, then the simulation might blithely step passed
+                     the specified time. Default: `true`.
+                     By `align_time_step = false` we ensure that the time-step does _not_ change within
+                     `time_step!(simulation)`
+
 - `wall_time_limit`: Stop the simulation if it's been running for longer than this many
                      seconds of wall clock time. Default: `Inf`.
 
@@ -196,6 +205,8 @@ function reset!(sim::Simulation)
     return nothing
 end
 
+set_clock!(sim::Simulation, new_clock) = set_clock!(sim.model, new_clock)
+
 #####
 ##### Default stop criteria callback functions
 #####