Implementing checkpointer for coupled model

Project.toml

@@ -49,7 +49,7 @@ JLD2 = "0.4, 0.5"
 KernelAbstractions = "0.9"
 MPI = "0.20"
 NCDatasets = "0.12, 0.13, 0.14"
-Oceananigans = "0.96 - 0.99"
+Oceananigans = "0.96.1 - 0.99"
 OffsetArrays = "1.14"
 PrecompileTools = "1"
 Reactant = "0.2.45"

examples/checkpointer_mwe.jl

@@ -0,0 +1,72 @@
+using ClimaOcean
+using Oceananigans
+using CFTime
+using Dates
+using Printf
+
+arch = CPU()
+
+Nx = 120
+Ny = 50
+Nz = 20
+
+grid = LatitudeLongitudeGrid(arch;
+                             size = (Nx, Ny, Nz),
+                             halo = (7, 7, 7),
+                             z = (-6000, 0),
+                             latitude  = (-75, 75),
+                             longitude = (0, 360))
+
+
+function testbed_coupled_simulation(grid; stop_iteration=8)
+    ocean = ocean_simulation(grid)
+
+    radiation = Radiation(arch)
+
+    atmosphere = JRA55PrescribedAtmosphere(arch; backend=JRA55NetCDFBackend(4))
+
+    coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation)
+
+    return Simulation(coupled_model; Δt=10, stop_iteration)
+end
+
+simulation = testbed_coupled_simulation(grid; stop_iteration=8)
+
+wall_time = Ref(time_ns())
+
+function progress(sim)
+    ocean = sim.model.ocean
+    atmosphere = sim.model.atmosphere
+
+    step_time = 1e-9 * (time_ns() - wall_time[])
+
+    msg = @sprintf("iteration: %d, sim time: %s, atmos time: %s, ocean time: %s, wall time: %s",
+                   iteration(sim), sim.model.clock.time, atmosphere.clock.time, ocean.model.clock.time, prettytime(step_time))
+
+    @info msg
+
+    wall_time[] = time_ns()
+end
+
+simulation.callbacks[:progress] = Callback(progress, IterationInterval(1))
+
+simulation.output_writers[:checkpoint] = Checkpointer(ocean.model;
+                                                      schedule = IterationInterval(3),
+                                                      prefix = "checkpointer",
+                                                      dir = ".",
+                                                      verbose = true,
+                                                      overwrite_existing = true)
+
+run!(simulation)
+
+new_simulation = testbed_coupled_simulation(grid; stop_iteration=14)
+
+new_simulation.output_writers[:checkpoint] = Checkpointer(ocean.model;
+                                                      schedule = IterationInterval(3),
+                                                      prefix = "checkpointer",
+                                                      dir = ".",
+                                                      verbose = true,
+                                                      overwrite_existing = true)
+
+
+run!(new_simulation, pickup=true)

examples/one_degree_simulation.jl

@@ -157,6 +157,14 @@ ocean.output_writers[:surface] = JLD2Writer(ocean.model, outputs;
                                             indices = (:, :, grid.Nz),
                                             overwrite_existing = true)
 
+# We also save a series of Checkpointers that will enable us to restart the simulation
+# from a later point; ideal for longer runs.
+
+simulation.output_writers[:checkpoint] = Checkpointer(ocean.model;
+                                                      schedule = TimeInterval(5days),
+                                                      prefix = "one_degree_checkpointer",
+                                                      overwrite_existing = true)
+
 # ### Ready to run
 
 # We are ready to press the big red button and run the simulation.
@@ -170,6 +178,16 @@ simulation.Δt = 20minutes
 simulation.stop_time = 360days
 run!(simulation)
 
+# After we had run long-enough to produce at least one checkpointer file, if we want to restart the
+# simulation we call
+#
+# ```julia
+# run!(simulation, pickup=true)
+# ```
+#
+# Doing so, simulation will be picked up from the latest checkpointer.
+
+
 # ### A pretty movie
 #
 # We load the saved output and make a pretty movie of the simulation. First we plot a snapshot:

