diff --git a/src/Models/Models.jl b/src/Models/Models.jl index e794370be4..b10bb10f90 100644 --- a/src/Models/Models.jl +++ b/src/Models/Models.jl @@ -21,6 +21,7 @@ using Oceananigans.Utils: Time import Oceananigans: initialize! import Oceananigans.Architectures: architecture +import Oceananigans.TimeSteppers: reset!, set_clock! import Oceananigans.Solvers: iteration import Oceananigans.Simulations: timestepper import Oceananigans.TimeSteppers: reset!, set_clock! diff --git a/src/OutputWriters/checkpointer.jl b/src/OutputWriters/checkpointer.jl index 456b37a5a6..00d79fa681 100644 --- a/src/OutputWriters/checkpointer.jl +++ b/src/OutputWriters/checkpointer.jl @@ -1,17 +1,16 @@ using Glob using Oceananigans -using Oceananigans: fields, prognostic_fields +using Oceananigans: AbstractModel, fields, prognostic_fields using Oceananigans.Fields: offset_data using Oceananigans.TimeSteppers: QuasiAdamsBashforth2TimeStepper import Oceananigans.Fields: set! -mutable struct Checkpointer{T, P} <: AbstractOutputWriter +mutable struct Checkpointer{T} <: AbstractOutputWriter schedule :: T dir :: String prefix :: String - properties :: P overwrite_existing :: Bool verbose :: Bool cleanup :: Bool @@ -19,6 +18,17 @@ end required_checkpoint_properties(model) = [:grid, :clock] +# Certain properties are required for `set!` to pickup from a checkpoint. +function required_checkpointed_properties(model) + properties = [:grid, :clock] + + if has_ab2_timestepper(model) + push!(properties, :timestepper) + end + + return properties +end + """ Checkpointer(model; schedule, @@ -30,16 +40,18 @@ required_checkpoint_properties(model) = [:grid, :clock] properties = required_checkpoint_properties(model)) Construct a `Checkpointer` that checkpoints the model to a JLD2 file on `schedule.` -The `model.clock.iteration` is included in the filename to distinguish between multiple checkpoint files. +The `model.clock.iteration` is included in the filename to distinguish between multiple +checkpoint files. -To restart or "pickup" a model from a checkpoint, specify `pickup = true` when calling `run!`, ensuring -that the checkpoint file is in directory `dir`. See [`run!`](@ref) for more details. +To restart or "pickup" a model from a checkpoint, specify `pickup = true` when +calling `run!`, ensuring that the checkpoint file is in directory `dir`. +See [`run!`](@ref) for more details. Note that extra model `properties` can be specified, but removing crucial properties -such as `:timestepper` will render restoring from the checkpoint impossible. +such as `:timestepper` might render restoring from the checkpoint impossible. The checkpointer attempts to serialize as much of the model to disk as possible, -but functions or objects containing functions cannot be serialized at this time. +but note that functions or objects containing functions cannot be serialized. Keyword arguments ================= @@ -93,7 +105,7 @@ function Checkpointer(model; schedule, mkpath(dir) - return Checkpointer(schedule, dir, prefix, properties, overwrite_existing, verbose, cleanup) + return Checkpointer(schedule, dir, prefix, overwrite_existing, verbose, cleanup) end ##### @@ -158,25 +170,16 @@ end ##### Writing checkpoints ##### -function write_output!(c::Checkpointer, model) +function write_output!(c::Checkpointer, model, addr=checkpointer_address(model)) filepath = checkpoint_path(model.clock.iteration, c) c.verbose && @info "Checkpointing to file $filepath..." - addr = checkpointer_address(model) t1 = time_ns() - jldopen(filepath, "w") do file - file["$addr/checkpointed_properties"] = c.properties - serializeproperties!(file, model, c.properties, addr) - model_fields = prognostic_fields(model) - field_names = keys(model_fields) - for name in field_names - full_address = "$addr/$name" - serializeproperty!(file, full_address, model_fields[name]) - end - end + write_output!(c, model, filepath, "w") t2, sz = time_ns(), filesize(filepath) + c.verbose && @info "Checkpointing done: time=$(prettytime((t2 - t1) * 1e-9)), size=$(pretty_filesize(sz))" c.cleanup && cleanup_checkpoints(c) @@ -184,6 +187,26 @@ function write_output!(c::Checkpointer, model) return nothing end +function write_output!(c, model, filepath::AbstractString, mode::AbstractString; + properties = default_checkpointed_properties(model)) + @show properties + @show model + + properties = validate_checkpointed_properties(model, properties) + addr = checkpointer_address(model) + + jldopen(filepath, mode) do file + file["$addr/checkpointed_properties"] = + serializeproperties!(file, model, properties, addr) + model_fields = prognostic_fields(model) + field_names = keys(model_fields) + for name in field_names + full_address = "$addr/$name" + serializeproperty!(file, full_address, model_fields[name]) + end + end +end + function cleanup_checkpoints(checkpointer) filepaths = glob(checkpoint_superprefix(checkpointer.prefix) * "*.jld2", checkpointer.dir) latest_checkpoint_filepath = latest_checkpoint(checkpointer, filepaths) @@ -197,12 +220,13 @@ end # Should this go in Models? """ - set!(model, filepath::AbstractString) + set!(model::AbstractModel, filepath::AbstractString) Set data in `model.velocities`, `model.tracers`, `model.timestepper.Gⁿ`, and `model.timestepper.G⁻` to checkpointed data stored at `filepath`. """ -function set!(model, filepath::AbstractString) +function set!(model::AbstractModel, filepath::AbstractString) + addr = checkpointer_address(model) jldopen(filepath, "r") do file @@ -225,7 +249,8 @@ function set!(model, filepath::AbstractString) end end - set_time_stepper!(model.timestepper, model.architecture, file, model_fields, addr) + set_time_stepper!(model.timestepper, file, model_fields, addr) + @show model.timestepper.Gⁿ[:u] if !isnothing(model.particles) copyto!(model.particles.properties, file["$addr/particles"]) @@ -251,7 +276,11 @@ function set_time_stepper_tendencies!(timestepper, arch, file, model_fields, add parent_data = on_architecture(arch, file["$addr/timestepper/Gⁿ/$name/data"]) tendencyⁿ_field = timestepper.Gⁿ[name] + @apply_regionally copyto!(parent(tendencyⁿ_field), parent_data) + if name==:u + @show tendencyⁿ_field + end # Tendency "n-1" parent_data = on_architecture(arch, file["$addr/timestepper/G⁻/$name/data"]) diff --git a/src/Simulations/run.jl b/src/Simulations/run.jl index eec31ecf4e..f619c64aef 100644 --- a/src/Simulations/run.jl +++ b/src/Simulations/run.jl @@ -132,7 +132,7 @@ function time_step!(sim::Simulation) sim.Δt end - initial_time_step = !(sim.initialized) + @show initial_time_step = !(sim.initialized) initial_time_step && initialize!(sim) if initial_time_step && sim.verbose @@ -191,13 +191,13 @@ we_want_to_pickup(pickup::String) = true we_want_to_pickup(pickup) = throw(ArgumentError("Cannot run! with pickup=$pickup")) """ - initialize!(sim::Simulation, pickup=false) + initialize!(sim::Simulation) Initialize a simulation: -- Update the auxiliary state of the simulation (filling halo regions, computing auxiliary fields) -- Evaluate all diagnostics, callbacks, and output writers if sim.model.clock.iteration == 0 -- Add diagnostics that "depend" on output writers +- Update the auxiliary state of the simulation (filling halo regions, computing auxiliary fields). +- Evaluate all diagnostics, callbacks, and output writers if `sim.model.clock.iteration == 0`. +- Add diagnostics that "depend" on output writers. """ function initialize!(sim::Simulation) if sim.verbose @@ -207,7 +207,7 @@ function initialize!(sim::Simulation) model = sim.model initialize!(model) - update_state!(model) + update_state!(model, compute_tendencies=true) # Output and diagnostics initialization [add_dependencies!(sim.diagnostics, writer) for writer in values(sim.output_writers)] @@ -253,4 +253,3 @@ function initialize!(sim::Simulation) return nothing end - diff --git a/src/Simulations/simulation.jl b/src/Simulations/simulation.jl index 4e66b1fb8d..a257b8791d 100644 --- a/src/Simulations/simulation.jl +++ b/src/Simulations/simulation.jl @@ -4,7 +4,7 @@ using Oceananigans.DistributedComputations: Distributed, all_reduce using Oceananigans.OutputWriters: JLD2Writer, NetCDFWriter import Oceananigans.Utils: prettytime -import Oceananigans.TimeSteppers: reset! +import Oceananigans.TimeSteppers: reset!, set_clock! import Oceananigans.OutputWriters: write_output! import Oceananigans.Solvers: iteration @@ -34,6 +34,7 @@ end stop_iteration = Inf, stop_time = Inf, wall_time_limit = Inf, + align_time_step = true, minimum_relative_step = 0) Construct a `Simulation` for a `model` with time step `Δt`. @@ -48,6 +49,14 @@ Keyword arguments - `stop_time`: Stop the simulation once this much model clock time has passed. Default: `Inf`. +- `align_time_step`: When `true` it implies that the simulation will automatically adjust the + time-step to meet a constraint imposed by various schedules like `ScheduledTimes`, + `TimeInterval`, `AveragedTimeInterval`, as well as a `stop_time` criterion. + If `false`, i.e., no time-step alignment, then the simulation might blithely step passed + the specified time. Default: `true`. + By `align_time_step = false` we ensure that the time-step does _not_ change within + `time_step!(simulation)` + - `wall_time_limit`: Stop the simulation if it's been running for longer than this many seconds of wall clock time. Default: `Inf`. @@ -196,6 +205,8 @@ function reset!(sim::Simulation) return nothing end +set_clock!(sim::Simulation, new_clock) = set_clock!(sim.model, new_clock) + ##### ##### Default stop criteria callback functions #####