Checkpointer updates

.buildkite/pipeline.yml

@@ -1,5 +1,5 @@
 env:
-  JULIA_VERSION: "1.10.8"
+  JULIA_VERSION: "1.10.9"
   JULIA_MINOR_VERSION: "1.10"
   TARTARUS_HOME: "/storage5/buildkite-agent"
   JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager
@@ -209,7 +209,7 @@ steps:
       architecture: CPU
     retry:
       automatic:
-        - exit_status: 1      
+        - exit_status: 1
           limit: 1
     depends_on: "init"
 
@@ -224,7 +224,7 @@ steps:
       architecture: CPU
     retry:
       automatic:
-        - exit_status: 1      
+        - exit_status: 1
           limit: 1
     depends_on: "init"
 
@@ -337,7 +337,7 @@ steps:
 ##### Turbulence Closures
 #####
 
-            
+
   - label: "🎣 gpu turbulence closures"
     env:
       JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER"
@@ -683,7 +683,7 @@ steps:
       architecture: CPU
     retry:
       automatic:
-        - exit_status: 1      
+        - exit_status: 1
           limit: 1
     depends_on: "init"
 
@@ -699,7 +699,7 @@ steps:
       architecture: CPU
     retry:
       automatic:
-        - exit_status: 1      
+        - exit_status: 1
           limit: 1
     depends_on: "init"
 

Project.toml

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

src/OutputWriters/checkpointer.jl

@@ -5,32 +5,60 @@ using Oceananigans: fields, prognostic_fields
 using Oceananigans.Fields: offset_data
 using Oceananigans.TimeSteppers: QuasiAdamsBashforth2TimeStepper
 
-import Oceananigans.Fields: set! 
+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
 
-function default_checkpointed_properties(model)
-    properties = [:grid, :particles, :clock, :timestepper]
-    #if has_ab2_timestepper(model)
-    #    push!(properties, :timestepper)
-    #end
-    return properties
-end
+default_checkpointed_properties(model) = [:grid, :particles, :clock, :timestepper]
 
 has_ab2_timestepper(model) = try
     model.timestepper isa QuasiAdamsBashforth2TimeStepper
 catch
     false
 end
 
+# Certain properties are required for `set!` to pickup from a checkpoint.
+function required_checkpointed_properties(model)
+    properties = [:grid, :particles, :clock]
+
+    if has_ab2_timestepper(model)
+        push!(properties, :timestepper)
+    end
+
+    return properties
+end
+
+function validate_checkpointed_properties(model, properties)
+    required_properties = required_checkpointed_properties(model)
+
+    for rp in required_properties
+        if rp ∉ properties
+            @warn "$rp is required for checkpointing. It will be added to checkpointed properties"
+            push!(properties, rp)
+        end
+    end
+
+    for p in properties
+        p isa Symbol || error("Property $p to be checkpointed must be a Symbol.")
+        p ∉ propertynames(model) && error("Cannot checkpoint $p, it is not a model property!")
+
+        if (p ∉ required_properties) && has_reference(Function, getproperty(model, p))
+            @warn "model.$p contains a function somewhere in its hierarchy and will not be checkpointed."
+            filter!(e -> e != p, properties)
+        end
+    end
+
+    return properties
+end
+
+
 """
     Checkpointer(model;
                  schedule,
@@ -79,36 +107,11 @@ function Checkpointer(model; schedule,
                       prefix = "checkpoint",
                       overwrite_existing = false,
                       verbose = false,
-                      cleanup = false,
-                      properties = default_checkpointed_properties(model))
-
-    # Certain properties are required for `set!` to pickup from a checkpoint.
-    required_properties = [:grid, :particles, :clock]
-
-    if has_ab2_timestepper(model)
-        push!(required_properties, :timestepper)
-    end
-
-    for rp in required_properties
-        if rp ∉ properties
-            @warn "$rp is required for checkpointing. It will be added to checkpointed properties"
-            push!(properties, rp)
-        end
-    end
-
-    for p in properties
-        p isa Symbol || error("Property $p to be checkpointed must be a Symbol.")
-        p ∉ propertynames(model) && error("Cannot checkpoint $p, it is not a model property!")
-
-        if (p ∉ required_properties) && has_reference(Function, getproperty(model, p))
-            @warn "model.$p contains a function somewhere in its hierarchy and will not be checkpointed."
-            filter!(e -> e != p, properties)
-        end
-    end
+                      cleanup = false)
 
     mkpath(dir)
 
-    return Checkpointer(schedule, dir, prefix, properties, overwrite_existing, verbose, cleanup)
+    return Checkpointer(schedule, dir, prefix, overwrite_existing, verbose, cleanup)
 end
 
 #####
@@ -177,29 +180,40 @@ end
 function write_output!(c::Checkpointer, 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)
@@ -212,13 +226,12 @@ end
 #####
 
 """
-    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::AbstractModel, filepath::AbstractString)
-
     addr = checkpointer_address(model)
 
     jldopen(filepath, "r") do file
@@ -280,9 +293,8 @@ function set_time_stepper_tendencies!(timestepper, file, model_fields, addr)
     return nothing
 end
 
-# For self-starting timesteppers like RK3 we do nothing 
+# For self-starting timesteppers like RK3 we do nothing
 set_time_stepper!(timestepper, args...) = nothing
 
 set_time_stepper!(timestepper::QuasiAdamsBashforth2TimeStepper, args...) =
     set_time_stepper_tendencies!(timestepper, args...)
-

src/Simulations/run.jl