src/ClimaOcean.jl

@@ -108,4 +108,3 @@ using PrecompileTools: @setup_workload, @compile_workload
 end
 
 end # module
-

src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl

@@ -418,4 +418,3 @@ function sea_ice_similarity_theory(sea_ice::SeaIceSimulation)
     interface_temperature_type = SkinTemperature(internal_flux)
     return SimilarityTheoryFluxes(; interface_temperature_type)
 end
-

src/OceanSeaIceModels/OceanSeaIceModels.jl

@@ -26,6 +26,11 @@ using ClimaOcean: stateindex
 using KernelAbstractions: @kernel, @index
 using KernelAbstractions.Extras.LoopInfo: @unroll
 
+import Oceananigans.OutputWriters: checkpointer_address,
+                                   required_checkpointed_properties,
+                                   default_checkpointed_properties,
+                                   prognostic_fields
+
 function downwelling_radiation end
 function freshwater_flux end
 function reference_density end
@@ -69,6 +74,7 @@ import .InterfaceComputations:
     compute_sea_ice_ocean_fluxes!
 
 include("ocean_sea_ice_model.jl")
+include("output_writers.jl")
 include("freezing_limited_ocean_temperature.jl")
 include("time_step_ocean_sea_ice_model.jl")
 

src/OceanSeaIceModels/PrescribedAtmospheres.jl

@@ -7,9 +7,15 @@ using Oceananigans.OutputReaders: FieldTimeSeries, update_field_time_series!, ex
 using Oceananigans.TimeSteppers: Clock, tick!
 
 using Adapt
+using JLD2
 using Thermodynamics.Parameters: AbstractThermodynamicsParameters
 
-import Oceananigans.TimeSteppers: time_step!
+import Oceananigans: initialize!, prognostic_fields
+import Oceananigans.Fields: set!
+import Oceananigans.TimeSteppers: time_step!, set_clock!
+import Oceananigans.OutputWriters: checkpointer_address,
+                                   required_checkpointed_properties,
+                                   default_checkpointed_properties
 
 import Thermodynamics.Parameters:
     gas_constant,   #
@@ -22,15 +28,15 @@ import Thermodynamics.Parameters:
     T_0,            # Enthalpy reference temperature
     LH_v0,          # Vaporization enthalpy at the reference temperature
     LH_s0,          # Sublimation enthalpy at the reference temperature
-    LH_f0,          # Fusionn enthalpy at the reference temperature
+    LH_f0,          # Fusion enthalpy at the reference temperature
     cp_d,           # Heat capacity of dry air at constant pressure
     cp_v,           # Isobaric specific heat capacity of gaseous water vapor
     cp_l,           # Isobaric specific heat capacity of liquid water
     cp_i,           # Isobaric specific heat capacity of water ice
     cv_v,           # Heat capacity of dry air at constant volume
     cv_l,           # Isobaric specific heat capacity of liquid water
     cv_i,           # Isobaric specific heat capacity of liquid water
-    e_int_v0,       # what? someting about reference internal energy of water vapor
+    e_int_v0,       # what? something about reference internal energy of water vapor
     T_freeze,       # Freezing temperature of _pure_ water
     T_triple,       # Triple point temperature of _pure_ water
     press_triple,   # Triple point pressure of pure water
@@ -75,7 +81,7 @@ Construct a set of parameters that define the density of moist air,
 ```
 
 where ``p`` is pressure, ``T`` is temperature, ``q`` defines the partition
-of total mass into vapor, liqiud, and ice mass fractions, and
+of total mass into vapor, liquid, and ice mass fractions, and
 ``Rᵐ`` is the effective specific gas constant for the mixture,
 
 ```math
@@ -318,6 +324,29 @@ function Base.show(io::IO, pa::PrescribedAtmosphere)
     print(io, "└── boundary_layer_height: ", prettysummary(pa.boundary_layer_height))
 end
 
