Adapt the checkpointer to a SeaIceModel (#90)

* try it out

* add a testset

* add the code credit

* Update sea_ice_momentum_equations.jl

* Fix auxiliary fields reference in prognostic_fields

* Update Oceananigans compat

* Update test_checkpointing.jl

* pass

* add the set to clock

---------

Co-authored-by: Navid C. Constantinou <[email protected]>

Author: simone-silvestri

Project.toml

@@ -16,7 +16,7 @@ SeawaterPolynomials = "d496a93d-167e-4197-9f49-d3af4ff8fe40"
 Adapt = "3, 4"
 KernelAbstractions = "0.9"
 JLD2 = "0.6.2"
-Oceananigans = "0.99, 0.100, 0.101, 0.102, 0.103, 0.104"
+Oceananigans = "0.100, 0.101, 0.102, 0.103, 0.104"
 RootSolvers = "0.3, 0.4"
 Roots = "2"
 SeawaterPolynomials = "0.3.4"

docs/Project.toml

@@ -9,4 +9,4 @@ Oceananigans = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 [compat]
 CairoMakie = "0.11.0, 0.12.0"
 Documenter = "1"
-Oceananigans = "0.95.3 - 0.99"
+Oceananigans = "0.100, 0.101, 0.102"

src/Rheologies/Rheologies.jl

@@ -27,9 +27,13 @@ sea ice stresses.
 """
 Auxiliaries(rheology, grid::AbstractGrid) = Auxiliaries(NamedTuple(), nothing)
 
+import Oceananigans: prognostic_fields
+
 # Nothing rheology
 initialize_rheology!(model, rheology) = nothing
+
 compute_stresses!(dynamics, fields, grid, rheology, Δt) = nothing
+prognostic_fields(mom, rheology) = NamedTuple()
 
 # Nothing rheology or viscous rheology
 @inline compute_substep_Δtᶠᶜᶜ(i, j, grid, Δt, rheology, substeps, fields) = Δt / substeps

src/Rheologies/elasto_visco_plastic_rheology.jl

@@ -141,6 +141,11 @@ function Auxiliaries(r::ElastoViscoPlasticRheology, grid::AbstractGrid)
     return Auxiliaries(fields, kernels)
 end
 
+prognostic_fields(mom, ::ElastoViscoPlasticRheology) =
+    (σ₁₁ = mom.auxiliaries.fields.σ₁₁,
+     σ₂₂ = mom.auxiliaries.fields.σ₂₂,
+     σ₁₂ = mom.auxiliaries.fields.σ₁₂)
+
 # Extend the `adapt_structure` function for the ElastoViscoPlasticRheology
 Adapt.adapt_structure(to, r::ElastoViscoPlasticRheology) = 
     ElastoViscoPlasticRheology(Adapt.adapt(to, r.ice_compressive_strength),

src/SeaIceDynamics/SeaIceDynamics.jl

@@ -25,8 +25,8 @@ using ClimaSeaIce.Rheologies: ∂ⱼ_σ₁ⱼ,
                               compute_substep_Δtᶜᶠᶜ,
                               sum_of_forcing_u,
                               sum_of_forcing_v
-                              
-import Oceananigans: fields
+
+import Oceananigans: fields, prognostic_fields
 
 ## A Framework to solve for the ice momentum equation, in the form:
 ## 

src/SeaIceDynamics/sea_ice_momentum_equations.jl

@@ -80,3 +80,4 @@ function SeaIceMomentumEquation(grid;
 end
 
 fields(mom::SeaIceMomentumEquation) = mom.auxiliaries.fields
+prognostic_fields(mom::SeaIceMomentumEquation) = prognostic_fields(mom, mom.rheology)

src/SeaIceThermodynamics/slab_sea_ice_thermodynamics.jl

@@ -1,4 +1,4 @@
-import Oceananigans: fields
+import Oceananigans: fields, prognostic_fields
 
 struct ProportionalEvolution end
 
@@ -32,6 +32,7 @@ function Base.show(io::IO, therm::SSIT)
 end
 
 fields(therm::SSIT) = (; Tu = therm.top_surface_temperature)
+prognostic_fields(therm::SSIT) = (; Gʰ = therm.thermodynamic_tendency)
 
 """
     SlabSeaIceThermodynamics(grid; kw...)

src/sea_ice_model.jl

@@ -152,6 +152,8 @@ function set!(model::SIM; h=nothing, ℵ=nothing)
     return nothing
 end
 
+set!(model::SIM, new_clock::Clock) = set!(model.clock, new_clock)
+
 Base.summary(model::SIM) = "SeaIceModel"
 prettytime(model::SIM) = prettytime(model.clock.time)
 iteration(model::SIM) = model.clock.iteration
@@ -186,8 +188,11 @@ fields(model::SIM) = merge((; h  = model.ice_thickness,
                            fields(model.ice_thermodynamics),
                            fields(model.dynamics))
 
+prognostic_fields(::Nothing) = NamedTuple()
+                           
 # TODO: make this correct
 prognostic_fields(model::SIM) = merge((; h  = model.ice_thickness,
                                          ℵ  = model.ice_concentration),
-                                      model.tracers,
-                                      model.velocities)
+                                      model.velocities,
+                                      prognostic_fields(model.dynamics),
+                                      prognostic_fields(model.ice_thermodynamics))

test/runtests.jl

@@ -25,4 +25,5 @@ end
 
 include("test_sea_ice_advection.jl")
 include("test_time_stepping.jl")
-include("test_distributed_sea_ice.jl")
+include("test_distributed_sea_ice.jl")
+include("test_checkpointing.jl")

test/test_checkpointing.jl

@@ -0,0 +1,135 @@
+using ClimaSeaIce
+using ClimaSeaIce.SeaIceDynamics
+using ClimaSeaIce.SeaIceThermodynamics
+using Test
+
+using Oceananigans.Fields: @allowscalar
+using Oceananigans: prognostic_fields
+
+# # Same test as in Oceananigans
+function test_model_equality(test_model, true_model)
+    @allowscalar begin
+        test_model_fields = prognostic_fields(test_model)
+        true_model_fields = prognostic_fields(true_model)
+        field_names = keys(test_model_fields)
+
+        for name in field_names
+            @test all(test_model_fields[name].data .≈ true_model_fields[name].data)
+        end
+    end
+
+    return nothing
+end
+
+# Test copied from Oceananigans.jl, code credit: 
+# https://github.com/CliMA/Oceananigans.jl/blob/main/test/test_checkpointer.jl#L94-L174
+function run_checkpointer_tests(true_model, test_model, Δt)
+    true_simulation = Simulation(true_model, Δt=Δt, stop_iteration=5)
+
+    checkpointer = Checkpointer(true_model, schedule=IterationInterval(5), overwrite_existing=true)
+    push!(true_simulation.output_writers, checkpointer)
+
+    run!(true_simulation) # for 5 iterations
+
+    checkpointed_model = deepcopy(true_simulation.model)
+
+    true_simulation.stop_iteration = 9
+    run!(true_simulation) # for 4 more iterations
+
+    #####
+    ##### Test `set!(model, checkpoint_file)`
+    #####
+
+    set!(test_model, "checkpoint_iteration5.jld2")
+
+    @test test_model.clock.iteration == checkpointed_model.clock.iteration
+    @test test_model.clock.time == checkpointed_model.clock.time
+    test_model_equality(test_model, checkpointed_model)
+
+    # This only applies to QuasiAdamsBashforthTimeStepper:
+    @test test_model.clock.last_Δt == checkpointed_model.clock.last_Δt
+
+    #####
+    ##### Test pickup from explicit checkpoint path
+    #####
+
+    test_simulation = Simulation(test_model, Δt=Δt, stop_iteration=9)
+
+    # Pickup from explicit checkpoint path
+    run!(test_simulation, pickup="checkpoint_iteration0.jld2")
+
+    @info "Testing model equality when running with pickup=checkpoint_iteration0.jld2."
+    @test test_simulation.model.clock.iteration == true_simulation.model.clock.iteration
+    @test test_simulation.model.clock.time == true_simulation.model.clock.time
+    test_model_equality(test_model, true_model)
+
+    run!(test_simulation, pickup="checkpoint_iteration5.jld2")
+    @info "Testing model equality when running with pickup=checkpoint_iteration5.jld2."
+
+    @test test_simulation.model.clock.iteration == true_simulation.model.clock.iteration
+    @test test_simulation.model.clock.time == true_simulation.model.clock.time
+    test_model_equality(test_model, true_model)
+
+    #####
+    ##### Test `run!(sim, pickup=true)
+    #####
+
+    # Pickup using existing checkpointer
+    test_simulation.output_writers[:checkpointer] =
+        Checkpointer(test_model, schedule=IterationInterval(5), overwrite_existing=true)
+
+    run!(test_simulation, pickup=true)
+    @info "    Testing model equality when running with pickup=true."
+
+    @test test_simulation.model.clock.iteration == true_simulation.model.clock.iteration
+    @test test_simulation.model.clock.time == true_simulation.model.clock.time
+    test_model_equality(test_model, true_model)
+
+    run!(test_simulation, pickup=0)
+    @info "    Testing model equality when running with pickup=0."
+
+    @test test_simulation.model.clock.iteration == true_simulation.model.clock.iteration
+    @test test_simulation.model.clock.time == true_simulation.model.clock.time
+    test_model_equality(test_model, true_model)
+
+    run!(test_simulation, pickup=5)
+    @info "    Testing model equality when running with pickup=5."
+
+    @test test_simulation.model.clock.iteration == true_simulation.model.clock.iteration
+    @test test_simulation.model.clock.time == true_simulation.model.clock.time
+    test_model_equality(test_model, true_model)
+
+    rm("checkpoint_iteration0.jld2", force=true)
+    rm("checkpoint_iteration5.jld2", force=true)
+
+    return nothing
+end
+
+function test_sea_ice_checkpointer_output(arch)
+    # Create and run "true model"
+    Nx, Ny = 16, 16
+    Lx, Ly = 100, 100
+    Δt = 1
+
+    grid = RectilinearGrid(arch, size=(Nx, Ny), x=(0, Lx), y=(0, Ly), topology=(Bounded, Bounded, Flat))
+    for ice_thermodynamics in (nothing, SlabSeaIceThermodynamics(grid))
+        for dynamics in (nothing, SeaIceMomentumEquation(grid))
+            
+            true_model = SeaIceModel(grid; dynamics, ice_thermodynamics)
+            test_model = deepcopy(true_model)
+
+            for field in merge(true_model.velocities, 
+                            (; h = true_model.ice_concentration, 
+                               ℵ = true_model.ice_thickness))
+
+                set!(field, (x, y) -> rand() * 1e-5)
+            end
+
+            run_checkpointer_tests(true_model, test_model, Δt)
+        end
+    end
+end
+
+@testset "Checkpointing Tests" begin
+    test_sea_ice_checkpointer_output(CPU())
+end