From 573a3666a22059870ea6e202ee7819e20bc27436 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 20 Mar 2025 15:23:09 +1100 Subject: [PATCH 01/33] add SplitExplicit tests for checkpointer --- test/test_checkpointer.jl | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/test/test_checkpointer.jl b/test/test_checkpointer.jl index 177d78f4bb..daf1200a8e 100644 --- a/test/test_checkpointer.jl +++ b/test/test_checkpointer.jl @@ -78,12 +78,8 @@ function test_thermal_bubble_checkpointer_output(arch) return run_checkpointer_tests(true_model, test_model, Δt) end -function test_hydrostatic_splash_checkpointer(arch, free_surface) +function test_hydrostatic_splash_checkpointer(grid, free_surface) # Create and run "true model" - Nx, Ny, Nz = 16, 16, 4 - Lx, Ly, Lz = 1, 1, 1 - - grid = RectilinearGrid(arch, size=(Nx, Ny, Nz), x=(-10, 10), y=(-10, 10), z=(-1, 0)) closure = ScalarDiffusivity(ν=1e-2, κ=1e-2) true_model = HydrostaticFreeSurfaceModel(; grid, free_surface, closure, buoyancy=nothing, tracers=()) test_model = deepcopy(true_model) @@ -196,12 +192,19 @@ end for arch in archs @testset "Checkpointer [$(typeof(arch))]" begin @info " Testing Checkpointer [$(typeof(arch))]..." - test_thermal_bubble_checkpointer_output(arch) - + # test_thermal_bubble_checkpointer_output(arch) + + Nx, Ny, Nz = 16, 16, 4 + Lx, Ly, Lz = 1, 1, 1 + + grid = RectilinearGrid(arch, size=(Nx, Ny, Nz), x=(-10, 10), y=(-10, 10), z=(-1, 0)) + for free_surface in [ExplicitFreeSurface(gravitational_acceleration=1), - ImplicitFreeSurface(gravitational_acceleration=1)] + ImplicitFreeSurface(gravitational_acceleration=1), + SplitExplicitFreeSurface(gravitational_acceleration=1, substeps=5), + SplitExplicitFreeSurface(grid; cfl = 0.7, gravitational_acceleration=1)] - test_hydrostatic_splash_checkpointer(arch, free_surface) + test_hydrostatic_splash_checkpointer(grid, free_surface) end run_checkpointer_cleanup_tests(arch) From 41cb11574e3e1a04e0956cb08d4e436bec8bc3b0 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 20 Mar 2025 15:37:55 +1100 Subject: [PATCH 02/33] add SplitExplicit tests for checkpointer --- test/test_checkpointer.jl | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/test/test_checkpointer.jl b/test/test_checkpointer.jl index daf1200a8e..d07a57d9e1 100644 --- a/test/test_checkpointer.jl +++ b/test/test_checkpointer.jl @@ -192,11 +192,12 @@ end for arch in archs @testset "Checkpointer [$(typeof(arch))]" begin @info " Testing Checkpointer [$(typeof(arch))]..." - # test_thermal_bubble_checkpointer_output(arch) + test_thermal_bubble_checkpointer_output(arch) + + # create a grid to test hydrostatic model Nx, Ny, Nz = 16, 16, 4 Lx, Ly, Lz = 1, 1, 1 - grid = RectilinearGrid(arch, size=(Nx, Ny, Nz), x=(-10, 10), y=(-10, 10), z=(-1, 0)) for free_surface in [ExplicitFreeSurface(gravitational_acceleration=1), From 0a52d4ecaca6c488088f691f4004a48852153b5e Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 20 Mar 2025 19:26:42 +1100 Subject: [PATCH 03/33] expose some checkpointer functionality; don't store properties as a Checkpointer field --- src/OutputWriters/checkpointer.jl | 112 ++++++++++++++++-------------- 1 file changed, 60 insertions(+), 52 deletions(-) diff --git a/src/OutputWriters/checkpointer.jl b/src/OutputWriters/checkpointer.jl index 25ba9a1270..587821ded4 100644 --- a/src/OutputWriters/checkpointer.jl +++ b/src/OutputWriters/checkpointer.jl @@ -5,25 +5,18 @@ using Oceananigans: fields, prognostic_fields using Oceananigans.Fields: offset_data using Oceananigans.TimeSteppers: QuasiAdamsBashforth2TimeStepper -import Oceananigans.Fields: set! +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 end -function default_checkpointed_properties(model) - properties = [:grid, :particles, :clock, :timestepper] - #if has_ab2_timestepper(model) - # push!(properties, :timestepper) - #end - return properties -end +default_checkpointed_properties(model) = [:grid, :particles, :clock, :timestepper] has_ab2_timestepper(model) = try model.timestepper isa QuasiAdamsBashforth2TimeStepper @@ -31,6 +24,41 @@ catch false end +# Certain properties are required for `set!` to pickup from a checkpoint. +function required_checkpointed_properties(model) + properties = [:grid, :particles, :clock] + + if has_ab2_timestepper(model) + push!(properties, :timestepper) + end + + return properties +end + +function validate_properties(model, properties) + required_properties = required_checkpointed_properties(model) + + for rp in required_properties + if rp ∉ properties + @warn "$rp is required for checkpointing. It will be added to checkpointed properties" + push!(properties, rp) + end + end + + for p in properties + p isa Symbol || error("Property $p to be checkpointed must be a Symbol.") + p ∉ propertynames(model) && error("Cannot checkpoint $p, it is not a model property!") + + if (p ∉ required_properties) && has_reference(Function, getproperty(model, p)) + @warn "model.$p contains a function somewhere in its hierarchy and will not be checkpointed." + filter!(e -> e != p, properties) + end + end + + return properties +end + + """ Checkpointer(model; schedule, @@ -79,36 +107,11 @@ function Checkpointer(model; schedule, prefix = "checkpoint", overwrite_existing = false, verbose = false, - cleanup = false, - properties = default_checkpointed_properties(model)) - - # Certain properties are required for `set!` to pickup from a checkpoint. - required_properties = [:grid, :particles, :clock] - - if has_ab2_timestepper(model) - push!(required_properties, :timestepper) - end - - for rp in required_properties - if rp ∉ properties - @warn "$rp is required for checkpointing. It will be added to checkpointed properties" - push!(properties, rp) - end - end - - for p in properties - p isa Symbol || error("Property $p to be checkpointed must be a Symbol.") - p ∉ propertynames(model) && error("Cannot checkpoint $p, it is not a model property!") - - if (p ∉ required_properties) && has_reference(Function, getproperty(model, p)) - @warn "model.$p contains a function somewhere in its hierarchy and will not be checkpointed." - filter!(e -> e != p, properties) - end - end + cleanup = false) mkpath(dir) - return Checkpointer(schedule, dir, prefix, properties, overwrite_existing, verbose, cleanup) + return Checkpointer(schedule, dir, prefix, overwrite_existing, verbose, cleanup) end ##### @@ -177,22 +180,13 @@ end function write_output!(c::Checkpointer, 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) @@ -200,6 +194,22 @@ function write_output!(c::Checkpointer, model) return nothing end +function write_output!(c, model, filepath::AbstractString, mode::AbstractString, properties) + properties = validate_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) @@ -212,13 +222,12 @@ end ##### """ - 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::AbstractModel, filepath::AbstractString) - addr = checkpointer_address(model) jldopen(filepath, "r") do file @@ -280,9 +289,8 @@ function set_time_stepper_tendencies!(timestepper, file, model_fields, addr) return nothing end -# For self-starting timesteppers like RK3 we do nothing +# For self-starting timesteppers like RK3 we do nothing set_time_stepper!(timestepper, args...) = nothing set_time_stepper!(timestepper::QuasiAdamsBashforth2TimeStepper, args...) = set_time_stepper_tendencies!(timestepper, args...) - From a77673093bc50cb95c9a435a4925df88b5b4fa1d Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 20 Mar 2025 19:26:58 +1100 Subject: [PATCH 04/33] use julia v1.10.9 --- .buildkite/pipeline.yml | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 96af4cc68d..8c43b42366 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -1,5 +1,5 @@ env: - JULIA_VERSION: "1.10.8" + JULIA_VERSION: "1.10.9" JULIA_MINOR_VERSION: "1.10" TARTARUS_HOME: "/storage5/buildkite-agent" JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager @@ -209,7 +209,7 @@ steps: architecture: CPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init" @@ -224,7 +224,7 @@ steps: architecture: CPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init" @@ -337,7 +337,7 @@ steps: ##### Turbulence Closures ##### - + - label: "🎣 gpu turbulence closures" env: JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER" @@ -683,7 +683,7 @@ steps: architecture: CPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init" @@ -699,7 +699,7 @@ steps: architecture: CPU retry: automatic: - - exit_status: 1 + - exit_status: 1 limit: 1 depends_on: "init" From b66ae1bf0067bd3017b85c55a8d9bfa9b04a2993 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 20 Mar 2025 19:32:13 +1100 Subject: [PATCH 05/33] pass properties to write_output! --- src/OutputWriters/checkpointer.jl | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/OutputWriters/checkpointer.jl b/src/OutputWriters/checkpointer.jl index 587821ded4..bf2b8a7d23 100644 --- a/src/OutputWriters/checkpointer.jl +++ b/src/OutputWriters/checkpointer.jl @@ -183,7 +183,8 @@ function write_output!(c::Checkpointer, model) t1 = time_ns() - write_output!(c, model, filepath, "w") + write_output!(c, model, filepath, "w"; + properties = default_checkpointed_properties(model)) t2, sz = time_ns(), filesize(filepath) @@ -194,7 +195,7 @@ function write_output!(c::Checkpointer, model) return nothing end -function write_output!(c, model, filepath::AbstractString, mode::AbstractString, properties) +function write_output!(c, model, filepath::AbstractString, mode::AbstractString; properties) properties = validate_properties(model, properties) addr = checkpointer_address(model) From be2206245c8eb05cea061d670eebcd5057f9edb5 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 20 Mar 2025 19:32:35 +1100 Subject: [PATCH 06/33] bump patch release --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 69f65d87f1..3d048eb142 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "Oceananigans" uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09" authors = ["Climate Modeling Alliance and contributors"] -version = "0.96.0" +version = "0.96.1" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" From 0b673770410470cfc040c8bbdd7a7ccbe88bd291 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 20 Mar 2025 19:53:49 +1100 Subject: [PATCH 07/33] validate_properties -> validate_checkpointed_properties --- src/OutputWriters/checkpointer.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/OutputWriters/checkpointer.jl b/src/OutputWriters/checkpointer.jl index bf2b8a7d23..da9a0aaf2f 100644 --- a/src/OutputWriters/checkpointer.jl +++ b/src/OutputWriters/checkpointer.jl @@ -35,7 +35,7 @@ function required_checkpointed_properties(model) return properties end -function validate_properties(model, properties) +function validate_checkpointed_properties(model, properties) required_properties = required_checkpointed_properties(model) for rp in required_properties @@ -196,7 +196,7 @@ function write_output!(c::Checkpointer, model) end function write_output!(c, model, filepath::AbstractString, mode::AbstractString; properties) - properties = validate_properties(model, properties) + properties = validate_checkpointed_properties(model, properties) addr = checkpointer_address(model) jldopen(filepath, mode) do file From fe8e70afd178a363b05f5760987859477f6d6551 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 20 Mar 2025 20:02:22 +1100 Subject: [PATCH 08/33] add default properties kwarg to write_output! --- src/OutputWriters/checkpointer.jl | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/OutputWriters/checkpointer.jl b/src/OutputWriters/checkpointer.jl index da9a0aaf2f..fc2f01803c 100644 --- a/src/OutputWriters/checkpointer.jl +++ b/src/OutputWriters/checkpointer.jl @@ -183,8 +183,7 @@ function write_output!(c::Checkpointer, model) t1 = time_ns() - write_output!(c, model, filepath, "w"; - properties = default_checkpointed_properties(model)) + write_output!(c, model, filepath, "w") t2, sz = time_ns(), filesize(filepath) @@ -195,7 +194,11 @@ function write_output!(c::Checkpointer, model) return nothing end -function write_output!(c, model, filepath::AbstractString, mode::AbstractString; properties) +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) From cbf323d69fdaf27c7bef71320c8eb729d3447d58 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Tue, 25 Mar 2025 11:43:17 +1100 Subject: [PATCH 09/33] remove stray spaces and add backticks --- src/Simulations/run.jl | 25 ++++++++++++------------- 1 file changed, 12 insertions(+), 13 deletions(-) diff --git a/src/Simulations/run.jl b/src/Simulations/run.jl index f2ecb454a8..739919a0b9 100644 --- a/src/Simulations/run.jl +++ b/src/Simulations/run.jl @@ -35,7 +35,7 @@ end """ aligned_time_step(sim, Δt) -Return a time step 'aligned' with `sim.stop_time`, output writer schedules, +Return a time step 'aligned' with `sim.stop_time`, output writer schedules, and callback schedules. Alignment with `sim.stop_time` takes precedence. """ function aligned_time_step(sim::Simulation, Δt) @@ -45,7 +45,7 @@ function aligned_time_step(sim::Simulation, Δt) # Align time step with output writing and callback execution aligned_Δt = schedule_aligned_time_step(sim, aligned_Δt) - + # Align time step with simulation stop time time_left = unit_time(sim.stop_time - clock.time) aligned_Δt = min(aligned_Δt, time_left) @@ -105,7 +105,7 @@ function run!(sim; pickup=false) time_step!(sim) end - for callback in values(sim.callbacks) + for callback in values(sim.callbacks) finalize!(callback, sim) end @@ -137,7 +137,7 @@ function time_step!(sim::Simulation) initial_time_step = !(sim.initialized) initial_time_step && initialize!(sim) - if initial_time_step && sim.verbose + if initial_time_step && sim.verbose @info "Executing initial time step..." start_time = time_ns() end @@ -153,7 +153,7 @@ function time_step!(sim::Simulation) # Callbacks and callback-like things for diag in values(sim.diagnostics) - diag.schedule(sim.model) && run_diagnostic!(diag, sim.model) + diag.schedule(sim.model) && run_diagnostic!(diag, sim.model) end for callback in values(sim.callbacks) @@ -162,7 +162,7 @@ function time_step!(sim::Simulation) end for writer in values(sim.output_writers) - writer.schedule(sim.model) && write_output!(writer, sim.model) + writer.schedule(sim.model) && write_output!(writer, sim.model) end if initial_time_step && sim.verbose @@ -188,14 +188,14 @@ we_want_to_pickup(pickup::Integer) = true 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 @@ -228,7 +228,7 @@ function initialize!(sim::Simulation) run_diagnostic!(diag, model) end - for callback in values(sim.callbacks) + for callback in values(sim.callbacks) callback.callsite isa TimeStepCallsite && callback(sim) end @@ -247,4 +247,3 @@ function initialize!(sim::Simulation) return nothing end - From 40dccba87a5e416df986084eb425faeb8888111b Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Tue, 25 Mar 2025 11:45:26 +1100 Subject: [PATCH 10/33] update docstring --- src/TimeSteppers/clock.jl | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/TimeSteppers/clock.jl b/src/TimeSteppers/clock.jl index 49bbe67516..e945185a14 100644 --- a/src/TimeSteppers/clock.jl +++ b/src/TimeSteppers/clock.jl @@ -7,9 +7,9 @@ import Base: show import Oceananigans.Units: Time """ - mutable struct Clock{T, FT} + mutable struct Clock{TT, DT, IT, S} -Keeps track of the current `time`, `last_Δt`, `iteration` number, and time-stepping `stage`. +Keep track of the current `time`, `last_Δt`, `iteration` number, and time-stepping `stage`. The `stage` is updated only for multi-stage time-stepping methods. The `time::T` is either a number or a `DateTime` object. """ @@ -121,5 +121,3 @@ Adapt.adapt_structure(to, clock::Clock) = (time = clock.time, last_stage_Δt = clock.last_stage_Δt, iteration = clock.iteration, stage = clock.stage) - - From 38ef5b67e978b0ac4510abd9469d9a7bbb0cb7bb Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Tue, 25 Mar 2025 11:46:08 +1100 Subject: [PATCH 11/33] code alignment --- test/test_grids.jl | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/test/test_grids.jl b/test/test_grids.jl index 18ae3c81c3..70e3b1f305 100644 --- a/test/test_grids.jl +++ b/test/test_grids.jl @@ -5,7 +5,7 @@ using Oceananigans.Grids: total_extent, ColumnEnsembleSize, xspacings, yspacings, zspacings, xnode, ynode, znode, λnode, φnode, λspacings, φspacings - + using Oceananigans.OrthogonalSphericalShellGrids: RotatedLatitudeLongitudeGrid using Oceananigans.Operators: Δx, Δy, Δz, Δλ, Δφ, Ax, Ay, Az, volume @@ -51,7 +51,7 @@ function test_regular_rectilinear_correct_coordinate_lengths(FT) @test length(grid.yᵃᶜᵃ) == Ny + 2Hy @test length(grid.xᶠᵃᵃ) == Nx + 2Hx @test length(grid.yᵃᶠᵃ) == Ny + 2Hy + 1 - + @test length(grid.z.cᵃᵃᶜ) == Nz + 2Hz @test length(grid.z.cᵃᵃᶠ) == Nz + 2Hz + 1 @@ -129,7 +129,7 @@ function test_regular_rectilinear_ranges_have_correct_length(FT) @test length(grid.yᵃᶜᵃ) == Ny + 2Hy @test length(grid.xᶠᵃᵃ) == Nx + 1 + 2Hx @test length(grid.yᵃᶠᵃ) == Ny + 1 + 2Hy - + @test length(grid.z.cᵃᵃᶜ) == Nz + 2Hz @test length(grid.z.cᵃᵃᶠ) == Nz + 1 + 2Hz @@ -154,7 +154,7 @@ function test_regular_rectilinear_grid_properties_are_same_type(FT) # Do this test two ways, one with defaults and one with explicit FT for method in (:explicit, :with_default) - if method === :explicit + if method === :explicit grid = RectilinearGrid(CPU(), FT, size=(10, 10, 10), extent=(1, 1//7, 2π)) elseif method === :with_default FT₀ = Oceananigans.defaults.FloatType @@ -349,7 +349,7 @@ function test_grid_equality(arch) grid2 = RectilinearGrid(arch, topology=topo, size=(Nx, Ny, Nz), x=(0, 1), y=(-1, 1), z=0:Nz) grid3 = RectilinearGrid(arch, topology=topo, size=(Nx, Ny, Nz), x=(0, 1), y=(-1, 1), z=0:Nz) - return grid1==grid1 && grid2 == grid3 && grid1 !== grid3 + return grid1 == grid1 && grid2 == grid3 && grid1 !== grid3 end function test_grid_equality_over_architectures() From 3aeeed7bb835dcab164da40a8437c367b2da1245 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Tue, 25 Mar 2025 21:56:25 +1100 Subject: [PATCH 12/33] =?UTF-8?q?set=5Fclock!=20+=20clock.last=5Fstage=5F?= =?UTF-8?q?=CE=94t,=20clock.last=5F=CE=94t=20in=20tick!(clock,=20=CE=94t)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/TimeSteppers/clock.jl | 27 ++++++++++++++++++++- src/TimeSteppers/quasi_adams_bashforth_2.jl | 6 ++--- src/TimeSteppers/runge_kutta_3.jl | 11 +++------ 3 files changed, 32 insertions(+), 12 deletions(-) diff --git a/src/TimeSteppers/clock.jl b/src/TimeSteppers/clock.jl index e945185a14..fb4c39aa03 100644 --- a/src/TimeSteppers/clock.jl +++ b/src/TimeSteppers/clock.jl @@ -30,10 +30,33 @@ function reset!(clock::Clock{TT, DT, IT, S}) where {TT, DT, IT, S} return nothing end +""" + set_clock!(clock1::Clock, clock2::Clock) + +Set `clock1` to the `clock2`. +""" +function set_clock!(clock1::Clock, clock2::Clock) + clock1.time = clock2.time + clock1.iteration = clock2.iteration + clock1.last_Δt = clock2.last_Δt + clock1.last_stage_Δt = clock2.last_stage_Δt + clock1.stage = clock2.stage + + return nothing +end + +function Base.:(==)(clock1::Clock, clock2::Clock) + return clock1.time == clock2.time && + clock1.iteration == clock2.iteration && + clock1.last_Δt == clock2.last_Δt && + clock1.last_stage_Δt == clock2.last_stage_Δt && + clock1.stage == clock2.stage +end + """ Clock(; time, last_Δt=Inf, last_stage_Δt=Inf, iteration=0, stage=1) -Returns a `Clock` object. By default, `Clock` is initialized to the zeroth `iteration` +Return a `Clock` object. By default, `Clock` is initialized to the zeroth `iteration` and first time step `stage` with `last_Δt=last_stage_Δt=Inf`. """ function Clock(; time, @@ -107,9 +130,11 @@ function tick!(clock, Δt; stage=false) if stage # tick a stage update clock.stage += 1 + clock.last_stage_Δt = Δt else # tick an iteration and reset stage clock.iteration += 1 clock.stage = 1 + clock.last_Δt = Δt end return nothing diff --git a/src/TimeSteppers/quasi_adams_bashforth_2.jl b/src/TimeSteppers/quasi_adams_bashforth_2.jl index d02a018fd1..4c4002ee22 100644 --- a/src/TimeSteppers/quasi_adams_bashforth_2.jl +++ b/src/TimeSteppers/quasi_adams_bashforth_2.jl @@ -15,7 +15,7 @@ end G⁻ = map(similar, prognostic_fields)) Return a 2nd-order quasi Adams-Bashforth (AB2) time stepper (`QuasiAdamsBashforth2TimeStepper`) -on `grid`, with `tracers`, and AB2 parameter `χ`. The tendency fields `Gⁿ` and `G⁻`, usually equal to +on `grid`, with `tracers`, and AB2 parameter `χ`. The tendency fields `Gⁿ` and `G⁻`, usually equal to the prognostic_fields passed as positional argument, can be specified via optional `kwargs`. The 2nd-order quasi Adams-Bashforth timestepper steps forward the state `Uⁿ` by `Δt` via @@ -99,9 +99,7 @@ function time_step!(model::AbstractModel{<:QuasiAdamsBashforth2TimeStepper}, Δt ab2_step!(model, Δt) tick!(model.clock, Δt) - model.clock.last_Δt = Δt - model.clock.last_stage_Δt = Δt # just one stage - + calculate_pressure_correction!(model, Δt) @apply_regionally correct_velocities_and_cache_previous_tendencies!(model, Δt) diff --git a/src/TimeSteppers/runge_kutta_3.jl b/src/TimeSteppers/runge_kutta_3.jl index bb3aac194a..5a66025102 100644 --- a/src/TimeSteppers/runge_kutta_3.jl +++ b/src/TimeSteppers/runge_kutta_3.jl @@ -51,7 +51,7 @@ function RungeKutta3TimeStepper(grid, prognostic_fields; G⁻ = map(similar, prognostic_fields)) where {TI, TG} !isnothing(implicit_solver) && - @warn("Implicit-explicit time-stepping with RungeKutta3TimeStepper is not tested. " * + @warn("Implicit-explicit time-stepping with RungeKutta3TimeStepper is not tested. " * "\n implicit_solver: $(typeof(implicit_solver))") γ¹ = 8 // 15 @@ -105,7 +105,6 @@ function time_step!(model::AbstractModel{<:RungeKutta3TimeStepper}, Δt; callbac rk3_substep!(model, Δt, γ¹, nothing) tick!(model.clock, first_stage_Δt; stage=true) - model.clock.last_stage_Δt = first_stage_Δt calculate_pressure_correction!(model, first_stage_Δt) pressure_correct_velocities!(model, first_stage_Δt) @@ -121,7 +120,6 @@ function time_step!(model::AbstractModel{<:RungeKutta3TimeStepper}, Δt; callbac rk3_substep!(model, Δt, γ², ζ²) tick!(model.clock, second_stage_Δt; stage=true) - model.clock.last_stage_Δt = second_stage_Δt calculate_pressure_correction!(model, second_stage_Δt) pressure_correct_velocities!(model, second_stage_Δt) @@ -133,21 +131,20 @@ function time_step!(model::AbstractModel{<:RungeKutta3TimeStepper}, Δt; callbac # # Third stage # - + rk3_substep!(model, Δt, γ³, ζ³) # This adjustment of the final time-step reduces the accumulation of - # round-off error when Δt is added to model.clock.time. Note that we still use + # round-off error when Δt is added to model.clock.time. Note that we still use # third_stage_Δt for the substep, pressure correction, and Lagrangian particles step. corrected_third_stage_Δt = tⁿ⁺¹ - model.clock.time tick!(model.clock, third_stage_Δt) model.clock.last_stage_Δt = corrected_third_stage_Δt - model.clock.last_Δt = Δt calculate_pressure_correction!(model, third_stage_Δt) pressure_correct_velocities!(model, third_stage_Δt) - + update_state!(model, callbacks; compute_tendencies = true) step_lagrangian_particles!(model, third_stage_Δt) From c3c007feaab1794b4bdea2c869a0706a895d5c0e Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Tue, 25 Mar 2025 22:00:11 +1100 Subject: [PATCH 13/33] add set_clock! for OceananigansModels --- src/Models/Models.jl | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/src/Models/Models.jl b/src/Models/Models.jl index c2b46d28f0..be12d154f6 100644 --- a/src/Models/Models.jl +++ b/src/Models/Models.jl @@ -20,7 +20,7 @@ using Oceananigans.Utils: Time import Oceananigans: initialize! import Oceananigans.Architectures: architecture -import Oceananigans.TimeSteppers: reset! +import Oceananigans.TimeSteppers: reset!, set_clock! import Oceananigans.Solvers: iteration # A prototype interface for AbstractModel. @@ -116,6 +116,8 @@ const OceananigansModels = Union{HydrostaticFreeSurfaceModel, NonhydrostaticModel, ShallowWaterModel} +set_clock!(model::OceananigansModels, new_clock) = set_clock!(model.clock, new_clock) + """ possible_field_time_series(model::HydrostaticFreeSurfaceModel) @@ -128,10 +130,10 @@ function possible_field_time_series(model::OceananigansModels) # such as model.diffusivity_fields return tuple(model_fields, forcing) end - -# Update _all_ `FieldTimeSeries`es in an `OceananigansModel`. + +# Update _all_ `FieldTimeSeries`es in an `OceananigansModel`. # Extract `FieldTimeSeries` from all property names that might contain a `FieldTimeSeries` -# Flatten the resulting tuple by extracting unique values and set! them to the +# Flatten the resulting tuple by extracting unique values and set! them to the # correct time range by looping over them function update_model_field_time_series!(model::OceananigansModels, clock::Clock) time = Time(clock.time) @@ -146,7 +148,7 @@ function update_model_field_time_series!(model::OceananigansModels, clock::Clock return nothing end - + import Oceananigans.TimeSteppers: reset! function reset!(model::OceananigansModels) @@ -162,7 +164,7 @@ function reset!(model::OceananigansModels) for field in model.timestepper.Gⁿ fill!(field, 0) end - + return nothing end @@ -170,7 +172,7 @@ end function default_nan_checker(model::OceananigansModels) model_fields = prognostic_fields(model) - if isempty(model_fields) + if isempty(model_fields) return nothing end @@ -182,7 +184,7 @@ end using Oceananigans.Models.HydrostaticFreeSurfaceModels: OnlyParticleTrackingModel -# Particle tracking models with prescribed velocities (and no tracers) +# Particle tracking models with prescribed velocities (and no tracers) # have no prognostic fields and no chance to producing a NaN. default_nan_checker(::OnlyParticleTrackingModel) = nothing From a21ff8ae3198fc1d34752f1e6b3cc23a66328ee9 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Tue, 25 Mar 2025 22:04:33 +1100 Subject: [PATCH 14/33] =?UTF-8?q?clock.last=5F=CE=94t=20=3D=20=CE=94t=20is?= =?UTF-8?q?=20part=20of=20tick!(clock,=20=CE=94t)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/TimeSteppers/split_hydrostatic_runge_kutta_3.jl | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/src/TimeSteppers/split_hydrostatic_runge_kutta_3.jl b/src/TimeSteppers/split_hydrostatic_runge_kutta_3.jl index df7f141d60..269d83560a 100644 --- a/src/TimeSteppers/split_hydrostatic_runge_kutta_3.jl +++ b/src/TimeSteppers/split_hydrostatic_runge_kutta_3.jl @@ -28,8 +28,8 @@ end Return a 3rd-order `SplitRungeKutta3TimeStepper` on `grid` and with `tracers`. The tendency fields `Gⁿ` and `G⁻`, and the previous state ` Ψ⁻` can be modified via optional `kwargs`. -The scheme described by [Lan2022](@citet). In a nutshell, the 3rd-order Runge Kutta timestepper -steps forward the state `Uⁿ` by `Δt` via 3 substeps. A barotropic velocity correction step is applied +The scheme described by [Lan2022](@citet). In a nutshell, the 3rd-order Runge Kutta timestepper +steps forward the state `Uⁿ` by `Δt` via 3 substeps. A barotropic velocity correction step is applied after at each substep. The state `U` after each substep `m` is @@ -56,7 +56,7 @@ function SplitRungeKutta3TimeStepper(grid, prognostic_fields, args...; "Use at own risk, and report any issues encountered.") !isnothing(implicit_solver) && - @warn("Implicit-explicit time-stepping with SplitRungeKutta3TimeStepper is not tested. " * + @warn("Implicit-explicit time-stepping with SplitRungeKutta3TimeStepper is not tested. " * "\n implicit_solver: $(typeof(implicit_solver))") γ² = 1 // 4 @@ -112,7 +112,7 @@ function time_step!(model::AbstractModel{<:SplitRungeKutta3TimeStepper}, Δt; ca #### model.clock.stage = 3 - + split_rk3_substep!(model, Δt, γ³, ζ³) calculate_pressure_correction!(model, Δt) pressure_correct_velocities!(model, Δt) @@ -121,7 +121,6 @@ function time_step!(model::AbstractModel{<:SplitRungeKutta3TimeStepper}, Δt; ca step_lagrangian_particles!(model, Δt) tick!(model.clock, Δt) - model.clock.last_Δt = Δt return nothing end @@ -160,10 +159,10 @@ function split_rk3_substep!(model, Δt, γⁿ, ζⁿ) end function cache_previous_fields!(model) - + previous_fields = model.timestepper.Ψ⁻ model_fields = prognostic_fields(model) - + for name in keys(previous_fields) if !isnothing(previous_fields[name]) parent(previous_fields[name]) .= parent(model_fields[name]) # Storing also the halos From 3329188bbbd0d5f0538b2fcbcf0a23b666190441 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Wed, 26 Mar 2025 10:19:30 +1100 Subject: [PATCH 15/33] add set_clock!(::Simulation, clock) --- src/Simulations/simulation.jl | 27 ++++++++++++++++----------- 1 file changed, 16 insertions(+), 11 deletions(-) diff --git a/src/Simulations/simulation.jl b/src/Simulations/simulation.jl index b00ee3b1c9..a7edec5a3d 100644 --- a/src/Simulations/simulation.jl +++ b/src/Simulations/simulation.jl @@ -4,7 +4,7 @@ using Oceananigans.DistributedComputations: Distributed, all_reduce import Oceananigans.Models: iteration import Oceananigans.Utils: prettytime -import Oceananigans.TimeSteppers: reset! +import Oceananigans.TimeSteppers: reset!, set_clock! default_progress(simulation) = nothing @@ -27,11 +27,12 @@ end """ Simulation(model; Δt, - verbose = true, - stop_iteration = Inf, - stop_time = Inf, - wall_time_limit = Inf, - minimum_relative_step = 0) + verbose = true, + 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`. @@ -116,6 +117,8 @@ function Base.show(io::IO, s::Simulation) "└── Diagnostics: $(ordered_dict_show(s.diagnostics, "│"))") end +set_clock!(sim::Simulation, new_clock) = set_clock!(sim.model, new_clock) + ##### ##### Utilities ##### @@ -184,6 +187,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 ##### @@ -194,11 +199,11 @@ function stop_iteration_exceeded(sim) if sim.model.clock.iteration >= sim.stop_iteration if sim.verbose msg = string("Model iteration ", iteration(sim), " equals or exceeds stop iteration ", Int(sim.stop_iteration), ".") - @info wall_time_msg(sim) + @info wall_time_msg(sim) @info msg end - sim.running = false + sim.running = false end return nothing @@ -208,11 +213,11 @@ function stop_time_exceeded(sim) if sim.model.clock.time >= sim.stop_time if sim.verbose msg = string("Simulation time ", prettytime(sim), " equals or exceeds stop time ", prettytime(sim.stop_time), ".") - @info wall_time_msg(sim) + @info wall_time_msg(sim) @info msg end - sim.running = false + sim.running = false end return nothing @@ -226,7 +231,7 @@ function wall_time_limit_exceeded(sim) @info msg end - sim.running = false + sim.running = false end return nothing From 2f6ef7d734ec429b8bf7d878e75fa200cdd708ff Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Wed, 26 Mar 2025 10:46:04 +1100 Subject: [PATCH 16/33] add docs for align_time_step --- src/Simulations/simulation.jl | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/Simulations/simulation.jl b/src/Simulations/simulation.jl index a7edec5a3d..65f4df98c6 100644 --- a/src/Simulations/simulation.jl +++ b/src/Simulations/simulation.jl @@ -46,6 +46,14 @@ Keyword arguments - `stop_time`: Stop the simulation once this much model clock time has passed. +- `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. - `minimum_relative_step`: time steps smaller than `Δt * minimum_relative_step` will be skipped. From 96286b17cf2fb9c5f89b276ee240215491eb7712 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Fri, 25 Apr 2025 08:41:34 +1000 Subject: [PATCH 17/33] wip --- ...ute_hydrostatic_free_surface_tendencies.jl | 6 +-- .../hydrostatic_free_surface_model.jl | 4 +- .../hydrostatic_free_surface_rk3_step.jl | 54 +++++++++---------- ...te_hydrostatic_free_surface_model_state.jl | 18 +++---- src/OutputWriters/checkpointer.jl | 20 ++++--- src/Simulations/run.jl | 4 +- 6 files changed, 56 insertions(+), 50 deletions(-) diff --git a/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl b/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl index 8003321119..36cc7aa7f2 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl @@ -10,7 +10,7 @@ using Oceananigans.Fields: immersed_boundary_condition using Oceananigans.Biogeochemistry: update_tendencies! using Oceananigans.TurbulenceClosures.TKEBasedVerticalDiffusivities: FlavorOfCATKE, FlavorOfTD -using Oceananigans.ImmersedBoundaries: get_active_cells_map, ActiveInteriorIBG, +using Oceananigans.ImmersedBoundaries: get_active_cells_map, ActiveInteriorIBG, linear_index_to_tuple """ @@ -147,12 +147,12 @@ function compute_hydrostatic_momentum_tendencies!(model, velocities, kernel_para v_kernel_args = tuple(start_momentum_kernel_args..., v_immersed_bc, end_momentum_kernel_args..., v_forcing) launch!(arch, grid, kernel_parameters, - compute_hydrostatic_free_surface_Gu!, model.timestepper.Gⁿ.u, grid, + compute_hydrostatic_free_surface_Gu!, model.timestepper.Gⁿ.u, grid, active_cells_map, u_kernel_args; active_cells_map) launch!(arch, grid, kernel_parameters, - compute_hydrostatic_free_surface_Gv!, model.timestepper.Gⁿ.v, grid, + compute_hydrostatic_free_surface_Gv!, model.timestepper.Gⁿ.v, grid, active_cells_map, v_kernel_args; active_cells_map) diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl index 17bce5fe7c..4211218edb 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl @@ -96,7 +96,7 @@ Keyword arguments preallocated `CenterField`s. - `forcing`: `NamedTuple` of user-defined forcing functions that contribute to solution tendencies. - `closure`: The turbulence closure for `model`. See `Oceananigans.TurbulenceClosures`. - - `timestepper`: A symbol that specifies the time-stepping method. + - `timestepper`: A symbol that specifies the time-stepping method. Either `:QuasiAdamsBashforth2` (default) or `:SplitRungeKutta3`. - `boundary_conditions`: `NamedTuple` containing field boundary conditions. - `particles`: Lagrangian particles to be advected with the flow. Default: `nothing`. @@ -213,7 +213,7 @@ function HydrostaticFreeSurfaceModel(; grid, # Regularize forcing for model tracer and velocity fields. model_fields = merge(prognostic_fields, auxiliary_fields) forcing = model_forcing(model_fields; forcing...) - + model = HydrostaticFreeSurfaceModel(arch, grid, clock, advection, buoyancy, coriolis, free_surface, forcing, closure, particles, biogeochemistry, velocities, tracers, pressure, diffusivity_fields, timestepper, auxiliary_fields, vertical_coordinate) diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_rk3_step.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_rk3_step.jl index a8e9eeef40..53cc249c3a 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_rk3_step.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_rk3_step.jl @@ -5,11 +5,11 @@ using Oceananigans.ImmersedBoundaries: get_active_cells_map, get_active_column_m import Oceananigans.TimeSteppers: split_rk3_substep!, _split_rk3_substep_field!, cache_previous_fields! function split_rk3_substep!(model::HydrostaticFreeSurfaceModel, Δt, γⁿ, ζⁿ) - + grid = model.grid timestepper = model.timestepper free_surface = model.free_surface - + compute_free_surface_tendency!(grid, model, free_surface) rk3_substep_velocities!(model.velocities, model, Δt, γⁿ, ζⁿ) @@ -22,7 +22,7 @@ function split_rk3_substep!(model::HydrostaticFreeSurfaceModel, Δt, γⁿ, ζ if model.clock.stage == 2 rk3_average_free_surface!(free_surface, grid, timestepper, γⁿ, ζⁿ) end - + return nothing end @@ -31,11 +31,11 @@ rk3_average_free_surface!(free_surface, args...) = nothing function rk3_average_free_surface!(free_surface::ImplicitFreeSurface, grid, timestepper, γⁿ, ζⁿ) arch = architecture(grid) - ηⁿ⁻¹ = timestepper.Ψ⁻.η - ηⁿ = free_surface.η - + ηⁿ⁻¹ = timestepper.Ψ⁻.η + ηⁿ = free_surface.η + launch!(arch, grid, :xy, _rk3_average_free_surface!, ηⁿ, grid, ηⁿ⁻¹, γⁿ, ζⁿ) - + return nothing end @@ -45,25 +45,25 @@ function rk3_average_free_surface!(free_surface::SplitExplicitFreeSurface, grid, Uⁿ⁻¹ = timestepper.Ψ⁻.U Vⁿ⁻¹ = timestepper.Ψ⁻.V - ηⁿ⁻¹ = timestepper.Ψ⁻.η + ηⁿ⁻¹ = timestepper.Ψ⁻.η Uⁿ = free_surface.barotropic_velocities.U Vⁿ = free_surface.barotropic_velocities.V - ηⁿ = free_surface.η + ηⁿ = free_surface.η launch!(arch, grid, :xy, _rk3_average_free_surface!, Uⁿ, grid, Uⁿ⁻¹, γⁿ, ζⁿ) launch!(arch, grid, :xy, _rk3_average_free_surface!, Vⁿ, grid, Vⁿ⁻¹, γⁿ, ζⁿ) - - # Averaging the free surface is only required for a grid with Mutable vertical coordinates, - # which needs to update the grid based on the value of the free surface + + # Averaging the free surface is only required for a grid with Mutable vertical coordinates, + # which needs to update the grid based on the value of the free surface launch!(arch, grid, :xy, _rk3_average_free_surface!, ηⁿ, grid, ηⁿ⁻¹, γⁿ, ζⁿ) return nothing end -@kernel function _rk3_average_free_surface!(η, grid, η⁻, γⁿ, ζⁿ) +@kernel function _rk3_average_free_surface!(η, grid, η⁻, γⁿ, ζⁿ) i, j = @index(Global, NTuple) k = grid.Nz + 1 - @inbounds η[i, j, k] = ζⁿ * η⁻[i, j, k] + γⁿ * η[i, j, k] + @inbounds η[i, j, k] = ζⁿ * η⁻[i, j, k] + γⁿ * η[i, j, k] end ##### @@ -85,7 +85,7 @@ function rk3_substep_velocities!(velocities, model, Δt, γⁿ, ζⁿ) model.closure, model.diffusivity_fields, nothing, - model.clock, + model.clock, Δt) end @@ -106,7 +106,7 @@ function rk3_substep_tracers!(tracers, model, Δt, γⁿ, ζⁿ) # Tracer update kernels for (tracer_index, tracer_name) in enumerate(propertynames(tracers)) - + Gⁿ = model.timestepper.Gⁿ[tracer_name] Ψ⁻ = model.timestepper.Ψ⁻[tracer_name] θ = tracers[tracer_name] @@ -128,39 +128,39 @@ function rk3_substep_tracers!(tracers, model, Δt, γⁿ, ζⁿ) end ##### -##### Tracer update in mutable vertical coordinates +##### Tracer update in mutable vertical coordinates ##### -# σθ is the evolved quantity. -# We store temporarily σθ in θ. Once σⁿ⁺¹ is known we can retrieve θⁿ⁺¹ +# σθ is the evolved quantity. +# We store temporarily σθ in θ. Once σⁿ⁺¹ is known we can retrieve θⁿ⁺¹ # with the `unscale_tracers!` function. Ψ⁻ is the previous tracer already scaled # by the vertical coordinate scaling factor: ψ⁻ = σ * θ -@kernel function _split_rk3_substep_tracer_field!(θ, grid, Δt, γⁿ, ζⁿ, Gⁿ, Ψ⁻) +@kernel function _split_rk3_substep_tracer_field!(θ, grid, Δt, γⁿ, ζⁿ, Gⁿ, Ψ⁻) i, j, k = @index(Global, NTuple) σᶜᶜⁿ = σⁿ(i, j, k, grid, Center(), Center(), Center()) @inbounds θ[i, j, k] = ζⁿ * Ψ⁻[i, j, k] + γⁿ * σᶜᶜⁿ * (θ[i, j, k] + Δt * Gⁿ[i, j, k]) end -# We store temporarily σθ in θ. +# We store temporarily σθ in θ. # The unscaled θ will be retrieved with `unscale_tracers!` -@kernel function _split_rk3_substep_tracer_field!(θ, grid, Δt, ::Nothing, ::Nothing, Gⁿ, Ψ⁻) +@kernel function _split_rk3_substep_tracer_field!(θ, grid, Δt, ::Nothing, ::Nothing, Gⁿ, Ψ⁻) i, j, k = @index(Global, NTuple) @inbounds θ[i, j, k] = Ψ⁻[i, j, k] + Δt * Gⁿ[i, j, k] * σⁿ(i, j, k, grid, Center(), Center(), Center()) end -##### +##### ##### Storing previous fields for the RK3 update -##### +##### # Tracers are multiplied by the vertical coordinate scaling factor -@kernel function _cache_tracer_fields!(Ψ⁻, grid, Ψⁿ) +@kernel function _cache_tracer_fields!(Ψ⁻, grid, Ψⁿ) i, j, k = @index(Global, NTuple) @inbounds Ψ⁻[i, j, k] = Ψⁿ[i, j, k] * σⁿ(i, j, k, grid, Center(), Center(), Center()) end function cache_previous_fields!(model::HydrostaticFreeSurfaceModel) - + previous_fields = model.timestepper.Ψ⁻ model_fields = prognostic_fields(model) grid = model.grid @@ -172,7 +172,7 @@ function cache_previous_fields!(model::HydrostaticFreeSurfaceModel) if name ∈ keys(model.tracers) # Tracers are stored with the grid scaling launch!(arch, grid, :xyz, _cache_tracer_fields!, Ψ⁻, grid, Ψⁿ) else # Velocities and free surface are stored without the grid scaling - parent(Ψ⁻) .= parent(Ψⁿ) + parent(Ψ⁻) .= parent(Ψⁿ) end end diff --git a/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl b/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl index 7ab34ff634..93c52e8b1b 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl @@ -23,7 +23,7 @@ compute_auxiliary_fields!(auxiliary_fields) = Tuple(compute!(a) for a in auxilia Update peripheral aspects of the model (auxiliary fields, halo regions, diffusivities, hydrostatic pressure) to the current model state. If `callbacks` are provided (in an array), -they are called in the end. Finally, the tendencies for the new time-step are computed if +they are called in the end. Finally, the tendencies for the new time-step are computed if `compute_tendencies = true`. """ update_state!(model::HydrostaticFreeSurfaceModel, callbacks=[]; compute_tendencies = true) = @@ -50,7 +50,7 @@ function update_state!(model::HydrostaticFreeSurfaceModel, grid, callbacks; comp update_biogeochemical_state!(model.biogeochemistry, model) - compute_tendencies && + compute_tendencies && @apply_regionally compute_tendencies!(model, callbacks) return nothing @@ -58,7 +58,7 @@ end # Mask immersed fields function mask_immersed_model_fields!(model, grid) - η = displacement(model.free_surface) + η = displacement(model.free_surface) fields_to_mask = merge(model.auxiliary_fields, prognostic_fields(model)) foreach(fields_to_mask) do field @@ -67,22 +67,22 @@ function mask_immersed_model_fields!(model, grid) end end mask_immersed_field_xy!(η, k=size(grid, 3)+1) - + return nothing end function compute_auxiliaries!(model::HydrostaticFreeSurfaceModel; w_parameters = w_kernel_parameters(model.grid), p_parameters = p_kernel_parameters(model.grid), - κ_parameters = :xyz) - + κ_parameters = :xyz) + grid = model.grid closure = model.closure tracers = model.tracers diffusivity = model.diffusivity_fields buoyancy = model.buoyancy - + P = model.pressure.pHY′ - arch = architecture(grid) + arch = architecture(grid) # Update the grid and unscale the tracers update_grid!(model, grid, model.vertical_coordinate; parameters = w_parameters) @@ -94,6 +94,6 @@ function compute_auxiliaries!(model::HydrostaticFreeSurfaceModel; w_parameters = # Update closure diffusivities compute_diffusivities!(diffusivity, closure, model; parameters = κ_parameters) - + return nothing end diff --git a/src/OutputWriters/checkpointer.jl b/src/OutputWriters/checkpointer.jl index fc2f01803c..62ed642737 100644 --- a/src/OutputWriters/checkpointer.jl +++ b/src/OutputWriters/checkpointer.jl @@ -70,16 +70,18 @@ end properties = default_checkpointed_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 ================= @@ -99,8 +101,8 @@ Keyword arguments Default: `false`. - `properties`: List of model properties to checkpoint. This list _must_ contain - `:grid`, `:particles` and :clock`, and if using AB2 timestepping then also - `:timestepper`. Default: default_checkpointed_properties(model) + `:grid`, `:particles` and ``:clock`, and if using AB2 timestepping then also + `:timestepper`. Default: `default_checkpointed_properties(model)`. """ function Checkpointer(model; schedule, dir = ".", @@ -255,6 +257,7 @@ function set!(model::AbstractModel, filepath::AbstractString) end 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"]) @@ -279,6 +282,9 @@ function set_time_stepper_tendencies!(timestepper, file, model_fields, addr) tendencyⁿ_field = timestepper.Gⁿ[name] copyto!(tendencyⁿ_field.data.parent, parent_data) + if name==:u + @show tendencyⁿ_field + end # Tendency "n-1" parent_data = file["$addr/timestepper/G⁻/$name/data"] diff --git a/src/Simulations/run.jl b/src/Simulations/run.jl index 739919a0b9..a3eaff58a8 100644 --- a/src/Simulations/run.jl +++ b/src/Simulations/run.jl @@ -134,7 +134,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 @@ -205,7 +205,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)] From 847c95350340d0b7d9d2e097ade67555a2b923d3 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 26 Jun 2025 10:05:09 +1000 Subject: [PATCH 18/33] Update Project.toml --- Project.toml | 4 ---- 1 file changed, 4 deletions(-) diff --git a/Project.toml b/Project.toml index 71d8bfdf4e..c4b7f63ec8 100644 --- a/Project.toml +++ b/Project.toml @@ -1,11 +1,7 @@ name = "Oceananigans" uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09" authors = ["Climate Modeling Alliance and contributors"] -<<<<<<< ncc/checkopointer-shenanigans-2 -version = "0.96.28" -======= version = "0.96.33" ->>>>>>> main [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" From 635550bf4f441ae2d6034925953849f101219b19 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 26 Jun 2025 10:06:05 +1000 Subject: [PATCH 19/33] Update src/OutputWriters/checkpointer.jl Co-authored-by: Simone Silvestri --- src/OutputWriters/checkpointer.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/OutputWriters/checkpointer.jl b/src/OutputWriters/checkpointer.jl index 2409a31519..be63a1d8e5 100644 --- a/src/OutputWriters/checkpointer.jl +++ b/src/OutputWriters/checkpointer.jl @@ -26,7 +26,7 @@ end # Certain properties are required for `set!` to pickup from a checkpoint. function required_checkpointed_properties(model) - properties = [:grid, :particles, :clock] + properties = [:grid, :clock] if has_ab2_timestepper(model) push!(properties, :timestepper) From 0ca81dd7e17412ce4cf19096e5240c96d61aeaf1 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 26 Jun 2025 10:06:44 +1000 Subject: [PATCH 20/33] Apply suggestions from code review Co-authored-by: Gregory L. Wagner --- src/OutputWriters/checkpointer.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/OutputWriters/checkpointer.jl b/src/OutputWriters/checkpointer.jl index be63a1d8e5..911200f213 100644 --- a/src/OutputWriters/checkpointer.jl +++ b/src/OutputWriters/checkpointer.jl @@ -179,7 +179,7 @@ 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..." From 666dfcc9ad55c3072418dcc494bff089bdd14f5b Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 26 Jun 2025 10:07:32 +1000 Subject: [PATCH 21/33] Update checkpointer.jl --- src/OutputWriters/checkpointer.jl | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/OutputWriters/checkpointer.jl b/src/OutputWriters/checkpointer.jl index 911200f213..89001c6524 100644 --- a/src/OutputWriters/checkpointer.jl +++ b/src/OutputWriters/checkpointer.jl @@ -234,12 +234,8 @@ end Set data in `model.velocities`, `model.tracers`, `model.timestepper.Gⁿ`, and `model.timestepper.G⁻` to checkpointed data stored at `filepath`. """ -<<<<<<< ncc/checkopointer-shenanigans-2 function set!(model::AbstractModel, filepath::AbstractString) -======= -function set!(model, filepath::AbstractString) ->>>>>>> main addr = checkpointer_address(model) jldopen(filepath, "r") do file From 97e3edead726ca4e629f247490f3a1ec3a808944 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 26 Jun 2025 10:08:50 +1000 Subject: [PATCH 22/33] Update simulation.jl --- src/Simulations/simulation.jl | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/src/Simulations/simulation.jl b/src/Simulations/simulation.jl index 3d6e5d92f6..33fbfcee85 100644 --- a/src/Simulations/simulation.jl +++ b/src/Simulations/simulation.jl @@ -28,23 +28,14 @@ mutable struct Simulation{ML, DT, ST, DI, OW, CB, FT, BL} end """ -<<<<<<< ncc/checkopointer-shenanigans-2 - Simulation(model; Δt, - verbose = true, - stop_iteration = Inf, - stop_time = Inf, - wall_time_limit = Inf, - align_time_step = true, - minimum_relative_step = 0) -======= Simulation(model; Δt, verbose = true, stop_iteration = Inf, stop_time = Inf, wall_time_limit = Inf, + align_time_step = true, minimum_relative_step = 0) ->>>>>>> main Construct a `Simulation` for a `model` with time step `Δt`. From 09920382aca09e14bf1599d6887318d82a673b04 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 26 Jun 2025 10:09:11 +1000 Subject: [PATCH 23/33] Update clock.jl --- src/TimeSteppers/clock.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/src/TimeSteppers/clock.jl b/src/TimeSteppers/clock.jl index 2c3b031903..0f9a0f831b 100644 --- a/src/TimeSteppers/clock.jl +++ b/src/TimeSteppers/clock.jl @@ -1,4 +1,3 @@ -<<<<<<< ncc/checkopointer-shenanigans-2 using Adapt using Dates: AbstractTime, DateTime, Nanosecond, Millisecond using Oceananigans.Utils: prettytime From a9e66572eb2bf9221c8496f4b74f4b147b3d8ae3 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 26 Jun 2025 10:10:11 +1000 Subject: [PATCH 24/33] Update clock.jl --- src/TimeSteppers/clock.jl | 125 -------------------------------------- 1 file changed, 125 deletions(-) diff --git a/src/TimeSteppers/clock.jl b/src/TimeSteppers/clock.jl index 0f9a0f831b..5520237962 100644 --- a/src/TimeSteppers/clock.jl +++ b/src/TimeSteppers/clock.jl @@ -146,128 +146,3 @@ Adapt.adapt_structure(to, clock::Clock) = (time = clock.time, last_stage_Δt = clock.last_stage_Δt, iteration = clock.iteration, stage = clock.stage) -======= -using Adapt -using Dates: AbstractTime, DateTime, Nanosecond, Millisecond -using Oceananigans.Utils: prettytime -using Oceananigans.Grids: AbstractGrid - -import Base: show -import Oceananigans.Units: Time - -""" - mutable struct Clock{T, FT} - -Keeps track of the current `time`, `last_Δt`, `iteration` number, and time-stepping `stage`. -The `stage` is updated only for multi-stage time-stepping methods. The `time::T` is -either a number or a `DateTime` object. -""" -mutable struct Clock{TT, DT, IT, S} - time :: TT - last_Δt :: DT - last_stage_Δt :: DT - iteration :: IT - stage :: S -end - -function reset!(clock::Clock{TT, DT, IT, S}) where {TT, DT, IT, S} - clock.time = zero(TT) - clock.iteration = zero(IT) - clock.stage = zero(S) - clock.last_Δt = Inf - clock.last_stage_Δt = Inf - return nothing -end - -""" - Clock(; time, last_Δt=Inf, last_stage_Δt=Inf, iteration=0, stage=1) - -Returns a `Clock` object. By default, `Clock` is initialized to the zeroth `iteration` -and first time step `stage` with `last_Δt=last_stage_Δt=Inf`. -""" -function Clock(; time, - last_Δt = Inf, - last_stage_Δt = Inf, - iteration = 0, - stage = 1) - - TT = typeof(time) - DT = typeof(last_Δt) - IT = typeof(iteration) - last_stage_Δt = convert(DT, last_Δt) - return Clock{TT, DT, IT, typeof(stage)}(time, last_Δt, last_stage_Δt, iteration, stage) -end - -# TODO: when supporting DateTime, this function will have to be extended -time_step_type(TT) = TT - -function Clock{TT}(; time, - last_Δt = Inf, - last_stage_Δt = Inf, - iteration = 0, - stage = 1) where TT - - DT = time_step_type(TT) - last_Δt = convert(DT, last_Δt) - last_stage_Δt = convert(DT, last_stage_Δt) - IT = typeof(iteration) - - return Clock{TT, DT, IT, typeof(stage)}(time, last_Δt, last_stage_Δt, iteration, stage) -end - -# helpful default -Clock(grid::AbstractGrid) = Clock{Float64}(time=0) - -function Base.summary(clock::Clock) - TT = typeof(clock.time) - DT = typeof(clock.last_Δt) - return string("Clock{", TT, ", ", DT, "}", - "(time=", prettytime(clock.time), - ", iteration=", clock.iteration, - ", last_Δt=", prettytime(clock.last_Δt), ")") -end - -function Base.show(io::IO, clock::Clock) - return print(io, summary(clock), '\n', - "├── stage: ", clock.stage, '\n', - "└── last_stage_Δt: ", prettytime(clock.last_stage_Δt)) -end - -next_time(clock, Δt) = clock.time + Δt -next_time(clock::Clock{<:AbstractTime}, Δt) = clock.time + Nanosecond(round(Int, 1e9 * Δt)) - -tick_time!(clock, Δt) = clock.time += Δt -tick_time!(clock::Clock{<:AbstractTime}, Δt) = clock.time += Nanosecond(round(Int, 1e9 * Δt)) - -Time(clock::Clock) = Time(clock.time) - -# Convert the time to units of clock.time (assumed to be seconds if using DateTime or TimeDate). -unit_time(t) = t -unit_time(t::Millisecond) = t.value / 1_000 -unit_time(t::Nanosecond) = t.value / 1_000_000_000 - -# Convert to a base Julia type (a float or DateTime). Mainly used by NetCDFWriter. -float_or_date_time(t) = t -float_or_date_time(t::AbstractTime) = DateTime(t) - -function tick!(clock, Δt; stage=false) - - tick_time!(clock, Δt) - - if stage # tick a stage update - clock.stage += 1 - else # tick an iteration and reset stage - clock.iteration += 1 - clock.stage = 1 - end - - return nothing -end - -"""Adapt `Clock` for GPU.""" -Adapt.adapt_structure(to, clock::Clock) = (time = clock.time, - last_Δt = clock.last_Δt, - last_stage_Δt = clock.last_stage_Δt, - iteration = clock.iteration, - stage = clock.stage) ->>>>>>> main From e4e07e08c5b4ff7b853a6c7912bc2bdf17cbcf17 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 26 Jun 2025 10:12:24 +1000 Subject: [PATCH 25/33] clock from main --- src/TimeSteppers/clock.jl | 25 ------------------------- 1 file changed, 25 deletions(-) diff --git a/src/TimeSteppers/clock.jl b/src/TimeSteppers/clock.jl index 5520237962..d8752975d4 100644 --- a/src/TimeSteppers/clock.jl +++ b/src/TimeSteppers/clock.jl @@ -30,29 +30,6 @@ function reset!(clock::Clock{TT, DT, IT, S}) where {TT, DT, IT, S} return nothing end -""" - set_clock!(clock1::Clock, clock2::Clock) - -Set `clock1` to the `clock2`. -""" -function set_clock!(clock1::Clock, clock2::Clock) - clock1.time = clock2.time - clock1.iteration = clock2.iteration - clock1.last_Δt = clock2.last_Δt - clock1.last_stage_Δt = clock2.last_stage_Δt - clock1.stage = clock2.stage - - return nothing -end - -function Base.:(==)(clock1::Clock, clock2::Clock) - return clock1.time == clock2.time && - clock1.iteration == clock2.iteration && - clock1.last_Δt == clock2.last_Δt && - clock1.last_stage_Δt == clock2.last_stage_Δt && - clock1.stage == clock2.stage -end - """ Clock(; time, last_Δt=Inf, last_stage_Δt=Inf, iteration=0, stage=1) @@ -130,11 +107,9 @@ function tick!(clock, Δt; stage=false) if stage # tick a stage update clock.stage += 1 - clock.last_stage_Δt = Δt else # tick an iteration and reset stage clock.iteration += 1 clock.stage = 1 - clock.last_Δt = Δt end return nothing From 3da6865620aee4f72eefbaab0e2f7231e629787a Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 26 Jun 2025 10:12:45 +1000 Subject: [PATCH 26/33] updates in clock --- src/TimeSteppers/clock.jl | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/src/TimeSteppers/clock.jl b/src/TimeSteppers/clock.jl index d8752975d4..5520237962 100644 --- a/src/TimeSteppers/clock.jl +++ b/src/TimeSteppers/clock.jl @@ -30,6 +30,29 @@ function reset!(clock::Clock{TT, DT, IT, S}) where {TT, DT, IT, S} return nothing end +""" + set_clock!(clock1::Clock, clock2::Clock) + +Set `clock1` to the `clock2`. +""" +function set_clock!(clock1::Clock, clock2::Clock) + clock1.time = clock2.time + clock1.iteration = clock2.iteration + clock1.last_Δt = clock2.last_Δt + clock1.last_stage_Δt = clock2.last_stage_Δt + clock1.stage = clock2.stage + + return nothing +end + +function Base.:(==)(clock1::Clock, clock2::Clock) + return clock1.time == clock2.time && + clock1.iteration == clock2.iteration && + clock1.last_Δt == clock2.last_Δt && + clock1.last_stage_Δt == clock2.last_stage_Δt && + clock1.stage == clock2.stage +end + """ Clock(; time, last_Δt=Inf, last_stage_Δt=Inf, iteration=0, stage=1) @@ -107,9 +130,11 @@ function tick!(clock, Δt; stage=false) if stage # tick a stage update clock.stage += 1 + clock.last_stage_Δt = Δt else # tick an iteration and reset stage clock.iteration += 1 clock.stage = 1 + clock.last_Δt = Δt end return nothing From b91aa9bef73ce8e8d3d3402112d6ac2045677b82 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 26 Jun 2025 10:13:56 +1000 Subject: [PATCH 27/33] clock from main --- src/TimeSteppers/clock.jl | 25 ------------------------- 1 file changed, 25 deletions(-) diff --git a/src/TimeSteppers/clock.jl b/src/TimeSteppers/clock.jl index 5520237962..d8752975d4 100644 --- a/src/TimeSteppers/clock.jl +++ b/src/TimeSteppers/clock.jl @@ -30,29 +30,6 @@ function reset!(clock::Clock{TT, DT, IT, S}) where {TT, DT, IT, S} return nothing end -""" - set_clock!(clock1::Clock, clock2::Clock) - -Set `clock1` to the `clock2`. -""" -function set_clock!(clock1::Clock, clock2::Clock) - clock1.time = clock2.time - clock1.iteration = clock2.iteration - clock1.last_Δt = clock2.last_Δt - clock1.last_stage_Δt = clock2.last_stage_Δt - clock1.stage = clock2.stage - - return nothing -end - -function Base.:(==)(clock1::Clock, clock2::Clock) - return clock1.time == clock2.time && - clock1.iteration == clock2.iteration && - clock1.last_Δt == clock2.last_Δt && - clock1.last_stage_Δt == clock2.last_stage_Δt && - clock1.stage == clock2.stage -end - """ Clock(; time, last_Δt=Inf, last_stage_Δt=Inf, iteration=0, stage=1) @@ -130,11 +107,9 @@ function tick!(clock, Δt; stage=false) if stage # tick a stage update clock.stage += 1 - clock.last_stage_Δt = Δt else # tick an iteration and reset stage clock.iteration += 1 clock.stage = 1 - clock.last_Δt = Δt end return nothing From 0ac44271629564bd323f065aefb699bc2e07cf62 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 26 Jun 2025 10:18:14 +1000 Subject: [PATCH 28/33] updates in clock --- src/TimeSteppers/clock.jl | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/src/TimeSteppers/clock.jl b/src/TimeSteppers/clock.jl index d8752975d4..5520237962 100644 --- a/src/TimeSteppers/clock.jl +++ b/src/TimeSteppers/clock.jl @@ -30,6 +30,29 @@ function reset!(clock::Clock{TT, DT, IT, S}) where {TT, DT, IT, S} return nothing end +""" + set_clock!(clock1::Clock, clock2::Clock) + +Set `clock1` to the `clock2`. +""" +function set_clock!(clock1::Clock, clock2::Clock) + clock1.time = clock2.time + clock1.iteration = clock2.iteration + clock1.last_Δt = clock2.last_Δt + clock1.last_stage_Δt = clock2.last_stage_Δt + clock1.stage = clock2.stage + + return nothing +end + +function Base.:(==)(clock1::Clock, clock2::Clock) + return clock1.time == clock2.time && + clock1.iteration == clock2.iteration && + clock1.last_Δt == clock2.last_Δt && + clock1.last_stage_Δt == clock2.last_stage_Δt && + clock1.stage == clock2.stage +end + """ Clock(; time, last_Δt=Inf, last_stage_Δt=Inf, iteration=0, stage=1) @@ -107,9 +130,11 @@ function tick!(clock, Δt; stage=false) if stage # tick a stage update clock.stage += 1 + clock.last_stage_Δt = Δt else # tick an iteration and reset stage clock.iteration += 1 clock.stage = 1 + clock.last_Δt = Δt end return nothing From cb01cf49ff566fb10c52c25559dec30ac42a6650 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 26 Jun 2025 10:21:20 +1000 Subject: [PATCH 29/33] updates in clock --- src/TimeSteppers/clock.jl | 38 +++++++++++++++++++------------------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/src/TimeSteppers/clock.jl b/src/TimeSteppers/clock.jl index 5520237962..d4a6553c0f 100644 --- a/src/TimeSteppers/clock.jl +++ b/src/TimeSteppers/clock.jl @@ -21,6 +21,25 @@ mutable struct Clock{TT, DT, IT, S} stage :: S end +""" + Clock(; time, last_Δt=Inf, last_stage_Δt=Inf, iteration=0, stage=1) + +Return a `Clock` object. By default, `Clock` is initialized to the zeroth `iteration` +and first time step `stage` with `last_Δt=last_stage_Δt=Inf`. +""" +function Clock(; time, + last_Δt = Inf, + last_stage_Δt = Inf, + iteration = 0, + stage = 1) + + TT = typeof(time) + DT = typeof(last_Δt) + IT = typeof(iteration) + last_stage_Δt = convert(DT, last_Δt) + return Clock{TT, DT, IT, typeof(stage)}(time, last_Δt, last_stage_Δt, iteration, stage) +end + function reset!(clock::Clock{TT, DT, IT, S}) where {TT, DT, IT, S} clock.time = zero(TT) clock.iteration = zero(IT) @@ -53,25 +72,6 @@ function Base.:(==)(clock1::Clock, clock2::Clock) clock1.stage == clock2.stage end -""" - Clock(; time, last_Δt=Inf, last_stage_Δt=Inf, iteration=0, stage=1) - -Returns a `Clock` object. By default, `Clock` is initialized to the zeroth `iteration` -and first time step `stage` with `last_Δt=last_stage_Δt=Inf`. -""" -function Clock(; time, - last_Δt = Inf, - last_stage_Δt = Inf, - iteration = 0, - stage = 1) - - TT = typeof(time) - DT = typeof(last_Δt) - IT = typeof(iteration) - last_stage_Δt = convert(DT, last_Δt) - return Clock{TT, DT, IT, typeof(stage)}(time, last_Δt, last_stage_Δt, iteration, stage) -end - # TODO: when supporting DateTime, this function will have to be extended time_step_type(TT) = TT From 2ac952f816d9f030a16496b5cd26e66459d3e252 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 26 Jun 2025 10:23:38 +1000 Subject: [PATCH 30/33] clock from main --- src/TimeSteppers/clock.jl | 271 +++++++++++++++++--------------------- 1 file changed, 123 insertions(+), 148 deletions(-) diff --git a/src/TimeSteppers/clock.jl b/src/TimeSteppers/clock.jl index d4a6553c0f..24c4a2364f 100644 --- a/src/TimeSteppers/clock.jl +++ b/src/TimeSteppers/clock.jl @@ -1,148 +1,123 @@ -using Adapt -using Dates: AbstractTime, DateTime, Nanosecond, Millisecond -using Oceananigans.Utils: prettytime -using Oceananigans.Grids: AbstractGrid - -import Base: show -import Oceananigans.Units: Time - -""" - mutable struct Clock{T, FT} - -Keeps track of the current `time`, `last_Δt`, `iteration` number, and time-stepping `stage`. -The `stage` is updated only for multi-stage time-stepping methods. The `time::T` is -either a number or a `DateTime` object. -""" -mutable struct Clock{TT, DT, IT, S} - time :: TT - last_Δt :: DT - last_stage_Δt :: DT - iteration :: IT - stage :: S -end - -""" - Clock(; time, last_Δt=Inf, last_stage_Δt=Inf, iteration=0, stage=1) - -Return a `Clock` object. By default, `Clock` is initialized to the zeroth `iteration` -and first time step `stage` with `last_Δt=last_stage_Δt=Inf`. -""" -function Clock(; time, - last_Δt = Inf, - last_stage_Δt = Inf, - iteration = 0, - stage = 1) - - TT = typeof(time) - DT = typeof(last_Δt) - IT = typeof(iteration) - last_stage_Δt = convert(DT, last_Δt) - return Clock{TT, DT, IT, typeof(stage)}(time, last_Δt, last_stage_Δt, iteration, stage) -end - -function reset!(clock::Clock{TT, DT, IT, S}) where {TT, DT, IT, S} - clock.time = zero(TT) - clock.iteration = zero(IT) - clock.stage = zero(S) - clock.last_Δt = Inf - clock.last_stage_Δt = Inf - return nothing -end - -""" - set_clock!(clock1::Clock, clock2::Clock) - -Set `clock1` to the `clock2`. -""" -function set_clock!(clock1::Clock, clock2::Clock) - clock1.time = clock2.time - clock1.iteration = clock2.iteration - clock1.last_Δt = clock2.last_Δt - clock1.last_stage_Δt = clock2.last_stage_Δt - clock1.stage = clock2.stage - - return nothing -end - -function Base.:(==)(clock1::Clock, clock2::Clock) - return clock1.time == clock2.time && - clock1.iteration == clock2.iteration && - clock1.last_Δt == clock2.last_Δt && - clock1.last_stage_Δt == clock2.last_stage_Δt && - clock1.stage == clock2.stage -end - -# TODO: when supporting DateTime, this function will have to be extended -time_step_type(TT) = TT - -function Clock{TT}(; time, - last_Δt = Inf, - last_stage_Δt = Inf, - iteration = 0, - stage = 1) where TT - - DT = time_step_type(TT) - last_Δt = convert(DT, last_Δt) - last_stage_Δt = convert(DT, last_stage_Δt) - IT = typeof(iteration) - - return Clock{TT, DT, IT, typeof(stage)}(time, last_Δt, last_stage_Δt, iteration, stage) -end - -# helpful default -Clock(grid::AbstractGrid) = Clock{Float64}(time=0) - -function Base.summary(clock::Clock) - TT = typeof(clock.time) - DT = typeof(clock.last_Δt) - return string("Clock{", TT, ", ", DT, "}", - "(time=", prettytime(clock.time), - ", iteration=", clock.iteration, - ", last_Δt=", prettytime(clock.last_Δt), ")") -end - -function Base.show(io::IO, clock::Clock) - return print(io, summary(clock), '\n', - "├── stage: ", clock.stage, '\n', - "└── last_stage_Δt: ", prettytime(clock.last_stage_Δt)) -end - -next_time(clock, Δt) = clock.time + Δt -next_time(clock::Clock{<:AbstractTime}, Δt) = clock.time + Nanosecond(round(Int, 1e9 * Δt)) - -tick_time!(clock, Δt) = clock.time += Δt -tick_time!(clock::Clock{<:AbstractTime}, Δt) = clock.time += Nanosecond(round(Int, 1e9 * Δt)) - -Time(clock::Clock) = Time(clock.time) - -# Convert the time to units of clock.time (assumed to be seconds if using DateTime or TimeDate). -unit_time(t) = t -unit_time(t::Millisecond) = t.value / 1_000 -unit_time(t::Nanosecond) = t.value / 1_000_000_000 - -# Convert to a base Julia type (a float or DateTime). Mainly used by NetCDFWriter. -float_or_date_time(t) = t -float_or_date_time(t::AbstractTime) = DateTime(t) - -function tick!(clock, Δt; stage=false) - - tick_time!(clock, Δt) - - if stage # tick a stage update - clock.stage += 1 - clock.last_stage_Δt = Δt - else # tick an iteration and reset stage - clock.iteration += 1 - clock.stage = 1 - clock.last_Δt = Δt - end - - return nothing -end - -"""Adapt `Clock` for GPU.""" -Adapt.adapt_structure(to, clock::Clock) = (time = clock.time, - last_Δt = clock.last_Δt, - last_stage_Δt = clock.last_stage_Δt, - iteration = clock.iteration, - stage = clock.stage) +using Adapt +using Dates: AbstractTime, DateTime, Nanosecond, Millisecond +using Oceananigans.Utils: prettytime +using Oceananigans.Grids: AbstractGrid + +import Base: show +import Oceananigans.Units: Time + +""" + mutable struct Clock{T, FT} + +Keeps track of the current `time`, `last_Δt`, `iteration` number, and time-stepping `stage`. +The `stage` is updated only for multi-stage time-stepping methods. The `time::T` is +either a number or a `DateTime` object. +""" +mutable struct Clock{TT, DT, IT, S} + time :: TT + last_Δt :: DT + last_stage_Δt :: DT + iteration :: IT + stage :: S +end + +function reset!(clock::Clock{TT, DT, IT, S}) where {TT, DT, IT, S} + clock.time = zero(TT) + clock.iteration = zero(IT) + clock.stage = zero(S) + clock.last_Δt = Inf + clock.last_stage_Δt = Inf + return nothing +end + +""" + Clock(; time, last_Δt=Inf, last_stage_Δt=Inf, iteration=0, stage=1) + +Returns a `Clock` object. By default, `Clock` is initialized to the zeroth `iteration` +and first time step `stage` with `last_Δt=last_stage_Δt=Inf`. +""" +function Clock(; time, + last_Δt = Inf, + last_stage_Δt = Inf, + iteration = 0, + stage = 1) + + TT = typeof(time) + DT = typeof(last_Δt) + IT = typeof(iteration) + last_stage_Δt = convert(DT, last_Δt) + return Clock{TT, DT, IT, typeof(stage)}(time, last_Δt, last_stage_Δt, iteration, stage) +end + +# TODO: when supporting DateTime, this function will have to be extended +time_step_type(TT) = TT + +function Clock{TT}(; time, + last_Δt = Inf, + last_stage_Δt = Inf, + iteration = 0, + stage = 1) where TT + + DT = time_step_type(TT) + last_Δt = convert(DT, last_Δt) + last_stage_Δt = convert(DT, last_stage_Δt) + IT = typeof(iteration) + + return Clock{TT, DT, IT, typeof(stage)}(time, last_Δt, last_stage_Δt, iteration, stage) +end + +# helpful default +Clock(grid::AbstractGrid) = Clock{Float64}(time=0) + +function Base.summary(clock::Clock) + TT = typeof(clock.time) + DT = typeof(clock.last_Δt) + return string("Clock{", TT, ", ", DT, "}", + "(time=", prettytime(clock.time), + ", iteration=", clock.iteration, + ", last_Δt=", prettytime(clock.last_Δt), ")") +end + +function Base.show(io::IO, clock::Clock) + return print(io, summary(clock), '\n', + "├── stage: ", clock.stage, '\n', + "└── last_stage_Δt: ", prettytime(clock.last_stage_Δt)) +end + +next_time(clock, Δt) = clock.time + Δt +next_time(clock::Clock{<:AbstractTime}, Δt) = clock.time + Nanosecond(round(Int, 1e9 * Δt)) + +tick_time!(clock, Δt) = clock.time += Δt +tick_time!(clock::Clock{<:AbstractTime}, Δt) = clock.time += Nanosecond(round(Int, 1e9 * Δt)) + +Time(clock::Clock) = Time(clock.time) + +# Convert the time to units of clock.time (assumed to be seconds if using DateTime or TimeDate). +unit_time(t) = t +unit_time(t::Millisecond) = t.value / 1_000 +unit_time(t::Nanosecond) = t.value / 1_000_000_000 + +# Convert to a base Julia type (a float or DateTime). Mainly used by NetCDFWriter. +float_or_date_time(t) = t +float_or_date_time(t::AbstractTime) = DateTime(t) + +function tick!(clock, Δt; stage=false) + + tick_time!(clock, Δt) + + if stage # tick a stage update + clock.stage += 1 + else # tick an iteration and reset stage + clock.iteration += 1 + clock.stage = 1 + end + + return nothing +end + +"""Adapt `Clock` for GPU.""" +Adapt.adapt_structure(to, clock::Clock) = (time = clock.time, + last_Δt = clock.last_Δt, + last_stage_Δt = clock.last_stage_Δt, + iteration = clock.iteration, + stage = clock.stage) From a673011853aa12b01a24276d417859d5cc1633c6 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 26 Jun 2025 10:24:59 +1000 Subject: [PATCH 31/33] Update clock.jl --- src/TimeSteppers/clock.jl | 45 ++++++++++++++++++++++++++++++--------- 1 file changed, 35 insertions(+), 10 deletions(-) diff --git a/src/TimeSteppers/clock.jl b/src/TimeSteppers/clock.jl index 24c4a2364f..8bba049290 100644 --- a/src/TimeSteppers/clock.jl +++ b/src/TimeSteppers/clock.jl @@ -21,19 +21,10 @@ mutable struct Clock{TT, DT, IT, S} stage :: S end -function reset!(clock::Clock{TT, DT, IT, S}) where {TT, DT, IT, S} - clock.time = zero(TT) - clock.iteration = zero(IT) - clock.stage = zero(S) - clock.last_Δt = Inf - clock.last_stage_Δt = Inf - return nothing -end - """ Clock(; time, last_Δt=Inf, last_stage_Δt=Inf, iteration=0, stage=1) -Returns a `Clock` object. By default, `Clock` is initialized to the zeroth `iteration` +Return a `Clock` object. By default, `Clock` is initialized to the zeroth `iteration` and first time step `stage` with `last_Δt=last_stage_Δt=Inf`. """ function Clock(; time, @@ -49,6 +40,38 @@ function Clock(; time, return Clock{TT, DT, IT, typeof(stage)}(time, last_Δt, last_stage_Δt, iteration, stage) end +function reset!(clock::Clock{TT, DT, IT, S}) where {TT, DT, IT, S} + clock.time = zero(TT) + clock.iteration = zero(IT) + clock.stage = zero(S) + clock.last_Δt = Inf + clock.last_stage_Δt = Inf + return nothing +end + +""" + set_clock!(clock1::Clock, clock2::Clock) + +Set `clock1` to the `clock2`. +""" +function set_clock!(clock1::Clock, clock2::Clock) + clock1.time = clock2.time + clock1.iteration = clock2.iteration + clock1.last_Δt = clock2.last_Δt + clock1.last_stage_Δt = clock2.last_stage_Δt + clock1.stage = clock2.stage + + return nothing +end + +function Base.:(==)(clock1::Clock, clock2::Clock) + return clock1.time == clock2.time && + clock1.iteration == clock2.iteration && + clock1.last_Δt == clock2.last_Δt && + clock1.last_stage_Δt == clock2.last_stage_Δt && + clock1.stage == clock2.stage +end + # TODO: when supporting DateTime, this function will have to be extended time_step_type(TT) = TT @@ -107,9 +130,11 @@ function tick!(clock, Δt; stage=false) if stage # tick a stage update clock.stage += 1 + clock.last_stage_Δt = Δt else # tick an iteration and reset stage clock.iteration += 1 clock.stage = 1 + clock.last_Δt = Δt end return nothing From 28f6e39f55a9e87eca468dfbea300a29d000b786 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Fri, 27 Jun 2025 12:20:29 +1000 Subject: [PATCH 32/33] import AbstractModel --- src/OutputWriters/checkpointer.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/OutputWriters/checkpointer.jl b/src/OutputWriters/checkpointer.jl index 89001c6524..abf67eddb7 100644 --- a/src/OutputWriters/checkpointer.jl +++ b/src/OutputWriters/checkpointer.jl @@ -1,7 +1,7 @@ using Glob using Oceananigans -using Oceananigans: fields, prognostic_fields +using Oceananigans: AbstractModel, fields, prognostic_fields using Oceananigans.Fields: offset_data using Oceananigans.TimeSteppers: QuasiAdamsBashforth2TimeStepper @@ -227,7 +227,7 @@ end ##### set! for checkpointer filepaths ##### -# Should this go in Models? +# Should this go in Models? """ set!(model::AbstractModel, filepath::AbstractString) @@ -278,7 +278,7 @@ end function set_time_stepper_tendencies!(timestepper, file, model_fields, addr) for name in propertynames(model_fields) - tendency_in_model = hasproperty(timestepper.Gⁿ, name) + tendency_in_model = hasproperty(timestepper.Gⁿ, name) tendency_in_checkpoint = string(name) ∈ keys(file["$addr/timestepper/Gⁿ"]) if tendency_in_model && tendency_in_checkpoint # Tendency "n" From db853557a385bcde1f5a80c5867e99dba140efbe Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Wed, 9 Jul 2025 16:11:23 +1000 Subject: [PATCH 33/33] Update runge_kutta_3.jl --- src/TimeSteppers/runge_kutta_3.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/TimeSteppers/runge_kutta_3.jl b/src/TimeSteppers/runge_kutta_3.jl index 8513caae86..318a9759f3 100644 --- a/src/TimeSteppers/runge_kutta_3.jl +++ b/src/TimeSteppers/runge_kutta_3.jl @@ -151,6 +151,7 @@ function time_step!(model::AbstractModel{<:RungeKutta3TimeStepper}, Δt; callbac # now model.clock.last_Δt = clock.last_stage_Δt = third_stage_Δt # we correct those below model.clock.last_stage_Δt = corrected_third_stage_Δt + model.clock.last_Δt = Δt compute_pressure_correction!(model, third_stage_Δt) make_pressure_correction!(model, third_stage_Δt)