+checkpointer_address(::PrescribedAtmosphere)             = "PrescribedAtmosphere"
+default_checkpointed_properties(::PrescribedAtmosphere)  = [:clock]
+required_checkpointed_properties(::PrescribedAtmosphere) = [:clock]
+prognostic_fields(::PrescribedAtmosphere)                = ()
+
+function set!(model::PrescribedAtmosphere, checkpoint_file_path::AbstractString)
+    addr = checkpointer_address(model)
+
+    checkpointed_clock = nothing
+    jldopen(checkpoint_file_path, "r") do file
+        checkpointed_clock = file["$addr/clock"]
+    end
+
+    # Update model clock
+    set_clock!(model, checkpointed_clock)
+
+    @info "hi, I set the clock of the atmosphere"
+    return nothing
+end
+
+set_clock!(atmos::PrescribedAtmosphere, new_clock) =
+    set_clock!(atmos.clock, new_clock)
+
 function default_atmosphere_velocities(grid, times)
     ua = FieldTimeSeries{Center, Center, Nothing}(grid, times)
     va = FieldTimeSeries{Center, Center, Nothing}(grid, times)
@@ -327,7 +356,7 @@ end
 function default_atmosphere_tracers(grid, times)
     Ta = FieldTimeSeries{Center, Center, Nothing}(grid, times)
     qa = FieldTimeSeries{Center, Center, Nothing}(grid, times)
-    parent(Ta) .= 273.15 + 20
+    parent(Ta) .= 273.15 + 20 # ᵒK
     return (T=Ta, q=qa)
 end
 
@@ -347,13 +376,20 @@ end
 This is approximately equal to the mean sea-level atmospheric pressure on Earth. """
 function default_atmosphere_pressure(grid, times)
     pa = FieldTimeSeries{Center, Center, Nothing}(grid, times)
-    parent(pa) .= 101325
+    parent(pa) .= 101325 # Pa
     return pa
 end
 
 @inline function time_step!(atmos::PrescribedAtmosphere, Δt)
     tick!(atmos.clock, Δt)
+    update_state!(atmos)
+
+    return nothing
+end
+
+initialize!(atmos::PrescribedAtmosphere) = update_state!(atmos)
 
+function update_state!(atmos::PrescribedAtmosphere)
     time = Time(atmos.clock.time)
     ftses = extract_field_time_series(atmos)
 
@@ -422,7 +458,7 @@ struct TwoBandDownwellingRadiation{SW, LW}
 end
 
 """
-    TwoBandDownwellingRadiation(shortwave=nothing, longwave=nothing)
+    TwoBandDownwellingRadiation(; shortwave=nothing, longwave=nothing)
 
 Return a two-band model for downwelling radiation (split into a shortwave band
 and a longwave band) that passes through the atmosphere and arrives at the surface of ocean

src/OceanSeaIceModels/ocean_sea_ice_model.jl

@@ -3,21 +3,25 @@ using Oceananigans.TimeSteppers: Clock
 using Oceananigans: SeawaterBuoyancy
 using ClimaSeaIce.SeaIceThermodynamics: melting_temperature
 using KernelAbstractions: @kernel, @index
-
 using SeawaterPolynomials: TEOS10EquationOfState
 
 import Thermodynamics as AtmosphericThermodynamics
 
 # Simulations interface
-import Oceananigans: fields, prognostic_fields
+import Oceananigans: fields,
+                     prognostic_fields
 import Oceananigans.Architectures: architecture
 import Oceananigans.Fields: set!
-import Oceananigans.Models: timestepper, NaNChecker, default_nan_checker
-import Oceananigans.OutputWriters: default_included_properties
-import Oceananigans.Simulations: reset!, initialize!, iteration
-import Oceananigans.TimeSteppers: time_step!, update_state!, time
-import Oceananigans.Utils: prettytime
 import Oceananigans.Models: timestepper, NaNChecker, default_nan_checker, initialization_update_state!
