Implementing checkpointer for coupled model

examples/checkpointer_mwe.jl

@@ -0,0 +1,88 @@
+using ClimaOcean
+using Oceananigans
+using Oceananigans.Units
+using CFTime
+using Dates
+using Printf
+
+arch = CPU()
+
+Nx = 144
+Ny = 60
+Nz = 40
+
+grid = LatitudeLongitudeGrid(arch;
+                             size = (Nx, Ny, Nz),
+                             halo = (7, 7, 7),
+                             z = (-6000, 0),
+                             latitude  = (-75, 75),
+                             longitude = (0, 360))
+
+ocean = ocean_simulation(grid)
+
+# date = DateTimeProlepticGregorian(1993, 1, 1)
+# set!(ocean.model, T=ECCOMetadata(:temperature; dates=date),
+#                   S=ECCOMetadata(:salinity; dates=date))
+
+radiation = Radiation(arch)
+
+atmosphere = JRA55PrescribedAtmosphere(arch; backend=JRA55NetCDFBackend(41))
+
+coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation)
+
+simulation = Simulation(coupled_model; Δt=10, stop_iteration=8)
+
+wall_time = Ref(time_ns())
+
+function progress(sim)
+    ocean = sim.model.ocean
+    atmosphere = sim.model.atmosphere
+
+    u, v, w = ocean.model.velocities
+    T = ocean.model.tracers.T
+
+    Tmax = maximum(interior(T))
+    Tmin = minimum(interior(T))
+
+    umax = (maximum(abs, interior(u)),
+            maximum(abs, interior(v)),
+            maximum(abs, interior(w)))
+
+    step_time = 1e-9 * (time_ns() - wall_time[])
+
+    msg = @sprintf("Iter: %d, sim time: %s, atmos time: %s, ocean time: %s", iteration(sim), sim.model.clock.time, atmosphere.clock.time, ocean.model.clock.time)
+    msg *= @sprintf(", max|u|: (%.2e, %.2e, %.2e) m s⁻¹, extrema(T): (%.2f, %.2f) ᵒC, wall time: %s",
+                    umax..., Tmax, Tmin, prettytime(step_time))
+
+    @info msg
+
+    wall_time[] = time_ns()
+end
+
+simulation.callbacks[:progress] = Callback(progress, IterationInterval(1))
+
+outputs = merge(ocean.model.tracers, ocean.model.velocities)
+
+simulation.output_writers[:surface] = JLD2OutputWriter(ocean.model, outputs;
+                                                       schedule = IterationInterval(2),
+                                                       filename = "surface",
+                                                       indices = (:, :, grid.Nz),
+                                                       with_halos = true,
+                                                       overwrite_existing = true,
+                                                       array_type = Array{Float32})
+
+output_dir = "."
+prefix = "checkpointer_mwe"
+
+simulation.output_writers[:checkpoint] = Checkpointer(ocean.model;
+                                                      schedule = IterationInterval(3),
+                                                      prefix = prefix,
+                                                      dir = output_dir,
+                                                      verbose = true,
+                                                      overwrite_existing = true)
+
+run!(simulation)
+
+simulation.stop_iteration += 5
+
+run!(simulation, pickup=true)

examples/near_global_ocean_simulation.jl

@@ -30,10 +30,10 @@ using Printf
 # The total depth of the domain is set to 6000 meters.
 # Finally, we specify the architecture for the simulation, which in this case is a GPU.
 
-arch = GPU()
+arch = CPU()
 
-Nx = 1440
-Ny = 600
+Nx = 144
+Ny = 60
 Nz = 40
 
 depth = 6000meters
@@ -53,21 +53,21 @@ grid = LatitudeLongitudeGrid(arch;
 # (i) all the minor enclosed basins except the 3 largest `major_basins`, as well as
 # (ii) regions that are shallower than `minimum_depth`.
 
-bottom_height = regrid_bathymetry(grid;
-                                  minimum_depth = 10meters,
-                                  interpolation_passes = 5,
-                                  major_basins = 3)
+# bottom_height = regrid_bathymetry(grid;
+#                                   minimum_depth = 10meters,
+#                                   interpolation_passes = 5,
+#                                   major_basins = 3)
 
-grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map=true)
+# grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map=true)
 
 # Let's see what the bathymetry looks like:
 
-h = grid.immersed_boundary.bottom_height
+# h = grid.immersed_boundary.bottom_height
 
-fig, ax, hm = heatmap(h, colormap=:deep, colorrange=(-depth, 0))
-Colorbar(fig[0, 1], hm, label="Bottom height (m)", vertical=false)
-save("bathymetry.png", fig)
-nothing #hide
+# fig, ax, hm = heatmap(h, colormap=:deep, colorrange=(-depth, 0))
+# Colorbar(fig[0, 1], hm, label="Bottom height (m)", vertical=false)
+# save("bathymetry.png", fig)
+# nothing #hide
 
 # ![](bathymetry.png)
 
@@ -83,9 +83,9 @@ ocean.model
 
 # We initialize the ocean model with ECCO2 temperature and salinity for January 1, 1993.
 
-date = DateTimeProlepticGregorian(1993, 1, 1)
-set!(ocean.model, T=Metadata(:temperature; dates=date, dataset=ECCO4Monthly()),
-                  S=Metadata(:salinity; dates=date, dataset=ECCO4Monthly()))
+# date = DateTimeProlepticGregorian(1993, 1, 1)
+# set!(ocean.model, T=ECCOMetadata(:temperature; dates=date),
+#                   S=ECCOMetadata(:salinity; dates=date))
 
 # ### Prescribed atmosphere and radiation
 #