@@ -35,7 +35,7 @@ end
 """
     aligned_time_step(sim, Δt)
 
-Return a time step 'aligned' with `sim.stop_time`, output writer schedules, 
+Return a time step 'aligned' with `sim.stop_time`, output writer schedules,
 and callback schedules. Alignment with `sim.stop_time` takes precedence.
 """
 function aligned_time_step(sim::Simulation, Δt)
@@ -45,7 +45,7 @@ function aligned_time_step(sim::Simulation, Δt)
 
     # Align time step with output writing and callback execution
     aligned_Δt = schedule_aligned_time_step(sim, aligned_Δt)
-    
+
     # Align time step with simulation stop time
     time_left = unit_time(sim.stop_time - clock.time)
     aligned_Δt = min(aligned_Δt, time_left)
@@ -105,7 +105,7 @@ function run!(sim; pickup=false)
         time_step!(sim)
     end
 
-    for callback in values(sim.callbacks) 
+    for callback in values(sim.callbacks)
         finalize!(callback, sim)
     end
 
@@ -137,7 +137,7 @@ function time_step!(sim::Simulation)
     initial_time_step = !(sim.initialized)
     initial_time_step && initialize!(sim)
 
-    if initial_time_step && sim.verbose 
+    if initial_time_step && sim.verbose
         @info "Executing initial time step..."
         start_time = time_ns()
     end
@@ -153,7 +153,7 @@ function time_step!(sim::Simulation)
 
     # Callbacks and callback-like things
     for diag in values(sim.diagnostics)
-        diag.schedule(sim.model) && run_diagnostic!(diag, sim.model) 
+        diag.schedule(sim.model) && run_diagnostic!(diag, sim.model)
     end
 
     for callback in values(sim.callbacks)
@@ -162,7 +162,7 @@ function time_step!(sim::Simulation)
     end
 
     for writer in values(sim.output_writers)
-        writer.schedule(sim.model) && write_output!(writer, sim.model) 
+        writer.schedule(sim.model) && write_output!(writer, sim.model)
     end
 
     if initial_time_step && sim.verbose
@@ -188,14 +188,14 @@ we_want_to_pickup(pickup::Integer) = true
 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
@@ -228,7 +228,7 @@ function initialize!(sim::Simulation)
             run_diagnostic!(diag, model)
         end
 
-        for callback in values(sim.callbacks) 
+        for callback in values(sim.callbacks)
             callback.callsite isa TimeStepCallsite && callback(sim)
         end
 
@@ -247,4 +247,3 @@ function initialize!(sim::Simulation)
 
     return nothing
 end
-

src/TimeSteppers/clock.jl

@@ -7,9 +7,9 @@ import Base: show
 import Oceananigans.Units: Time
 
 """
-    mutable struct Clock{T, FT}
+    mutable struct Clock{TT, DT, IT, S}
 
-Keeps track of the current `time`, `last_Δt`, `iteration` number, and time-stepping `stage`.
+Keep track of the current `time`, `last_Δt`, `iteration` number, and time-stepping `stage`.
 The `stage` is updated only for multi-stage time-stepping methods. The `time::T` is
 either a number or a `DateTime` object.
 """
@@ -121,5 +121,3 @@ Adapt.adapt_structure(to, clock::Clock) = (time          = clock.time,
                                            last_stage_Δt = clock.last_stage_Δt,
                                            iteration     = clock.iteration,
                                            stage         = clock.stage)
-
-

test/test_checkpointer.jl

@@ -78,12 +78,8 @@ function test_thermal_bubble_checkpointer_output(arch)
     return run_checkpointer_tests(true_model, test_model, Δt)
 end
 
-function test_hydrostatic_splash_checkpointer(arch, free_surface)
+function test_hydrostatic_splash_checkpointer(grid, free_surface)
     # Create and run "true model"
-    Nx, Ny, Nz = 16, 16, 4
-    Lx, Ly, Lz = 1, 1, 1
-
-    grid = RectilinearGrid(arch, size=(Nx, Ny, Nz), x=(-10, 10), y=(-10, 10), z=(-1, 0))
     closure = ScalarDiffusivity(ν=1e-2, κ=1e-2)
     true_model = HydrostaticFreeSurfaceModel(; grid, free_surface, closure, buoyancy=nothing, tracers=())
     test_model = deepcopy(true_model)
@@ -196,12 +192,20 @@ end
 for arch in archs
     @testset "Checkpointer [$(typeof(arch))]" begin
         @info "  Testing Checkpointer [$(typeof(arch))]..."
+
         test_thermal_bubble_checkpointer_output(arch)
-
+
+        # create a grid to test hydrostatic model
+        Nx, Ny, Nz = 16, 16, 4
+        Lx, Ly, Lz = 1, 1, 1
+        grid = RectilinearGrid(arch, size=(Nx, Ny, Nz), x=(-10, 10), y=(-10, 10), z=(-1, 0))
+
         for free_surface in [ExplicitFreeSurface(gravitational_acceleration=1),
-                             ImplicitFreeSurface(gravitational_acceleration=1)]
+                             ImplicitFreeSurface(gravitational_acceleration=1),
+                             SplitExplicitFreeSurface(gravitational_acceleration=1, substeps=5),
+                             SplitExplicitFreeSurface(grid; cfl = 0.7, gravitational_acceleration=1)]
 
-            test_hydrostatic_splash_checkpointer(arch, free_surface)
+            test_hydrostatic_splash_checkpointer(grid, free_surface)
         end
 
         run_checkpointer_cleanup_tests(arch)