Implementing checkpointer for coupled model

examples/one_degree_simulation.jl

@@ -170,6 +170,14 @@ ocean.output_writers[:surface] = JLD2OutputWriter(ocean.model, outputs;
                                                   overwrite_existing = true,
                                                   array_type = Array{Float32})
 
+# 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.
@@ -184,6 +192,16 @@ 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

@@ -27,7 +27,7 @@ export
     Metadata,
     all_dates,
     JRA55FieldTimeSeries,
-    ECCO_field, 
+    ECCO_field,
     ECCORestoring,
     LinearlyTaperedPolarMask,
     ocean_simulation,
@@ -51,7 +51,7 @@ const SKOFTS = SomeKindOfFieldTimeSeries
 @inline stateindex(a::SKOFTS, i, j, k, grid, time, args...) = @inbounds a[i, j, k, time]
 
 @inline function stateindex(a::Function, i, j, k, grid, time, loc)
-    LX, LY, LZ = loc 
+    LX, LY, LZ = loc
     λ, φ, z = node(i, j, k, grid, LX(), LY(), LZ())
     return a(λ, φ, z, time)
 end

src/OceanSeaIceModels/PrescribedAtmospheres.jl

@@ -7,8 +7,11 @@ using Oceananigans.OutputReaders: FieldTimeSeries, update_field_time_series!, ex
 using Oceananigans.TimeSteppers: Clock, tick!
 
 using Adapt
+using JLD2
 using Thermodynamics.Parameters: AbstractThermodynamicsParameters
 
+import Oceananigans.Fields: set!
+import Oceananigans.OutputWriters: checkpointer_address
 import Oceananigans.TimeSteppers: time_step!
 
 import Thermodynamics.Parameters:
@@ -22,15 +25,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 +78,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
@@ -221,7 +224,7 @@ Base.summary(::PATP{FT}) where FT = "PrescribedAtmosphereThermodynamicsParameter
 function Base.show(io::IO, p::PrescribedAtmosphereThermodynamicsParameters)
     FT = eltype(p)
 
-    cp = p.constitutive 
+    cp = p.constitutive
     hc = p.heat_capacity
     pt = p.phase_transitions
 
@@ -238,9 +241,9 @@ function Base.show(io::IO, p::PrescribedAtmosphereThermodynamicsParameters)
         "└── PhaseTransitionParameters{$FT}", '\n',
         "    ├── reference_vaporization_enthalpy (ℒᵛ⁰): ", prettysummary(pt.reference_vaporization_enthalpy), '\n',
         "    ├── reference_sublimation_enthalpy  (ℒˢ⁰): ", prettysummary(pt.reference_sublimation_enthalpy), '\n',
-        "    ├── reference_temperature (T⁰):            ", prettysummary(pt.reference_temperature), '\n',    
+        "    ├── reference_temperature (T⁰):            ", prettysummary(pt.reference_temperature), '\n',
         "    ├── triple_point_temperature (Tᵗʳ):        ", prettysummary(pt.triple_point_temperature), '\n',
-        "    ├── triple_point_pressure (pᵗʳ):           ", prettysummary(pt.triple_point_pressure), '\n',   
+        "    ├── triple_point_pressure (pᵗʳ):           ", prettysummary(pt.triple_point_pressure), '\n',
         "    ├── water_freezing_temperature (Tᶠ):       ", prettysummary(pt.water_freezing_temperature), '\n',
         "    └── total_ice_nucleation_temperature (Tⁱ): ", prettysummary(pt.total_ice_nucleation_temperature))
 end
@@ -318,6 +321,36 @@ function Base.show(io::IO, pa::PrescribedAtmosphere)
     print(io, "└── boundary_layer_height: ", prettysummary(pa.boundary_layer_height))
 end
 
+# set the clock to be the same as the ocean model
+function set!(model::PrescribedAtmosphere, checkpoint_file_path)
+    addr = checkpointer_address(model)
+
+    jldopen(checkpoint_file_path, "r") do file
+        checkpointed_clock = file["$addr/clock"]
+
+        # Update model clock
+        set_clock!(model, checkpointed_clock)
+    end
+
+    return nothing
+end
+
+checkpointer_address(::PrescribedAtmosphere) = "HydrostaticFreeSurfaceModel"
+
+"""
+    set_clock!(sim, clock)
+
+Set the clock of `sim`ulation to match the values of `clock`.
+"""
+function set_clock!(sim::PrescribedAtmosphere, clock)
+    sim.clock.time = clock.time
+    sim.clock.iteration = clock.iteration
+    sim.clock.last_Δt = clock.last_Δt
+    sim.clock.last_stage_Δt = clock.last_stage_Δt
+    sim.clock.stage = clock.stage
+    return nothing
+end
+
 function default_atmosphere_velocities(grid, times)
     ua = FieldTimeSeries{Center, Center, Nothing}(grid, times)
     va = FieldTimeSeries{Center, Center, Nothing}(grid, times)