@@ -117,7 +117,7 @@ coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation)
 # We then create a coupled simulation. We start with a small-ish time step of 90 seconds.
 # We run the simulation for 10 days with this small-ish time step.
 
-simulation = Simulation(coupled_model; Δt=90, stop_time=10days)
+simulation = Simulation(coupled_model; Δt=90, stop_iteration=10)
 
 # We define a callback function to monitor the simulation's progress,
 
@@ -146,7 +146,7 @@ function progress(sim)
     wall_time[] = time_ns()
 end
 
-simulation.callbacks[:progress] = Callback(progress, TimeInterval(5days))
+simulation.callbacks[:progress] = Callback(progress, IterationInterval(1))
 
 # ### Set up output writers
 #
@@ -156,13 +156,23 @@ simulation.callbacks[:progress] = Callback(progress, TimeInterval(5days))
 # Below, we use `indices` to save only the values of the variables at the surface, which corresponds to `k = grid.Nz`
 
 outputs = merge(ocean.model.tracers, ocean.model.velocities)
-ocean.output_writers[:surface] = JLD2OutputWriter(ocean.model, outputs;
-                                                  schedule = TimeInterval(1days),
-                                                  filename = "near_global_surface_fields",
-                                                  indices = (:, :, grid.Nz),
-                                                  with_halos = true,
-                                                  overwrite_existing = true,
-                                                  array_type = Array{Float32})
+# ocean.output_writers[:surface] = JLD2OutputWriter(ocean.model, outputs;
+#                                                   schedule = IterationInterval(1),
+#                                                   filename = "near_global_surface_fields",
+#                                                   indices = (:, :, grid.Nz),
+#                                                   with_halos = true,
+#                                                   overwrite_existing = true,
+#                                                   array_type = Array{Float32})
+
+output_dir = "."
+prefix = "near_global"
+
+ocean.output_writers[:checkpoint] = Checkpointer(ocean.model;
+                                                schedule = TimeInterval(3minutes),
+                                                prefix = prefix,
+                                                dir = output_dir)#,
+                                                # verbose = true,
+                                                # overwrite_existing = true)
 
 # ### Spinning up the simulation
 #
@@ -175,11 +185,12 @@ run!(simulation)
 # ### Running the simulation for real
 
 # After the initial spin up of 10 days, we can increase the time-step and run for longer.
-
+#=
 simulation.stop_time = 60days
 simulation.Δt = 10minutes
 run!(simulation)
 
+
 # ## A pretty movie
 #
 # It's time to make a pretty movie of the simulation. First we load the output we've been saving on
@@ -259,3 +270,4 @@ end
 nothing #hide
 
 # ![](near_global_ocean_surface.mp4)
+=#

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/ocean_sea_ice_model.jl

@@ -1,9 +1,9 @@
 using Oceananigans
+using Oceananigans.OutputWriters: checkpoint_path
 using Oceananigans.TimeSteppers: Clock
 using Oceananigans: SeawaterBuoyancy
 using ClimaSeaIce.SeaIceThermodynamics: melting_temperature
 using KernelAbstractions: @kernel, @index
-
 using SeawaterPolynomials: TEOS10EquationOfState
 
 import Thermodynamics as AtmosphericThermodynamics
@@ -13,11 +13,11 @@ 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.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
 
 struct OceanSeaIceModel{I, A, O, F, C, Arch} <: AbstractModel{Nothing, Arch}
     architecture :: Arch
@@ -29,6 +29,7 @@ struct OceanSeaIceModel{I, A, O, F, C, Arch} <: AbstractModel{Nothing, Arch}
 end
 
 const OSIM = OceanSeaIceModel
+const OSIMSIM = Simulation{<:OceanSeaIceModel}
 
 function Base.summary(model::OSIM)
     A = nameof(typeof(architecture(model)))
@@ -59,9 +60,11 @@ 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)
@@ -73,6 +76,26 @@ 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!(sim::OSIMSIM, pickup::Union{Bool, Integer, String})
+    checkpoint_file_path = checkpoint_path(pickup, sim.output_writers)
+
+    set!(sim.model.ocean.model, checkpoint_file_path)
+
+    sim.model.clock.iteration = sim.model.ocean.model.clock.iteration
+    sim.model.clock.time = sim.model.ocean.model.clock.time
+
+    # Setting the atmosphere time to the ocean time
+    sim.model.atmosphere.clock.iteration = sim.model.ocean.model.clock.iteration
+    sim.model.atmosphere.clock.time = sim.model.ocean.model.clock.time
+
+    return nothing
+end
+
 reference_density(unsupported) =
     throw(ArgumentError("Cannot extract reference density from $(typeof(unsupported))"))
 
@@ -114,7 +137,7 @@ function OceanSeaIceModel(ocean, sea_ice=FreezingLimitedOceanTemperature(eltype(
         pop!(ocean.callbacks, :wall_time_limit_exceeded, nothing)
         pop!(ocean.callbacks, :nan_checker, nothing)
     end
-    
+
     if sea_ice isa SeaIceSimulation
         if !isnothing(sea_ice.callbacks)
             pop!(sea_ice.callbacks, :stop_time_exceeded, nothing)
@@ -151,10 +174,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