+import Oceananigans.OutputWriters: checkpointer_address,
+                                   required_checkpointed_properties,
+                                   default_checkpointed_properties
+
+import Oceananigans.Simulations: reset!, initialize!, iteration, run!
+import Oceananigans.TimeSteppers: time_step!, update_state!, time, set_clock!
+import Oceananigans.Utils: prettytime
+
+import .PrescribedAtmospheres: set_clock!
 
 mutable struct OceanSeaIceModel{I, A, O, F, C, Arch} <: AbstractModel{Nothing, Arch}
     architecture :: Arch
@@ -29,6 +33,7 @@ mutable struct OceanSeaIceModel{I, A, O, F, C, Arch} <: AbstractModel{Nothing, A
 end
 
 const OSIM = OceanSeaIceModel
+const OSIMPA = OceanSeaIceModel{<:Any, <:PrescribedAtmosphere}
 
 function Base.summary(model::OSIM)
     A = nameof(typeof(architecture(model)))
@@ -53,20 +58,23 @@ function Base.show(io::IO, cm::OSIM)
 end
 
 # Assumption: We have an ocean!
-architecture(model::OSIM)           = architecture(model.ocean.model)
-Base.eltype(model::OSIM)            = Base.eltype(model.ocean.model)
-prettytime(model::OSIM)             = prettytime(model.clock.time)
-iteration(model::OSIM)              = model.clock.iteration
-timestepper(::OSIM)                 = nothing
-default_included_properties(::OSIM) = tuple()
-prognostic_fields(cm::OSIM)         = nothing
-fields(::OSIM)                      = NamedTuple()
-default_clock(TT)                   = Oceananigans.TimeSteppers.Clock{TT}(0, 0, 1)
-
-function reset!(model::OSIM)
-    reset!(model.ocean)
-    return nothing
-end
+architecture(model::OSIM)                = architecture(model.ocean.model)
+Base.eltype(model::OSIM)                 = Base.eltype(model.ocean.model)
+prettytime(model::OSIM)                  = prettytime(model.clock.time)
+iteration(model::OSIM)                   = model.clock.iteration
+timestepper(::OSIM)                      = nothing
+default_included_properties(::OSIM)      = tuple()
+prognostic_fields(::OSIM)                = tuple()
+fields(::OSIM)                           = NamedTuple()
+default_clock(TT)                        = Oceananigans.TimeSteppers.Clock{TT}(0, 0, 1)
+time(model::OSIM)                        = model.clock.time
+required_checkpointed_properties(::OSIM) = [:clock]
+default_checkpointed_properties(::OSIM)  = [:clock]
+checkpointer_address(::OSIM)             = "OceanSeaIceModel"
+
+reset!(model::OSIM) = reset!(model.ocean)
+
+components(model::OSIM) = (model.atmosphere, model.ocean.model)
 
 # Make sure to initialize the exchanger here
 function initialization_update_state!(model::OSIM)
@@ -77,10 +85,25 @@ end
 
 function initialize!(model::OSIM)
     initialize!(model.ocean)
+    initialize!(model.atmosphere)
     initialize!(model.interfaces.exchanger, model.atmosphere)
     return nothing
 end
 
+function set_clock!(model::OSIM, clock)
+
+    set_clock!(model.clock, clock)
+
+    # set the component clocks
+    atmos, ocean = model.atmosphere, model.ocean
+    set_clock!(atmos, clock)
+    set_clock!(ocean, clock)
+
+    return nothing
+end
+
+set_clock!(sim::Simulation, new_clock) = set_clock!(sim.model, new_clock)
+
 reference_density(unsupported) =
     throw(ArgumentError("Cannot extract reference density from $(typeof(unsupported))"))
 
@@ -159,10 +182,8 @@ function OceanSeaIceModel(ocean, sea_ice=FreezingLimitedOceanTemperature(eltype(
     return ocean_sea_ice_model
 end
 
-time(coupled_model::OceanSeaIceModel) = coupled_model.clock.time
-
 # Check for NaNs in the first prognostic field (generalizes to prescribed velocities).
-function default_nan_checker(model::OceanSeaIceModel)
+function default_nan_checker(model::OSIM)
     u_ocean = model.ocean.model.velocities.u
     nan_checker = NaNChecker((; u_ocean))
     return nan_checker