Checkpointer revival

src/Oceananigans.jl

@@ -109,8 +109,10 @@ export
  BinaryOutputWriter,
  NetCDFOutputWriter,
  JLD2OutputWriter,
+ Checkpointer,
  write_output,
  read_output,
+ restore_from_checkpoint,
 
  # Model diagnostics
  Diagnostic,

src/output_writers.jl

@@ -324,3 +324,72 @@ function time_to_write(clock, out::JLD2OutputWriter)
  return false
  end
 end
+
+mutable struct Checkpointer <: OutputWriter
+ dir :: String
+ prefix :: String
+ output_frequency :: Int
+end
+
+function Checkpointer(; output_frequency, dir=".", prefix="checkpoint", force=false)
+ mkpath(dir)
+ return Checkpointer(dir, prefix, output_frequency)
+end
+
+function savesubfields!(file, model, name, flds=propertynames(getproperty(model, name)))
+ for f in flds
+ file["$name/$f"] = Array(getproperty(getproperty(model, name), f).data.parent)
+ end
+ return nothing
+end
+
+checkpointed_structs = [:arch, :boundary_conditions, :grid, :clock, :eos, :constants, :closure]
+checkpointed_fieldsets = [:velocities, :tracers, :G, :Gp]
+
+function write_output(model, c::Checkpointer)
+ @warn "Checkpointer will not save forcing functions, output writers, or diagnostics. They will need to be " *
+ "restored manually."
+
+ filepath = joinpath(c.dir, c.prefix * string(model.clock.iteration) * ".jld2")
+
+ jldopen(filepath, "w") do file
+ # Checkpointing model properties that we can just serialize.
+ [file["$p"] = getproperty(model, p) for p in checkpointed_structs]
+
+ # Checkpointing structs containing fields.
+ [savesubfields!(file, model, p) for p in checkpointed_fieldsets]
+ end
+end
+
+_arr(::CPU, a) = a
+_arr(::GPU, a) = CuArray(a)
+
+function restore_from_checkpoint(filepath)
+ @warn "Checkpointer cannot restore forcing functions, output writers, or diagnostics. They will need to be " *
+ "restored manually."
+
+ kwargs = Dict{Symbol, Any}() # We'll store all the kwargs we need to initialize a Model.
+
+ file = jldopen(filepath, "r")
+
+ # Restore model properties that were just serialized.
+ for p in checkpointed_structs
+ kwargs[Symbol(p)] = file["$p"]
+ end
+
+ # The Model constructor needs N and L.
+ kwargs[:N] = (kwargs[:grid].Nx, kwargs[:grid].Ny, kwargs[:grid].Nz)
+ kwargs[:L] = (kwargs[:grid].Lx, kwargs[:grid].Ly, kwargs[:grid].Lz)
+
+ model = Model(; kwargs...)
+
+ for p in checkpointed_fieldsets
+ for subp in propertynames(getproperty(model, p))
+ getproperty(getproperty(model, p), subp).data.parent .= _arr(model.arch, file["$p/$subp"])
+ end
+ end
+
+ close(file)
+
+ return model
+end

test/test_output_writers.jl

@@ -35,11 +35,65 @@ function run_thermal_bubble_netcdf_tests()
  @test all(S .≈ data(model.tracers.S))
 end
 
+"""
+Run two coarse rising thermal bubble simulations and make sure that when
+restarting from a checkpoint, the restarted simulation matches the non-restarted
+simulation numerically.
+"""
+function run_thermal_bubble_checkpointer_tests(arch)
+ Nx, Ny, Nz = 16, 16, 16
+ Lx, Ly, Lz = 100, 100, 100
+ Δt = 6
+
+ true_model = Model(N=(Nx, Ny, Nz), L=(Lx, Ly, Lz), ν=4e-2, κ=4e-2, arch=arch)
+
+ # Add a cube-shaped warm temperature anomaly that takes up the middle 50%
+ # of the domain volume.
+ i1, i2 = round(Int, Nx/4), round(Int, 3Nx/4)
+ j1, j2 = round(Int, Ny/4), round(Int, 3Ny/4)
+ k1, k2 = round(Int, Nz/4), round(Int, 3Nz/4)
+ true_model.tracers.T.data[i1:i2, j1:j2, k1:k2] .+= 0.01
+
+ checkpointed_model = deepcopy(true_model)
+
+ time_step!(true_model, 9, Δt)
+
+ checkpointer = Checkpointer(output_frequency=5)
+ push!(checkpointed_model.output_writers, checkpointer)
+
+ # Checkpoint should be saved as "test_model_checkpoint_5.jld" after the 5th iteration.
+ time_step!(checkpointed_model, 5, Δt)
+
+ # Remove all knowledge of the checkpointed model.
+ checkpointed_model = nothing
+
+ restored_model = restore_from_checkpoint("checkpoint5.jld2")
+
+ time_step!(restored_model, 4, Δt; adams_bashforth_parameter = n->0.125)
+
+ # Now the true_model and restored_model should be identical.
+ @test all(restored_model.velocities.u.data .≈ true_model.velocities.u.data)
+ @test all(restored_model.velocities.v.data .≈ true_model.velocities.v.data)
+ @test all(restored_model.velocities.w.data .≈ true_model.velocities.w.data)
+ @test all(restored_model.tracers.T.data .≈ true_model.tracers.T.data)
+ @test all(restored_model.tracers.S.data .≈ true_model.tracers.S.data)
+ @test all(restored_model.G.Gu.data .≈ true_model.G.Gu.data)
+ @test all(restored_model.G.Gv.data .≈ true_model.G.Gv.data)
+ @test all(restored_model.G.Gw.data .≈ true_model.G.Gw.data)
+ @test all(restored_model.G.GT.data .≈ true_model.G.GT.data)
+ @test all(restored_model.G.GS.data .≈ true_model.G.GS.data)
+end
+
 @testset "Output writers" begin
  println("Testing output writers...")
 
  @testset "NetCDF" begin
  println(" Testing NetCDF output writer...")
  run_thermal_bubble_netcdf_tests()
  end
+
+ @testset "Checkpointer" begin
+ println(" Testing Checkpointer...")
+ run_thermal_bubble_checkpointer_tests(CPU())
+ end
 end