@@ -357,15 +390,15 @@ end
 
     for fts in ftses
         update_field_time_series!(fts, time)
-    end    
-    
+    end
+
     return nothing
 end
 
 @inline thermodynamics_parameters(atmos::Nothing) = nothing
 @inline thermodynamics_parameters(atmos::PrescribedAtmosphere) = atmos.thermodynamics_parameters
 @inline surface_layer_height(atmos::PrescribedAtmosphere) = atmos.surface_layer_height
-@inline boundary_layer_height(atmos::PrescribedAtmosphere) = atmos.boundary_layer_height    
+@inline boundary_layer_height(atmos::PrescribedAtmosphere) = atmos.boundary_layer_height
 
 """
     PrescribedAtmosphere(grid, times;
@@ -420,7 +453,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 in a shortwave band
 and a longwave band) that passes through the atmosphere and arrives at the surface of ocean
@@ -434,4 +467,3 @@ Adapt.adapt_structure(to, tsdr::TwoBandDownwellingRadiation) =
                                 adapt(to, tsdr.longwave))
 
 end # module
-

src/OceanSeaIceModels/ocean_sea_ice_model.jl

@@ -3,7 +3,6 @@ using Oceananigans.TimeSteppers: Clock
 using Oceananigans: SeawaterBuoyancy
 using ClimaSeaIce.SeaIceThermodynamics: melting_temperature
 using KernelAbstractions: @kernel, @index
-
 using SeawaterPolynomials: TEOS10EquationOfState
 
 import Thermodynamics as AtmosphericThermodynamics
@@ -12,12 +11,14 @@ import Thermodynamics as AtmosphericThermodynamics
 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.Models: timestepper, NaNChecker, default_nan_checker, initialization_update_state!
+import Oceananigans.OutputWriters: default_included_properties, checkpointer_address,
+                                   write_output!, initialize_jld2_file!
+import Oceananigans.Simulations: reset!, initialize!, iteration, run!
 import Oceananigans.TimeSteppers: time_step!, update_state!, time
 import Oceananigans.Utils: prettytime
-import Oceananigans.Models: timestepper, NaNChecker, default_nan_checker, initialization_update_state!
+
+import .PrescribedAtmospheres: set_clock!
 
 mutable struct OceanSeaIceModel{I, A, O, F, C, Arch} <: AbstractModel{Nothing, Arch}
     architecture :: Arch
@@ -59,14 +60,13 @@ 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
+prognostic_fields(::OSIM)           = nothing
 fields(::OSIM)                      = NamedTuple()
 default_clock(TT)                   = Oceananigans.TimeSteppers.Clock{TT}(0, 0, 1)
+time(model::OSIM)                   = model.clock.time
+checkpointer_address(::OSIM)        = "HydrostaticFreeSurfaceModel"
 
-function reset!(model::OSIM)
-    reset!(model.ocean)
-    return nothing
-end
+reset!(model::OSIM) = reset!(model.ocean)
 
 # Make sure to initialize the exchanger here
 function initialization_update_state!(model::OSIM)
@@ -81,6 +81,32 @@ function initialize!(model::OSIM)
     return nothing
 end
 
+initialize_jld2_file!(filepath, init, jld2_kw, including, outputs, model::OSIM) =
+    initialize_jld2_file!(filepath, init, jld2_kw, including, outputs, model.ocean.model)
+
+write_output!(c::Checkpointer, model::OSIM) = write_output!(c, model.ocean.model)
+
+function set_clock!(model::OSIM, clock)
+    model.clock.time = clock.time
+    model.clock.iteration = clock.iteration
+    model.clock.last_Δt = clock.last_Δt
+    model.clock.last_stage_Δt = clock.last_stage_Δt
+    model.clock.stage = clock.stage
+    return nothing
+end
+
+function set!(model::OSIM, checkpoint_file_path::AbstractString)
+    atmosphere = model.atmosphere
+    ocean = model.ocean.model
+
+    set!(ocean, checkpoint_file_path)
+    set!(atmosphere, checkpoint_file_path)
+
+    set_clock!(model, ocean.clock)
+
+    return nothing
+end
+
 reference_density(unsupported) =
     throw(ArgumentError("Cannot extract reference density from $(typeof(unsupported))"))
 
@@ -159,10 +185,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
@@ -199,4 +223,3 @@ function above_freezing_ocean_temperature!(ocean, sea_ice::SeaIceSimulation)
 
     return nothing
 end
-

test/test_ocean_sea_ice_model.jl

@@ -43,15 +43,15 @@ using ClimaSeaIce.Rheologies
                             halo = (7, 7, 7),
                             z = (-6000, 0))
 
-        bottom_height = regrid_bathymetry(grid; 
+        bottom_height = regrid_bathymetry(grid;
                                           minimum_depth = 10,
                                           interpolation_passes = 20,
                                           major_basins = 1)
-        
+
         grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map=true)
 
         free_surface = SplitExplicitFreeSurface(grid; substeps=20)
-        ocean = ocean_simulation(grid; free_surface) 
+        ocean = ocean_simulation(grid; free_surface)
 
         backend = JRA55NetCDFBackend(4)
         atmosphere = JRA55PrescribedAtmosphere(arch; backend)
@@ -69,18 +69,18 @@ using ClimaSeaIce.Rheologies
         #####
 
         # Adding a sea ice model to the coupled model
-        τua = Field{Face, Center, Nothing}(grid) 
+        τua = Field{Face, Center, Nothing}(grid)
         τva = Field{Center, Face, Nothing}(grid)
 
-        dynamics = SeaIceMomentumEquation(grid; 
+        dynamics = SeaIceMomentumEquation(grid;
                                           coriolis = ocean.model.coriolis,
                                           top_momentum_stress = (u=τua, v=τva),
                                           rheology = ElastoViscoPlasticRheology(),
                                           solver = SplitExplicitSolver(120))
 
-        sea_ice  = sea_ice_simulation(grid; dynamics, advection=WENO(order=7)) 
+        sea_ice  = sea_ice_simulation(grid; dynamics, advection=WENO(order=7))
         liquidus = sea_ice.model.ice_thermodynamics.phase_transitions.liquidus
-        
+
         # Set the ocean temperature and salinity
         set!(ocean.model, T=temperature_metadata[1], S=salinity_metadata[1])
         above_freezing_ocean_temperature!(ocean, sea_ice)
@@ -102,3 +102,56 @@ using ClimaSeaIce.Rheologies
     end
 end
 
+"""
+    testbed_coupled_simulation(grid; stop_iteration=8)
+
+Return a test-bed coupled simulation with a Checkpointer.
+"""
+function testbed_coupled_simulation(grid; stop_iteration=8)
+    ocean = ocean_simulation(grid)
+
+    atmosphere = JRA55PrescribedAtmosphere(arch; backend=JRA55NetCDFBackend(4))
+
+    coupled_model = OceanSeaIceModel(ocean; atmosphere)
+
+    simulation = Simulation(coupled_model; Δt=10, stop_iteration)
+
+    simulation.output_writers[:checkpoint] = Checkpointer(ocean.model;
+                                                          schedule = IterationInterval(3),
+                                                          prefix = "checkpointer",
+                                                          dir = ".",
+                                                          verbose = true,
+                                                          overwrite_existing = true)
+
+    return simulation
+end
+
+@testset "Checkpointer test" begin
+    for arch in test_architectures
+
+        Nx, Ny, Nz = 40, 25, 10
+
+        grid = LatitudeLongitudeGrid(arch;
+                                     size = (Nx, Ny, Nz),
+                                     halo = (7, 7, 7),
+                                     z = (-6000, 0),
+                                     latitude  = (-75, 75),
+                                     longitude = (0, 360))
+
+        simulation = testbed_coupled_simulation(grid; stop_iteration=8)
+
+        run!(simulation)
+
+        # create a new simulation and pick up from the last checkpointer
+        new_simulation = testbed_coupled_simulation(grid; stop_iteration=13)
+
+        run!(new_simulation, pickup=true)
+
+        # ensure the ocean, atmosphere, and coupled model are all at same time and iteration
+        clock = new_simulation.model.ocean.model.clock
+        @test new_simulation.model.atmosphere.clock.iteration ≈ clock.iteration
+        @test new_simulation.model.atmosphere.clock.time ≈ clock.time
+        @test new_simulation.model.clock.iteration ≈ clock.iteration
+        @test new_simulation.model.clock.time ≈ clock.time
+    end
+end