From 806be1eb36d21dc7784895fcafbbe4a4238dde93 Mon Sep 17 00:00:00 2001 From: Taimoor Sohail Date: Tue, 11 Mar 2025 17:56:53 +1100 Subject: [PATCH 01/54] Added a run_coupled! function to implement checkpointers in coupled simulations --- examples/checkpointer_mwe.jl | 107 +++++++++++++++++++++++++++++++++++ 1 file changed, 107 insertions(+) create mode 100644 examples/checkpointer_mwe.jl diff --git a/examples/checkpointer_mwe.jl b/examples/checkpointer_mwe.jl new file mode 100644 index 000000000..7613d63e1 --- /dev/null +++ b/examples/checkpointer_mwe.jl @@ -0,0 +1,107 @@ +using ClimaOcean +using Oceananigans +using Oceananigans.Units +using CFTime +using Dates +using Printf + +arch = CPU() + +Nx = 144 +Ny = 60 +Nz = 40 + +depth = 6000meters +z_faces = exponential_z_faces(; Nz, depth) + +grid = LatitudeLongitudeGrid(arch; + size = (Nx, Ny, Nz), + halo = (7, 7, 7), + z = z_faces, + latitude = (-75, 75), + longitude = (0, 360)) + +ocean = ocean_simulation(grid) + +# date = DateTimeProlepticGregorian(1993, 1, 1) +# set!(ocean.model, T=ECCOMetadata(:temperature; dates=date), +# S=ECCOMetadata(:salinity; dates=date)) + +radiation = Radiation(arch) + +atmosphere = JRA55PrescribedAtmosphere(arch; backend=JRA55NetCDFBackend(41)) + +coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation) + +simulation = Simulation(coupled_model; Δt=10, stop_iteration=10) + +wall_time = Ref(time_ns()) + +function progress(sim) + ocean = sim.model.ocean + u, v, w = ocean.model.velocities + T = ocean.model.tracers.T + + Tmax = maximum(interior(T)) + Tmin = minimum(interior(T)) + + umax = (maximum(abs, interior(u)), + maximum(abs, interior(v)), + maximum(abs, interior(w))) + + step_time = 1e-9 * (time_ns() - wall_time[]) + + msg = @sprintf("Iter: %d, simulation time: %s, atmosphere time: %s, Δt: %s", iteration(sim), prettytime(sim), prettytime(atmosphere.clock.time), prettytime(sim.Δt)) + msg *= @sprintf(", max|u|: (%.2e, %.2e, %.2e) m s⁻¹, extrema(T): (%.2f, %.2f) ᵒC, wall time: %s", + umax..., Tmax, Tmin, prettytime(step_time)) + + @info msg + + wall_time[] = time_ns() +end + +simulation.callbacks[:progress] = Callback(progress, IterationInterval(1)) + +outputs = merge(ocean.model.tracers, ocean.model.velocities) + +ocean.output_writers[:surface] = JLD2OutputWriter(ocean.model, outputs; + schedule = IterationInterval(2), + filename = "checkpointer_mwe_surface", + indices = (:, :, grid.Nz), + with_halos = true, + overwrite_existing = true, + array_type = Array{Float32}) + +output_dir = "." +prefix = "checkpointer_mwe" + +ocean.output_writers[:checkpoint] = Checkpointer(ocean.model; + schedule = IterationInterval(2), + prefix = prefix, + dir = output_dir, + verbose = true, + overwrite_existing = true) + +# coupled_checkpointer = Checkpointer(coupled_model; +# schedule = IterationInterval(4), +# prefix = prefix, +# dir = output_dir, +# verbose = true, +# overwrite_existing = true) + +# @show simulation + +run_coupled!(simulation, pickup=true) + +# @info "simulation run for 10 iterations; you should have a checkpointer at 8" + +# checkpoint_file = prefix * "_iteration0.jld2" + +# set!(simulation, checkpoint_file) + +# coupled_model = OceanSeaIceModel(simulation.model.ocean; atmosphere, radiation) + +# simulation = Simulation(coupled_model; Δt=10, stop_iteration=20) +# simulation.callbacks[:progress] = Callback(progress, IterationInterval(1)) + +# run!(simulation) \ No newline at end of file From 46b7bb283373a5d3fc02a67f6afaa01a9fa82080 Mon Sep 17 00:00:00 2001 From: Taimoor Sohail Date: Tue, 11 Mar 2025 17:58:39 +1100 Subject: [PATCH 02/54] Added a run_coupled! function to implement checkpointers in coupled simulations --- src/ClimaOcean.jl | 1 + src/OceanSeaIceModels/ocean_sea_ice_model.jl | 30 ++++++++++++++++++-- 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/src/ClimaOcean.jl b/src/ClimaOcean.jl index 2205ede93..48b8fa056 100644 --- a/src/ClimaOcean.jl +++ b/src/ClimaOcean.jl @@ -32,6 +32,7 @@ export ocean_simulation, sea_ice_simulation, initialize!, + run_coupled!, @root, @onrank, @distribute, diff --git a/src/OceanSeaIceModels/ocean_sea_ice_model.jl b/src/OceanSeaIceModels/ocean_sea_ice_model.jl index 2261457eb..58b8bf6a0 100644 --- a/src/OceanSeaIceModels/ocean_sea_ice_model.jl +++ b/src/OceanSeaIceModels/ocean_sea_ice_model.jl @@ -1,9 +1,11 @@ +export run_coupled! + using Oceananigans +using Oceananigans.OutputWriters: checkpoint_path using Oceananigans.TimeSteppers: Clock using Oceananigans: SeawaterBuoyancy using ClimaSeaIce.SeaIceThermodynamics: melting_temperature using KernelAbstractions: @kernel, @index - using SeawaterPolynomials: TEOS10EquationOfState import Thermodynamics as AtmosphericThermodynamics @@ -14,7 +16,7 @@ import Oceananigans.Architectures: architecture import Oceananigans.Fields: set! import Oceananigans.Models: timestepper, NaNChecker, default_nan_checker import Oceananigans.OutputWriters: default_included_properties -import Oceananigans.Simulations: reset!, initialize!, iteration +import Oceananigans.Simulations: reset!, initialize!, iteration, run! import Oceananigans.TimeSteppers: time_step!, update_state!, time import Oceananigans.Utils: prettytime import Oceananigans.Models: timestepper, NaNChecker, default_nan_checker @@ -29,6 +31,20 @@ struct OceanSeaIceModel{I, A, O, F, C, Arch} <: AbstractModel{Nothing, Arch} end const OSIM = OceanSeaIceModel +const OSIMSIM = Simulation{<:OceanSeaIceModel} + +function run_coupled!(sim::OSIMSIM; pickup) + if pickup + checkpoint_file_path = checkpoint_path(pickup, sim.model.ocean.output_writers) + set!(sim.model.ocean, checkpoint_file_path) + # Setting the atmosphere time to the ocean time + sim.model.atmosphere.clock.time = sim.model.ocean.model.clock.time + end + + @show sim.model.clock.time + run!(sim) +end + function Base.summary(model::OSIM) A = nameof(typeof(architecture(model))) @@ -160,6 +176,16 @@ function default_nan_checker(model::OceanSeaIceModel) return nan_checker end +# TODO: picking up OceanSeaIceModel simulations from a checkpoint is a WIP +function set!(sim::OSIMSIM, pickup::Union{Bool, Integer, String}) + checkpoint_file_path = checkpoint_path(pickup, sim.model.ocean.output_writers) + set!(sim.model.ocean, filepath = checkpoint_file_path) + set!(sim.model.ocean, pickup = pickup) + # Setting the atmosphere time to the ocean time + sim.model.atmosphere.clock.time = sim.model.ocean.model.clock.time + return nothing +end + @kernel function _above_freezing_ocean_temperature!(T, grid, S, ℵ, liquidus) i, j = @index(Global, NTuple) Nz = size(grid, 3) From 75b81821e1fe6099a8765175f84119e0fa34bad3 Mon Sep 17 00:00:00 2001 From: Taimoor Sohail Date: Wed, 12 Mar 2025 16:51:02 +1100 Subject: [PATCH 03/54] Simplified the near_global_ocean example to see why the checkpointing isn't working --- examples/near_global_ocean_simulation.jl | 64 ++++++++++++++---------- 1 file changed, 38 insertions(+), 26 deletions(-) diff --git a/examples/near_global_ocean_simulation.jl b/examples/near_global_ocean_simulation.jl index bd525f227..c23394c01 100644 --- a/examples/near_global_ocean_simulation.jl +++ b/examples/near_global_ocean_simulation.jl @@ -29,10 +29,10 @@ using Printf # The total depth of the domain is set to 6000 meters. # Finally, we specify the architecture for the simulation, which in this case is a GPU. -arch = GPU() +arch = CPU() -Nx = 1440 -Ny = 600 +Nx = 144 +Ny = 60 Nz = 40 depth = 6000meters @@ -52,21 +52,21 @@ grid = LatitudeLongitudeGrid(arch; # (i) all the minor enclosed basins except the 3 largest `major_basins`, as well as # (ii) regions that are shallower than `minimum_depth`. -bottom_height = regrid_bathymetry(grid; - minimum_depth = 10meters, - interpolation_passes = 5, - major_basins = 3) +# bottom_height = regrid_bathymetry(grid; +# minimum_depth = 10meters, +# interpolation_passes = 5, +# major_basins = 3) -grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map=true) +# grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map=true) # Let's see what the bathymetry looks like: -h = grid.immersed_boundary.bottom_height +# h = grid.immersed_boundary.bottom_height -fig, ax, hm = heatmap(h, colormap=:deep, colorrange=(-depth, 0)) -Colorbar(fig[0, 1], hm, label="Bottom height (m)", vertical=false) -save("bathymetry.png", fig) -nothing #hide +# fig, ax, hm = heatmap(h, colormap=:deep, colorrange=(-depth, 0)) +# Colorbar(fig[0, 1], hm, label="Bottom height (m)", vertical=false) +# save("bathymetry.png", fig) +# nothing #hide # ![](bathymetry.png) @@ -82,9 +82,9 @@ ocean.model # We initialize the ocean model with ECCO2 temperature and salinity for January 1, 1993. -date = DateTimeProlepticGregorian(1993, 1, 1) -set!(ocean.model, T=ECCOMetadata(:temperature; dates=date), - S=ECCOMetadata(:salinity; dates=date)) +# date = DateTimeProlepticGregorian(1993, 1, 1) +# set!(ocean.model, T=ECCOMetadata(:temperature; dates=date), +# S=ECCOMetadata(:salinity; dates=date)) # ### Prescribed atmosphere and radiation # @@ -116,7 +116,7 @@ coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation) # We then create a coupled simulation. We start with a small-ish time step of 90 seconds. # We run the simulation for 10 days with this small-ish time step. -simulation = Simulation(coupled_model; Δt=90, stop_time=10days) +simulation = Simulation(coupled_model; Δt=90, stop_iteration=10) # We define a callback function to monitor the simulation's progress, @@ -145,7 +145,7 @@ function progress(sim) wall_time[] = time_ns() end -simulation.callbacks[:progress] = Callback(progress, TimeInterval(5days)) +simulation.callbacks[:progress] = Callback(progress, IterationInterval(1)) # ### Set up output writers # @@ -155,13 +155,23 @@ simulation.callbacks[:progress] = Callback(progress, TimeInterval(5days)) # Below, we use `indices` to save only the values of the variables at the surface, which corresponds to `k = grid.Nz` outputs = merge(ocean.model.tracers, ocean.model.velocities) -ocean.output_writers[:surface] = JLD2OutputWriter(ocean.model, outputs; - schedule = TimeInterval(1days), - filename = "near_global_surface_fields", - indices = (:, :, grid.Nz), - with_halos = true, - overwrite_existing = true, - array_type = Array{Float32}) +# ocean.output_writers[:surface] = JLD2OutputWriter(ocean.model, outputs; +# schedule = IterationInterval(1), +# filename = "near_global_surface_fields", +# indices = (:, :, grid.Nz), +# with_halos = true, +# overwrite_existing = true, +# array_type = Array{Float32}) + +output_dir = "." +prefix = "near_global" + +ocean.output_writers[:checkpoint] = Checkpointer(ocean.model; + schedule = TimeInterval(3minutes), + prefix = prefix, + dir = output_dir)#, + # verbose = true, + # overwrite_existing = true) # ### Spinning up the simulation # @@ -174,11 +184,12 @@ run!(simulation) # ### Running the simulation for real # After the initial spin up of 10 days, we can increase the time-step and run for longer. - +#= simulation.stop_time = 60days simulation.Δt = 10minutes run!(simulation) + # ## A pretty movie # # It's time to make a pretty movie of the simulation. First we load the output we've been saving on @@ -258,3 +269,4 @@ end nothing #hide # ![](near_global_ocean_surface.mp4) +=# \ No newline at end of file From fb5d35d4ea0315a192d7c02278326368797212bd Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 13 Mar 2025 16:39:28 +1100 Subject: [PATCH 04/54] Testing the set! function --- examples/checkpointer_mwe.jl | 48 +++++++------------- src/OceanSeaIceModels/ocean_sea_ice_model.jl | 17 +------ 2 files changed, 18 insertions(+), 47 deletions(-) diff --git a/examples/checkpointer_mwe.jl b/examples/checkpointer_mwe.jl index 7613d63e1..07f78f116 100644 --- a/examples/checkpointer_mwe.jl +++ b/examples/checkpointer_mwe.jl @@ -64,44 +64,30 @@ simulation.callbacks[:progress] = Callback(progress, IterationInterval(1)) outputs = merge(ocean.model.tracers, ocean.model.velocities) -ocean.output_writers[:surface] = JLD2OutputWriter(ocean.model, outputs; - schedule = IterationInterval(2), - filename = "checkpointer_mwe_surface", - indices = (:, :, grid.Nz), - with_halos = true, - overwrite_existing = true, - array_type = Array{Float32}) +simulation.output_writers[:surface] = JLD2OutputWriter(ocean.model, outputs; + schedule = IterationInterval(2), + filename = "checkpointer_mwe_surface", + indices = (:, :, grid.Nz), + with_halos = true, + overwrite_existing = true, + array_type = Array{Float32}) output_dir = "." prefix = "checkpointer_mwe" -ocean.output_writers[:checkpoint] = Checkpointer(ocean.model; - schedule = IterationInterval(2), - prefix = prefix, - dir = output_dir, - verbose = true, - overwrite_existing = true) - -# coupled_checkpointer = Checkpointer(coupled_model; -# schedule = IterationInterval(4), -# prefix = prefix, -# dir = output_dir, -# verbose = true, -# overwrite_existing = true) +simulation.output_writers[:checkpoint] = Checkpointer(ocean.model; + schedule = IterationInterval(2), + prefix = prefix, + dir = output_dir, + verbose = true, + overwrite_existing = true) # @show simulation -run_coupled!(simulation, pickup=true) - -# @info "simulation run for 10 iterations; you should have a checkpointer at 8" - -# checkpoint_file = prefix * "_iteration0.jld2" - -# set!(simulation, checkpoint_file) +run!(simulation) -# coupled_model = OceanSeaIceModel(simulation.model.ocean; atmosphere, radiation) +@info "simulation run for 10 iterations; you should have a checkpointer at 8" -# simulation = Simulation(coupled_model; Δt=10, stop_iteration=20) -# simulation.callbacks[:progress] = Callback(progress, IterationInterval(1)) +simulation.stop_iteration = 20 -# run!(simulation) \ No newline at end of file +run!(simulation, pickup=true) diff --git a/src/OceanSeaIceModels/ocean_sea_ice_model.jl b/src/OceanSeaIceModels/ocean_sea_ice_model.jl index 931fc7670..55ed786a8 100644 --- a/src/OceanSeaIceModels/ocean_sea_ice_model.jl +++ b/src/OceanSeaIceModels/ocean_sea_ice_model.jl @@ -1,5 +1,3 @@ -export run_coupled! - using Oceananigans using Oceananigans.OutputWriters: checkpoint_path using Oceananigans.TimeSteppers: Clock @@ -33,19 +31,6 @@ end const OSIM = OceanSeaIceModel const OSIMSIM = Simulation{<:OceanSeaIceModel} -function run_coupled!(sim::OSIMSIM; pickup) - if pickup - checkpoint_file_path = checkpoint_path(pickup, sim.model.ocean.output_writers) - set!(sim.model.ocean, checkpoint_file_path) - # Setting the atmosphere time to the ocean time - sim.model.atmosphere.clock.time = sim.model.ocean.model.clock.time - end - - @show sim.model.clock.time - run!(sim) -end - - function Base.summary(model::OSIM) A = nameof(typeof(architecture(model))) return string("OceanSeaIceModel{$A}", @@ -130,7 +115,7 @@ function OceanSeaIceModel(ocean, sea_ice=FreezingLimitedOceanTemperature(eltype( pop!(ocean.callbacks, :wall_time_limit_exceeded, nothing) pop!(ocean.callbacks, :nan_checker, nothing) end - + if sea_ice isa SeaIceSimulation if !isnothing(sea_ice.callbacks) pop!(sea_ice.callbacks, :stop_time_exceeded, nothing) From d15439f4410d466b7254dd1eb9d69d74b798ff60 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 13 Mar 2025 17:25:31 +1100 Subject: [PATCH 05/54] managed to pick up! --- examples/checkpointer_mwe.jl | 30 ++++++++++++++++++++---------- 1 file changed, 20 insertions(+), 10 deletions(-) diff --git a/examples/checkpointer_mwe.jl b/examples/checkpointer_mwe.jl index 07f78f116..e01598106 100644 --- a/examples/checkpointer_mwe.jl +++ b/examples/checkpointer_mwe.jl @@ -1,3 +1,4 @@ +#= using ClimaOcean using Oceananigans using Oceananigans.Units @@ -11,13 +12,10 @@ Nx = 144 Ny = 60 Nz = 40 -depth = 6000meters -z_faces = exponential_z_faces(; Nz, depth) - grid = LatitudeLongitudeGrid(arch; size = (Nx, Ny, Nz), halo = (7, 7, 7), - z = z_faces, + z = (-6000, 0), latitude = (-75, 75), longitude = (0, 360)) @@ -28,12 +26,12 @@ ocean = ocean_simulation(grid) # S=ECCOMetadata(:salinity; dates=date)) radiation = Radiation(arch) - +=# atmosphere = JRA55PrescribedAtmosphere(arch; backend=JRA55NetCDFBackend(41)) coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation) -simulation = Simulation(coupled_model; Δt=10, stop_iteration=10) +simulation = Simulation(coupled_model; Δt=10, stop_iteration=50) wall_time = Ref(time_ns()) @@ -75,7 +73,7 @@ simulation.output_writers[:surface] = JLD2OutputWriter(ocean.model, outputs; output_dir = "." prefix = "checkpointer_mwe" -simulation.output_writers[:checkpoint] = Checkpointer(ocean.model; +ocean.output_writers[:checkpoint] = Checkpointer(ocean.model; schedule = IterationInterval(2), prefix = prefix, dir = output_dir, @@ -84,10 +82,22 @@ simulation.output_writers[:checkpoint] = Checkpointer(ocean.model; # @show simulation -run!(simulation) +using ClimaOcean.OceanSeaIceModels: OceanSeaIceModel, OSIMSIM +import Oceananigans.OutputWriters: write_output! +import Oceananigans.Simulations: pickup! + +write_output!(c::Checkpointer, model::OceanSeaIceModel) = write_output!(c::Checkpointer, model.ocean.model) + +function pickup!(sim::OSIMSIM, pickup) + @info "i try to properly pick up" + set!(sim.model.ocean, pickup) + return nothing +end + +# run!(simulation) -@info "simulation run for 10 iterations; you should have a checkpointer at 8" +# @info "simulation run for 10 iterations; you should have a checkpointer at 8" -simulation.stop_iteration = 20 +# simulation.stop_iteration = 20 run!(simulation, pickup=true) From a0dca49d2ff849bbb3439a1ebc542682f6200687 Mon Sep 17 00:00:00 2001 From: Taimoor Sohail Date: Thu, 13 Mar 2025 17:32:43 +1100 Subject: [PATCH 06/54] Changed checkointer_mwe.jl --- Project.toml | 2 +- examples/checkpointer_mwe.jl | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Project.toml b/Project.toml index 909454590..0bf2e45a4 100644 --- a/Project.toml +++ b/Project.toml @@ -41,7 +41,7 @@ JLD2 = "0.4, 0.5" KernelAbstractions = "0.9" MPI = "0.20" NCDatasets = "0.12, 0.13, 0.14" -Oceananigans = "0.95.26 - 0.99" +Oceananigans = "0.95.20 - 0.99" OffsetArrays = "1.14" Scratch = "1" SeawaterPolynomials = "0.3.4" diff --git a/examples/checkpointer_mwe.jl b/examples/checkpointer_mwe.jl index 7613d63e1..a3e0207ea 100644 --- a/examples/checkpointer_mwe.jl +++ b/examples/checkpointer_mwe.jl @@ -91,7 +91,7 @@ ocean.output_writers[:checkpoint] = Checkpointer(ocean.model; # @show simulation -run_coupled!(simulation, pickup=true) +run!(simulation) # @info "simulation run for 10 iterations; you should have a checkpointer at 8" From 3fdf7856a2ab2477ab55f39aff60ca6990fe5298 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Fri, 14 Mar 2025 00:41:08 +1100 Subject: [PATCH 07/54] extends methods to work with OSIM and OSIMSIM --- src/OceanSeaIceModels/ocean_sea_ice_model.jl | 39 ++++++++++++++------ 1 file changed, 27 insertions(+), 12 deletions(-) diff --git a/src/OceanSeaIceModels/ocean_sea_ice_model.jl b/src/OceanSeaIceModels/ocean_sea_ice_model.jl index 55ed786a8..31d7a03c9 100644 --- a/src/OceanSeaIceModels/ocean_sea_ice_model.jl +++ b/src/OceanSeaIceModels/ocean_sea_ice_model.jl @@ -13,11 +13,11 @@ import Oceananigans: fields, prognostic_fields import Oceananigans.Architectures: architecture import Oceananigans.Fields: set! import Oceananigans.Models: timestepper, NaNChecker, default_nan_checker -import Oceananigans.OutputWriters: default_included_properties +import Oceananigans.OutputWriters: default_included_properties, checkpointer_address, + write_output!, initialize_jld2_file! import Oceananigans.Simulations: reset!, initialize!, iteration, run! import Oceananigans.TimeSteppers: time_step!, update_state!, time import Oceananigans.Utils: prettytime -import Oceananigans.Models: timestepper, NaNChecker, default_nan_checker struct OceanSeaIceModel{I, A, O, F, C, Arch} <: AbstractModel{Nothing, Arch} architecture :: Arch @@ -74,6 +74,31 @@ function initialize!(model::OSIM) return nothing end + +initialize_jld2_file!(filepath, init, jld2_kw, including, outputs, model::OceanSeaIceModel) = + initialize_jld2_file!(filepath, init, jld2_kw, including, outputs, model.ocean.model) + +set!(model::OSIM, filepath::AbstractString) = set!(model.ocean.model, filepath) + +write_output!(c::Checkpointer, model::OSIM) = write_output!(c, model.ocean.model) + +function set!(sim::OSIMSIM, pickup::Union{Bool, Integer, String}) + checkpoint_file_path = checkpoint_path(pickup, sim.output_writers) + + set!(sim.model, checkpoint_file_path) + + sim.model.clock.iteration = sim.model.ocean.model.clock.iteration + sim.model.clock.time = sim.model.ocean.model.clock.time + + # Setting the atmosphere time to the ocean time + sim.model.atmosphere.clock.iteration = sim.model.ocean.model.clock.iteration + sim.model.atmosphere.clock.time = sim.model.ocean.model.clock.time + + return nothing +end + +checkpointer_address(::OSIM) = "HydrostaticFreeSurfaceModel" + reference_density(unsupported) = throw(ArgumentError("Cannot extract reference density from $(typeof(unsupported))")) @@ -161,16 +186,6 @@ function default_nan_checker(model::OceanSeaIceModel) return nan_checker end -# TODO: picking up OceanSeaIceModel simulations from a checkpoint is a WIP -function set!(sim::OSIMSIM, pickup::Union{Bool, Integer, String}) - checkpoint_file_path = checkpoint_path(pickup, sim.model.ocean.output_writers) - set!(sim.model.ocean, filepath = checkpoint_file_path) - set!(sim.model.ocean, pickup = pickup) - # Setting the atmosphere time to the ocean time - sim.model.atmosphere.clock.time = sim.model.ocean.model.clock.time - return nothing -end - @kernel function _above_freezing_ocean_temperature!(T, grid, S, ℵ, liquidus) i, j = @index(Global, NTuple) Nz = size(grid, 3) From 2328bb5916beada097d68c6df3dcf4a0568cae42 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Fri, 14 Mar 2025 00:41:44 +1100 Subject: [PATCH 08/54] mwe --- examples/checkpointer_mwe.jl | 35 ++++++++++------------------------- 1 file changed, 10 insertions(+), 25 deletions(-) diff --git a/examples/checkpointer_mwe.jl b/examples/checkpointer_mwe.jl index e01598106..b8e9b74a4 100644 --- a/examples/checkpointer_mwe.jl +++ b/examples/checkpointer_mwe.jl @@ -1,4 +1,3 @@ -#= using ClimaOcean using Oceananigans using Oceananigans.Units @@ -26,17 +25,19 @@ ocean = ocean_simulation(grid) # S=ECCOMetadata(:salinity; dates=date)) radiation = Radiation(arch) -=# + atmosphere = JRA55PrescribedAtmosphere(arch; backend=JRA55NetCDFBackend(41)) coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation) -simulation = Simulation(coupled_model; Δt=10, stop_iteration=50) +simulation = Simulation(coupled_model; Δt=10, stop_iteration=8) wall_time = Ref(time_ns()) function progress(sim) ocean = sim.model.ocean + atmosphere = sim.model.atmosphere + u, v, w = ocean.model.velocities T = ocean.model.tracers.T @@ -49,7 +50,7 @@ function progress(sim) step_time = 1e-9 * (time_ns() - wall_time[]) - msg = @sprintf("Iter: %d, simulation time: %s, atmosphere time: %s, Δt: %s", iteration(sim), prettytime(sim), prettytime(atmosphere.clock.time), prettytime(sim.Δt)) + msg = @sprintf("Iter: %d, sim time: %s, atmos time: %s, ocean time: %s", iteration(sim), sim.model.clock.time, atmosphere.clock.time, ocean.model.clock.time) msg *= @sprintf(", max|u|: (%.2e, %.2e, %.2e) m s⁻¹, extrema(T): (%.2f, %.2f) ᵒC, wall time: %s", umax..., Tmax, Tmin, prettytime(step_time)) @@ -64,7 +65,7 @@ outputs = merge(ocean.model.tracers, ocean.model.velocities) simulation.output_writers[:surface] = JLD2OutputWriter(ocean.model, outputs; schedule = IterationInterval(2), - filename = "checkpointer_mwe_surface", + filename = "surface", indices = (:, :, grid.Nz), with_halos = true, overwrite_existing = true, @@ -73,31 +74,15 @@ simulation.output_writers[:surface] = JLD2OutputWriter(ocean.model, outputs; output_dir = "." prefix = "checkpointer_mwe" -ocean.output_writers[:checkpoint] = Checkpointer(ocean.model; - schedule = IterationInterval(2), +simulation.output_writers[:checkpoint] = Checkpointer(ocean.model; + schedule = IterationInterval(3), prefix = prefix, dir = output_dir, verbose = true, overwrite_existing = true) -# @show simulation - -using ClimaOcean.OceanSeaIceModels: OceanSeaIceModel, OSIMSIM -import Oceananigans.OutputWriters: write_output! -import Oceananigans.Simulations: pickup! - -write_output!(c::Checkpointer, model::OceanSeaIceModel) = write_output!(c::Checkpointer, model.ocean.model) - -function pickup!(sim::OSIMSIM, pickup) - @info "i try to properly pick up" - set!(sim.model.ocean, pickup) - return nothing -end - -# run!(simulation) - -# @info "simulation run for 10 iterations; you should have a checkpointer at 8" +run!(simulation) -# simulation.stop_iteration = 20 +simulation.stop_iteration += 5 run!(simulation, pickup=true) From 9ab0df109e3cb05aef99f2f44d98b8246857faf9 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Fri, 14 Mar 2025 00:51:02 +1100 Subject: [PATCH 09/54] simplify --- src/OceanSeaIceModels/ocean_sea_ice_model.jl | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/src/OceanSeaIceModels/ocean_sea_ice_model.jl b/src/OceanSeaIceModels/ocean_sea_ice_model.jl index 31d7a03c9..8f440b49d 100644 --- a/src/OceanSeaIceModels/ocean_sea_ice_model.jl +++ b/src/OceanSeaIceModels/ocean_sea_ice_model.jl @@ -74,18 +74,15 @@ function initialize!(model::OSIM) return nothing end - -initialize_jld2_file!(filepath, init, jld2_kw, including, outputs, model::OceanSeaIceModel) = +initialize_jld2_file!(filepath, init, jld2_kw, including, outputs, model::OSIM) = initialize_jld2_file!(filepath, init, jld2_kw, including, outputs, model.ocean.model) -set!(model::OSIM, filepath::AbstractString) = set!(model.ocean.model, filepath) - write_output!(c::Checkpointer, model::OSIM) = write_output!(c, model.ocean.model) function set!(sim::OSIMSIM, pickup::Union{Bool, Integer, String}) checkpoint_file_path = checkpoint_path(pickup, sim.output_writers) - set!(sim.model, checkpoint_file_path) + set!(sim.model.ocean.model, checkpoint_file_path) sim.model.clock.iteration = sim.model.ocean.model.clock.iteration sim.model.clock.time = sim.model.ocean.model.clock.time @@ -177,10 +174,10 @@ function OceanSeaIceModel(ocean, sea_ice=FreezingLimitedOceanTemperature(eltype( return ocean_sea_ice_model end -time(coupled_model::OceanSeaIceModel) = coupled_model.clock.time +time(coupled_model::OSIM) = coupled_model.clock.time # Check for NaNs in the first prognostic field (generalizes to prescribed velocities). -function default_nan_checker(model::OceanSeaIceModel) +function default_nan_checker(model::OSIM) u_ocean = model.ocean.model.velocities.u nan_checker = NaNChecker((; u_ocean)) return nan_checker From 5bd34f81c8687279fe05e64517181c4da38b3975 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Fri, 14 Mar 2025 08:36:09 +1100 Subject: [PATCH 10/54] tidying up --- src/OceanSeaIceModels/ocean_sea_ice_model.jl | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/OceanSeaIceModels/ocean_sea_ice_model.jl b/src/OceanSeaIceModels/ocean_sea_ice_model.jl index 8f440b49d..aa5cd3b29 100644 --- a/src/OceanSeaIceModels/ocean_sea_ice_model.jl +++ b/src/OceanSeaIceModels/ocean_sea_ice_model.jl @@ -63,6 +63,8 @@ default_included_properties(::OSIM) = tuple() prognostic_fields(cm::OSIM) = nothing fields(::OSIM) = NamedTuple() default_clock(TT) = Oceananigans.TimeSteppers.Clock{TT}(0, 0, 1) +time(model::OSIM) = model.clock.time +checkpointer_address(::OSIM) = "HydrostaticFreeSurfaceModel" function reset!(model::OSIM) reset!(model.ocean) @@ -94,8 +96,6 @@ function set!(sim::OSIMSIM, pickup::Union{Bool, Integer, String}) return nothing end -checkpointer_address(::OSIM) = "HydrostaticFreeSurfaceModel" - reference_density(unsupported) = throw(ArgumentError("Cannot extract reference density from $(typeof(unsupported))")) @@ -174,8 +174,6 @@ function OceanSeaIceModel(ocean, sea_ice=FreezingLimitedOceanTemperature(eltype( return ocean_sea_ice_model end -time(coupled_model::OSIM) = coupled_model.clock.time - # Check for NaNs in the first prognostic field (generalizes to prescribed velocities). function default_nan_checker(model::OSIM) u_ocean = model.ocean.model.velocities.u From 9e2af01a4ae5739a2d5a773519cc39e04e9f9f9a Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Fri, 14 Mar 2025 08:44:46 +1100 Subject: [PATCH 11/54] bit cleaner --- src/OceanSeaIceModels/ocean_sea_ice_model.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/OceanSeaIceModels/ocean_sea_ice_model.jl b/src/OceanSeaIceModels/ocean_sea_ice_model.jl index aa5cd3b29..faea8e4c1 100644 --- a/src/OceanSeaIceModels/ocean_sea_ice_model.jl +++ b/src/OceanSeaIceModels/ocean_sea_ice_model.jl @@ -60,7 +60,7 @@ prettytime(model::OSIM) = prettytime(model.clock.time) iteration(model::OSIM) = model.clock.iteration timestepper(::OSIM) = nothing default_included_properties(::OSIM) = tuple() -prognostic_fields(cm::OSIM) = nothing +prognostic_fields(::OSIM) = nothing fields(::OSIM) = NamedTuple() default_clock(TT) = Oceananigans.TimeSteppers.Clock{TT}(0, 0, 1) time(model::OSIM) = model.clock.time From 4f5ff2e6ef6ae38e71ef1112316fc552205ce32f Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Fri, 14 Mar 2025 09:36:29 +1100 Subject: [PATCH 12/54] set!(sim::OSIMSIM{PrescribedAtmosphere}) --- src/OceanSeaIceModels/ocean_sea_ice_model.jl | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/src/OceanSeaIceModels/ocean_sea_ice_model.jl b/src/OceanSeaIceModels/ocean_sea_ice_model.jl index faea8e4c1..a9e9704ab 100644 --- a/src/OceanSeaIceModels/ocean_sea_ice_model.jl +++ b/src/OceanSeaIceModels/ocean_sea_ice_model.jl @@ -30,6 +30,7 @@ end const OSIM = OceanSeaIceModel const OSIMSIM = Simulation{<:OceanSeaIceModel} +const OSIMSIMPA = Simulation{<:OceanSeaIceModel{<:Any, <:PrescribedAtmosphere}} function Base.summary(model::OSIM) A = nameof(typeof(architecture(model))) @@ -66,22 +67,16 @@ default_clock(TT) = Oceananigans.TimeSteppers.Clock{TT}(0, 0, time(model::OSIM) = model.clock.time checkpointer_address(::OSIM) = "HydrostaticFreeSurfaceModel" -function reset!(model::OSIM) - reset!(model.ocean) - return nothing -end +reset!(model::OSIM) = reset!(model.ocean) -function initialize!(model::OSIM) - initialize!(model.ocean) - return nothing -end +initialize!(model::OSIM) = initialize!(model.ocean) initialize_jld2_file!(filepath, init, jld2_kw, including, outputs, model::OSIM) = initialize_jld2_file!(filepath, init, jld2_kw, including, outputs, model.ocean.model) write_output!(c::Checkpointer, model::OSIM) = write_output!(c, model.ocean.model) -function set!(sim::OSIMSIM, pickup::Union{Bool, Integer, String}) +function set!(sim::OSIMSIMPA, pickup::Union{Bool, Integer, String}) checkpoint_file_path = checkpoint_path(pickup, sim.output_writers) set!(sim.model.ocean.model, checkpoint_file_path) From ba2521ae81a7ccc6e5e94762aa12b150e647c11c Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Fri, 14 Mar 2025 09:47:07 +1100 Subject: [PATCH 13/54] cleaner mwe --- examples/checkpointer_mwe.jl | 45 +++++++----------------------------- 1 file changed, 8 insertions(+), 37 deletions(-) diff --git a/examples/checkpointer_mwe.jl b/examples/checkpointer_mwe.jl index b8e9b74a4..4a069a300 100644 --- a/examples/checkpointer_mwe.jl +++ b/examples/checkpointer_mwe.jl @@ -1,15 +1,14 @@ using ClimaOcean using Oceananigans -using Oceananigans.Units using CFTime using Dates using Printf arch = CPU() -Nx = 144 -Ny = 60 -Nz = 40 +Nx = 120 +Ny = 50 +Nz = 20 grid = LatitudeLongitudeGrid(arch; size = (Nx, Ny, Nz), @@ -20,13 +19,9 @@ grid = LatitudeLongitudeGrid(arch; ocean = ocean_simulation(grid) -# date = DateTimeProlepticGregorian(1993, 1, 1) -# set!(ocean.model, T=ECCOMetadata(:temperature; dates=date), -# S=ECCOMetadata(:salinity; dates=date)) - radiation = Radiation(arch) -atmosphere = JRA55PrescribedAtmosphere(arch; backend=JRA55NetCDFBackend(41)) +atmosphere = JRA55PrescribedAtmosphere(arch; backend=JRA55NetCDFBackend(4)) coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation) @@ -38,21 +33,10 @@ function progress(sim) ocean = sim.model.ocean atmosphere = sim.model.atmosphere - u, v, w = ocean.model.velocities - T = ocean.model.tracers.T - - Tmax = maximum(interior(T)) - Tmin = minimum(interior(T)) - - umax = (maximum(abs, interior(u)), - maximum(abs, interior(v)), - maximum(abs, interior(w))) - step_time = 1e-9 * (time_ns() - wall_time[]) - msg = @sprintf("Iter: %d, sim time: %s, atmos time: %s, ocean time: %s", iteration(sim), sim.model.clock.time, atmosphere.clock.time, ocean.model.clock.time) - msg *= @sprintf(", max|u|: (%.2e, %.2e, %.2e) m s⁻¹, extrema(T): (%.2f, %.2f) ᵒC, wall time: %s", - umax..., Tmax, Tmin, prettytime(step_time)) + msg = @sprintf("iteration: %d, sim time: %s, atmos time: %s, ocean time: %s, wall time: %s", + iteration(sim), sim.model.clock.time, atmosphere.clock.time, ocean.model.clock.time, prettytime(step_time)) @info msg @@ -61,23 +45,10 @@ end simulation.callbacks[:progress] = Callback(progress, IterationInterval(1)) -outputs = merge(ocean.model.tracers, ocean.model.velocities) - -simulation.output_writers[:surface] = JLD2OutputWriter(ocean.model, outputs; - schedule = IterationInterval(2), - filename = "surface", - indices = (:, :, grid.Nz), - with_halos = true, - overwrite_existing = true, - array_type = Array{Float32}) - -output_dir = "." -prefix = "checkpointer_mwe" - simulation.output_writers[:checkpoint] = Checkpointer(ocean.model; schedule = IterationInterval(3), - prefix = prefix, - dir = output_dir, + prefix = "checkpointer", + dir = ".", verbose = true, overwrite_existing = true) From b75d6e0aec1021292a6bdfe638fb71012b809e03 Mon Sep 17 00:00:00 2001 From: Taimoor Sohail Date: Fri, 14 Mar 2025 10:56:38 +1100 Subject: [PATCH 14/54] Changed the function to set! --- examples/checkpointer_mwe.jl | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/examples/checkpointer_mwe.jl b/examples/checkpointer_mwe.jl index 4cc6816e4..2ef86f18e 100644 --- a/examples/checkpointer_mwe.jl +++ b/examples/checkpointer_mwe.jl @@ -83,12 +83,17 @@ ocean.output_writers[:checkpoint] = Checkpointer(ocean.model; using ClimaOcean.OceanSeaIceModels: OSIMSIM # import Oceananigans.OutputWriters: write_output! -import Oceananigans.Simulations: pickup! +import Oceananigans.Simulations: pickup!, set! # write_output!(c::Checkpointer, model::OceanSeaIceModel) = write_output!(c::Checkpointer, model.ocean.model) -function pickup!(sim::OSIMSIM, pickup) - @info "i try to properly pick up" +# function pickup!(sim::OSIMSIM, pickup) +# set!(sim.model.ocean, pickup) +# return nothing +# end + +function set!(sim::OSIMSIM, pickup) + @info "I am trying the set! method instead" set!(sim.model.ocean, pickup) return nothing end @@ -97,6 +102,6 @@ run!(simulation) # @info "simulation run for 10 iterations; you should have a checkpointer at 8" -# simulation.stop_iteration = 20 +simulation.stop_iteration = 20 run!(simulation, pickup=true) From dfb57fa054bcaad46b588849fcca65c242d0f9d9 Mon Sep 17 00:00:00 2001 From: Taimoor Sohail Date: Fri, 14 Mar 2025 11:00:44 +1100 Subject: [PATCH 15/54] Merge NCC changes --- examples/checkpointer_mwe.jl | 8 -------- 1 file changed, 8 deletions(-) diff --git a/examples/checkpointer_mwe.jl b/examples/checkpointer_mwe.jl index c08c10146..4a069a300 100644 --- a/examples/checkpointer_mwe.jl +++ b/examples/checkpointer_mwe.jl @@ -21,19 +21,11 @@ ocean = ocean_simulation(grid) radiation = Radiation(arch) -<<<<<<< HEAD -atmosphere = JRA55PrescribedAtmosphere(arch; backend=JRA55NetCDFBackend(41)) - -coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation) - -simulation = Simulation(coupled_model; Δt=10, stop_iteration=10) -======= atmosphere = JRA55PrescribedAtmosphere(arch; backend=JRA55NetCDFBackend(4)) coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation) simulation = Simulation(coupled_model; Δt=10, stop_iteration=8) ->>>>>>> ba2521ae81a7ccc6e5e94762aa12b150e647c11c wall_time = Ref(time_ns()) From 76dd72074f400fdf7563fd16af73bb66068ef8ce Mon Sep 17 00:00:00 2001 From: Taimoor Sohail Date: Fri, 14 Mar 2025 11:06:23 +1100 Subject: [PATCH 16/54] reverting near_global_ocean.jl example --- examples/near_global_ocean_simulation.jl | 62 ++++++++++-------------- 1 file changed, 25 insertions(+), 37 deletions(-) diff --git a/examples/near_global_ocean_simulation.jl b/examples/near_global_ocean_simulation.jl index 5fb3c56fb..8c9fe84d7 100644 --- a/examples/near_global_ocean_simulation.jl +++ b/examples/near_global_ocean_simulation.jl @@ -32,8 +32,8 @@ using Printf arch = CPU() -Nx = 144 -Ny = 60 +Nx = 1440 +Ny = 600 Nz = 40 depth = 6000meters @@ -53,21 +53,21 @@ grid = LatitudeLongitudeGrid(arch; # (i) all the minor enclosed basins except the 3 largest `major_basins`, as well as # (ii) regions that are shallower than `minimum_depth`. -# bottom_height = regrid_bathymetry(grid; -# minimum_depth = 10meters, -# interpolation_passes = 5, -# major_basins = 3) +bottom_height = regrid_bathymetry(grid; + minimum_depth = 10meters, + interpolation_passes = 5, + major_basins = 3) -# grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map=true) +grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map=true) # Let's see what the bathymetry looks like: -# h = grid.immersed_boundary.bottom_height +h = grid.immersed_boundary.bottom_height -# fig, ax, hm = heatmap(h, colormap=:deep, colorrange=(-depth, 0)) -# Colorbar(fig[0, 1], hm, label="Bottom height (m)", vertical=false) -# save("bathymetry.png", fig) -# nothing #hide +fig, ax, hm = heatmap(h, colormap=:deep, colorrange=(-depth, 0)) +Colorbar(fig[0, 1], hm, label="Bottom height (m)", vertical=false) +save("bathymetry.png", fig) +nothing #hide # ![](bathymetry.png) @@ -83,9 +83,9 @@ ocean.model # We initialize the ocean model with ECCO2 temperature and salinity for January 1, 1993. -# date = DateTimeProlepticGregorian(1993, 1, 1) -# set!(ocean.model, T=ECCOMetadata(:temperature; dates=date), -# S=ECCOMetadata(:salinity; dates=date)) +date = DateTimeProlepticGregorian(1993, 1, 1) +set!(ocean.model, T=Metadata(:temperature; dates=date, dataset=ECCO4Monthly()), + S=Metadata(:salinity; dates=date, dataset=ECCO4Monthly())) # ### Prescribed atmosphere and radiation # @@ -117,7 +117,7 @@ coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation) # We then create a coupled simulation. We start with a small-ish time step of 90 seconds. # We run the simulation for 10 days with this small-ish time step. -simulation = Simulation(coupled_model; Δt=90, stop_iteration=10) +simulation = Simulation(coupled_model; Δt=90, stop_time=10days) # We define a callback function to monitor the simulation's progress, @@ -146,7 +146,7 @@ function progress(sim) wall_time[] = time_ns() end -simulation.callbacks[:progress] = Callback(progress, IterationInterval(1)) +simulation.callbacks[:progress] = Callback(progress, TimeInterval(5days)) # ### Set up output writers # @@ -156,23 +156,13 @@ simulation.callbacks[:progress] = Callback(progress, IterationInterval(1)) # Below, we use `indices` to save only the values of the variables at the surface, which corresponds to `k = grid.Nz` outputs = merge(ocean.model.tracers, ocean.model.velocities) -# ocean.output_writers[:surface] = JLD2OutputWriter(ocean.model, outputs; -# schedule = IterationInterval(1), -# filename = "near_global_surface_fields", -# indices = (:, :, grid.Nz), -# with_halos = true, -# overwrite_existing = true, -# array_type = Array{Float32}) - -output_dir = "." -prefix = "near_global" - -ocean.output_writers[:checkpoint] = Checkpointer(ocean.model; - schedule = TimeInterval(3minutes), - prefix = prefix, - dir = output_dir)#, - # verbose = true, - # overwrite_existing = true) +ocean.output_writers[:surface] = JLD2OutputWriter(ocean.model, outputs; + schedule = TimeInterval(1days), + filename = "near_global_surface_fields", + indices = (:, :, grid.Nz), + with_halos = true, + overwrite_existing = true, + array_type = Array{Float32}) # ### Spinning up the simulation # @@ -185,12 +175,11 @@ run!(simulation) # ### Running the simulation for real # After the initial spin up of 10 days, we can increase the time-step and run for longer. -#= + simulation.stop_time = 60days simulation.Δt = 10minutes run!(simulation) - # ## A pretty movie # # It's time to make a pretty movie of the simulation. First we load the output we've been saving on @@ -270,4 +259,3 @@ end nothing #hide # ![](near_global_ocean_surface.mp4) -=# \ No newline at end of file From 3177cf260dadd0cf8d7828c4d3483537ded6f4de Mon Sep 17 00:00:00 2001 From: Taimoor Sohail Date: Fri, 14 Mar 2025 11:12:06 +1100 Subject: [PATCH 17/54] Update near_global_ocean_simulation.jl --- examples/near_global_ocean_simulation.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/near_global_ocean_simulation.jl b/examples/near_global_ocean_simulation.jl index 8c9fe84d7..1f005b7bd 100644 --- a/examples/near_global_ocean_simulation.jl +++ b/examples/near_global_ocean_simulation.jl @@ -30,7 +30,7 @@ using Printf # The total depth of the domain is set to 6000 meters. # Finally, we specify the architecture for the simulation, which in this case is a GPU. -arch = CPU() +arch = GPU() Nx = 1440 Ny = 600 From 868c87076d8c6bc185b3672735aee172c68b354d Mon Sep 17 00:00:00 2001 From: Taimoor Sohail Date: Fri, 14 Mar 2025 11:17:58 +1100 Subject: [PATCH 18/54] Update Project.toml --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 0bf2e45a4..909454590 100644 --- a/Project.toml +++ b/Project.toml @@ -41,7 +41,7 @@ JLD2 = "0.4, 0.5" KernelAbstractions = "0.9" MPI = "0.20" NCDatasets = "0.12, 0.13, 0.14" -Oceananigans = "0.95.20 - 0.99" +Oceananigans = "0.95.26 - 0.99" OffsetArrays = "1.14" Scratch = "1" SeawaterPolynomials = "0.3.4" From 5f02123af9bc89b79fc7310e6912c8b55048aae1 Mon Sep 17 00:00:00 2001 From: Taimoor Sohail Date: Fri, 14 Mar 2025 16:07:36 +1100 Subject: [PATCH 19/54] Added checkpointing test; integrated checkpointer into one_degree example; added set_clock! function --- examples/checkpointer_mwe.jl | 59 -------------------- examples/one_degree_simulation.jl | 17 +++++- src/OceanSeaIceModels/ocean_sea_ice_model.jl | 14 +++-- test/test_ocean_sea_ice_model.jl | 46 +++++++++++++++ 4 files changed, 70 insertions(+), 66 deletions(-) delete mode 100644 examples/checkpointer_mwe.jl diff --git a/examples/checkpointer_mwe.jl b/examples/checkpointer_mwe.jl deleted file mode 100644 index 4a069a300..000000000 --- a/examples/checkpointer_mwe.jl +++ /dev/null @@ -1,59 +0,0 @@ -using ClimaOcean -using Oceananigans -using CFTime -using Dates -using Printf - -arch = CPU() - -Nx = 120 -Ny = 50 -Nz = 20 - -grid = LatitudeLongitudeGrid(arch; - size = (Nx, Ny, Nz), - halo = (7, 7, 7), - z = (-6000, 0), - latitude = (-75, 75), - longitude = (0, 360)) - -ocean = ocean_simulation(grid) - -radiation = Radiation(arch) - -atmosphere = JRA55PrescribedAtmosphere(arch; backend=JRA55NetCDFBackend(4)) - -coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation) - -simulation = Simulation(coupled_model; Δt=10, stop_iteration=8) - -wall_time = Ref(time_ns()) - -function progress(sim) - ocean = sim.model.ocean - atmosphere = sim.model.atmosphere - - step_time = 1e-9 * (time_ns() - wall_time[]) - - msg = @sprintf("iteration: %d, sim time: %s, atmos time: %s, ocean time: %s, wall time: %s", - iteration(sim), sim.model.clock.time, atmosphere.clock.time, ocean.model.clock.time, prettytime(step_time)) - - @info msg - - wall_time[] = time_ns() -end - -simulation.callbacks[:progress] = Callback(progress, IterationInterval(1)) - -simulation.output_writers[:checkpoint] = Checkpointer(ocean.model; - schedule = IterationInterval(3), - prefix = "checkpointer", - dir = ".", - verbose = true, - overwrite_existing = true) - -run!(simulation) - -simulation.stop_iteration += 5 - -run!(simulation, pickup=true) diff --git a/examples/one_degree_simulation.jl b/examples/one_degree_simulation.jl index d1c2f6ecc..69cccac3e 100644 --- a/examples/one_degree_simulation.jl +++ b/examples/one_degree_simulation.jl @@ -170,6 +170,19 @@ ocean.output_writers[:surface] = JLD2OutputWriter(ocean.model, outputs; overwrite_existing = true, array_type = Array{Float32}) +# We also save a series of Checkpointers which will enable the simulation to be restarted +# from a later point; ideal for longer runs + +directory_checkpointer = "." +prefix_checkpointer = "one_degree_checkpointer" + +simulation.output_writers[:checkpoint] = Checkpointer(ocean.model; + schedule = TimeInterval(5days), + prefix = prefix_checkpointer, + dir = directory_checkpointer, + verbose = true, + overwrite_existing = true) + # ### Ready to run # We are ready to press the big red button and run the simulation. @@ -177,12 +190,12 @@ ocean.output_writers[:surface] = JLD2OutputWriter(ocean.model, outputs; # After we run for a short time (here we set up the simulation with `stop_time = 10days`), # we increase the timestep and run for longer. -run!(simulation) +run!(simulation, pickup=true) simulation.Δt = 20minutes simulation.stop_time = 360days -run!(simulation) +run!(simulation, pickup=true) # ### A pretty movie # diff --git a/src/OceanSeaIceModels/ocean_sea_ice_model.jl b/src/OceanSeaIceModels/ocean_sea_ice_model.jl index a9e9704ab..71c9f8391 100644 --- a/src/OceanSeaIceModels/ocean_sea_ice_model.jl +++ b/src/OceanSeaIceModels/ocean_sea_ice_model.jl @@ -76,17 +76,21 @@ initialize_jld2_file!(filepath, init, jld2_kw, including, outputs, model::OSIM) write_output!(c::Checkpointer, model::OSIM) = write_output!(c, model.ocean.model) -function set!(sim::OSIMSIMPA, pickup::Union{Bool, Integer, String}) - checkpoint_file_path = checkpoint_path(pickup, sim.output_writers) - - set!(sim.model.ocean.model, checkpoint_file_path) - +function set_clock!(sim) sim.model.clock.iteration = sim.model.ocean.model.clock.iteration sim.model.clock.time = sim.model.ocean.model.clock.time # Setting the atmosphere time to the ocean time sim.model.atmosphere.clock.iteration = sim.model.ocean.model.clock.iteration sim.model.atmosphere.clock.time = sim.model.ocean.model.clock.time + return nothing +end + +function set!(sim::OSIMSIMPA, pickup::Union{Bool, Integer, String}) + checkpoint_file_path = checkpoint_path(pickup, sim.output_writers) + + set!(sim.model.ocean.model, checkpoint_file_path) + set_clock!(sim) return nothing end diff --git a/test/test_ocean_sea_ice_model.jl b/test/test_ocean_sea_ice_model.jl index 3d2395e11..841754997 100644 --- a/test/test_ocean_sea_ice_model.jl +++ b/test/test_ocean_sea_ice_model.jl @@ -93,3 +93,49 @@ using ClimaSeaIce.SeaIceThermodynamics: melting_temperature end end +function testbed_coupled_simulation(grid; stop_iteration=8) + ocean = ocean_simulation(grid) + + radiation = Radiation(arch) + + atmosphere = JRA55PrescribedAtmosphere(arch; backend=JRA55NetCDFBackend(4)) + + coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation) + + simulation = Simulation(coupled_model; Δt=10, stop_iteration) + + simulation.output_writers[:checkpoint] = Checkpointer(ocean.model; + schedule = IterationInterval(3), + prefix = "checkpointer", + dir = ".", + verbose = true, + overwrite_existing = true) + return simulation +end + +@testset "Checkpointer test" begin + for arch in test_architectures + + Nx, Ny, Nz = 40, 25, 10 + + grid = LatitudeLongitudeGrid(arch; + size = (Nx, Ny, Nz), + halo = (7, 7, 7), + z = (-6000, 0), + latitude = (-75, 75), + longitude = (0, 360)) + + simulation = testbed_coupled_simulation(grid; stop_iteration=8) + + run!(simulation, pickup=true) + + # create a new simulation and pick up from the last checkpointer + new_simulation = testbed_coupled_simulation(grid; stop_iteration=13) + + run!(new_simulation, pickup=true) + + # test that ocean tiem step and iteration is the same as the atmosphere + @test new_simulation.model.atmosphere.clock.iteration ≈ new_simulation.model.ocean.model.clock.iteration + @test new_simulation.model.atmosphere.clock.time ≈ new_simulation.model.ocean.model.clock.time + end +end From f9087aad4437c4d8322a5d5a3d292a5a63f0b8ec Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Fri, 14 Mar 2025 16:36:01 +1100 Subject: [PATCH 20/54] set clock method for each simulation type --- src/OceanSeaIceModels/ocean_sea_ice_model.jl | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/src/OceanSeaIceModels/ocean_sea_ice_model.jl b/src/OceanSeaIceModels/ocean_sea_ice_model.jl index 71c9f8391..29846e4a3 100644 --- a/src/OceanSeaIceModels/ocean_sea_ice_model.jl +++ b/src/OceanSeaIceModels/ocean_sea_ice_model.jl @@ -76,13 +76,15 @@ initialize_jld2_file!(filepath, init, jld2_kw, including, outputs, model::OSIM) write_output!(c::Checkpointer, model::OSIM) = write_output!(c, model.ocean.model) -function set_clock!(sim) - sim.model.clock.iteration = sim.model.ocean.model.clock.iteration - sim.model.clock.time = sim.model.ocean.model.clock.time +function set_clock!(sim::PrescribedAtmosphere, time, iter) + sim.clock.time = time + sim.clock.iteration = iter + return nothing +end - # Setting the atmosphere time to the ocean time - sim.model.atmosphere.clock.iteration = sim.model.ocean.model.clock.iteration - sim.model.atmosphere.clock.time = sim.model.ocean.model.clock.time +function set_clock!(sim::OSIMSIM, time, iter) + sim.model.clock.time = time + sim.model.clock.iteration = iter return nothing end @@ -90,7 +92,10 @@ function set!(sim::OSIMSIMPA, pickup::Union{Bool, Integer, String}) checkpoint_file_path = checkpoint_path(pickup, sim.output_writers) set!(sim.model.ocean.model, checkpoint_file_path) - set_clock!(sim) + + iteration, time = sim.model.ocean.model.clock.iteration, sim.model.ocean.model.clock.time + set_clock!(sim, time, iteration) + set_clock!(sim.model.atmosphere, time, iteration) return nothing end From 4b0ebeb398592eecd3a21a0baeeabe006f8dc074 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Fri, 14 Mar 2025 16:39:06 +1100 Subject: [PATCH 21/54] Apply suggestions from code review --- examples/one_degree_simulation.jl | 11 +++-------- test/test_ocean_sea_ice_model.jl | 6 +++--- 2 files changed, 6 insertions(+), 11 deletions(-) diff --git a/examples/one_degree_simulation.jl b/examples/one_degree_simulation.jl index 69cccac3e..0916c7887 100644 --- a/examples/one_degree_simulation.jl +++ b/examples/one_degree_simulation.jl @@ -170,17 +170,12 @@ ocean.output_writers[:surface] = JLD2OutputWriter(ocean.model, outputs; overwrite_existing = true, array_type = Array{Float32}) -# We also save a series of Checkpointers which will enable the simulation to be restarted -# from a later point; ideal for longer runs - -directory_checkpointer = "." -prefix_checkpointer = "one_degree_checkpointer" +# We also save a series of Checkpointers that will enable us to restart the simulation +# from a later point; ideal for longer runs. simulation.output_writers[:checkpoint] = Checkpointer(ocean.model; schedule = TimeInterval(5days), - prefix = prefix_checkpointer, - dir = directory_checkpointer, - verbose = true, + prefix = "one_degree_checkpointer", overwrite_existing = true) # ### Ready to run diff --git a/test/test_ocean_sea_ice_model.jl b/test/test_ocean_sea_ice_model.jl index 841754997..4158cf895 100644 --- a/test/test_ocean_sea_ice_model.jl +++ b/test/test_ocean_sea_ice_model.jl @@ -134,8 +134,8 @@ end run!(new_simulation, pickup=true) - # test that ocean tiem step and iteration is the same as the atmosphere - @test new_simulation.model.atmosphere.clock.iteration ≈ new_simulation.model.ocean.model.clock.iteration - @test new_simulation.model.atmosphere.clock.time ≈ new_simulation.model.ocean.model.clock.time + # ensure the ocean and atmosphere time step and iteration are the same + @test new_simulation.model.atmosphere.clock.iteration ≈ new_simulation.model.ocean.model.clock.iteration + @test new_simulation.model.atmosphere.clock.time ≈ new_simulation.model.ocean.model.clock.time end end From 3f0b9d55a6de6a581308ce7037270b705778b225 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Fri, 14 Mar 2025 16:41:11 +1100 Subject: [PATCH 22/54] Apply suggestions from code review --- test/test_ocean_sea_ice_model.jl | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/test/test_ocean_sea_ice_model.jl b/test/test_ocean_sea_ice_model.jl index 4158cf895..f04aada65 100644 --- a/test/test_ocean_sea_ice_model.jl +++ b/test/test_ocean_sea_ice_model.jl @@ -93,6 +93,11 @@ using ClimaSeaIce.SeaIceThermodynamics: melting_temperature end end +""" + testbed_coupled_simulation(grid; stop_iteration=8) + +Return a test-bed coupled simulation with a Checkpointer. +""" function testbed_coupled_simulation(grid; stop_iteration=8) ocean = ocean_simulation(grid) From f4aa21e12527d4d0fc6a8e2a746f0f25ca659399 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Fri, 14 Mar 2025 16:49:35 +1100 Subject: [PATCH 23/54] don't pickup by default; add explanation --- examples/one_degree_simulation.jl | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/examples/one_degree_simulation.jl b/examples/one_degree_simulation.jl index 69cccac3e..d73632c38 100644 --- a/examples/one_degree_simulation.jl +++ b/examples/one_degree_simulation.jl @@ -170,7 +170,7 @@ ocean.output_writers[:surface] = JLD2OutputWriter(ocean.model, outputs; overwrite_existing = true, array_type = Array{Float32}) -# We also save a series of Checkpointers which will enable the simulation to be restarted +# We also save a series of Checkpointers which will enable the simulation to be restarted # from a later point; ideal for longer runs directory_checkpointer = "." @@ -190,12 +190,22 @@ simulation.output_writers[:checkpoint] = Checkpointer(ocean.model; # After we run for a short time (here we set up the simulation with `stop_time = 10days`), # we increase the timestep and run for longer. -run!(simulation, pickup=true) +run!(simulation) simulation.Δt = 20minutes simulation.stop_time = 360days -run!(simulation, pickup=true) +run!(simulation) + +# After we had run long-enough to produce at least one checkpointer file, if we want to restart the +# simulation we call +# +# ```julia +# run!(simulation, pickup=true) +# ``` +# +# Doing so, simulation will be picked up from the latest checkpointer. + # ### A pretty movie # From f56ca8e2f49140acdfdeeb1f1e27e03dc5968ec9 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Fri, 14 Mar 2025 17:05:57 +1100 Subject: [PATCH 24/54] move set_clock! for PrescribedAtmosphere to where it belongs --- src/OceanSeaIceModels/PrescribedAtmospheres.jl | 18 ++++++++++++------ src/OceanSeaIceModels/ocean_sea_ice_model.jl | 6 ------ 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/src/OceanSeaIceModels/PrescribedAtmospheres.jl b/src/OceanSeaIceModels/PrescribedAtmospheres.jl index 4ce53cefc..ef08c58eb 100644 --- a/src/OceanSeaIceModels/PrescribedAtmospheres.jl +++ b/src/OceanSeaIceModels/PrescribedAtmospheres.jl @@ -221,7 +221,7 @@ Base.summary(::PATP{FT}) where FT = "PrescribedAtmosphereThermodynamicsParameter function Base.show(io::IO, p::PrescribedAtmosphereThermodynamicsParameters) FT = eltype(p) - cp = p.constitutive + cp = p.constitutive hc = p.heat_capacity pt = p.phase_transitions @@ -238,9 +238,9 @@ function Base.show(io::IO, p::PrescribedAtmosphereThermodynamicsParameters) "└── PhaseTransitionParameters{$FT}", '\n', " ├── reference_vaporization_enthalpy (ℒᵛ⁰): ", prettysummary(pt.reference_vaporization_enthalpy), '\n', " ├── reference_sublimation_enthalpy (ℒˢ⁰): ", prettysummary(pt.reference_sublimation_enthalpy), '\n', - " ├── reference_temperature (T⁰): ", prettysummary(pt.reference_temperature), '\n', + " ├── reference_temperature (T⁰): ", prettysummary(pt.reference_temperature), '\n', " ├── triple_point_temperature (Tᵗʳ): ", prettysummary(pt.triple_point_temperature), '\n', - " ├── triple_point_pressure (pᵗʳ): ", prettysummary(pt.triple_point_pressure), '\n', + " ├── triple_point_pressure (pᵗʳ): ", prettysummary(pt.triple_point_pressure), '\n', " ├── water_freezing_temperature (Tᶠ): ", prettysummary(pt.water_freezing_temperature), '\n', " └── total_ice_nucleation_temperature (Tⁱ): ", prettysummary(pt.total_ice_nucleation_temperature)) end @@ -318,6 +318,12 @@ function Base.show(io::IO, pa::PrescribedAtmosphere) print(io, "└── boundary_layer_height: ", prettysummary(pa.boundary_layer_height)) end +function set_clock!(sim::PrescribedAtmosphere, time, iter) + sim.clock.time = time + sim.clock.iteration = iter + return nothing +end + function default_atmosphere_velocities(grid, times) ua = FieldTimeSeries{Center, Center, Nothing}(grid, times) va = FieldTimeSeries{Center, Center, Nothing}(grid, times) @@ -357,14 +363,14 @@ end for fts in ftses update_field_time_series!(fts, time) - end - + end + return nothing end @inline thermodynamics_parameters(atmos::PrescribedAtmosphere) = atmos.thermodynamics_parameters @inline surface_layer_height(atmos::PrescribedAtmosphere) = atmos.surface_layer_height -@inline boundary_layer_height(atmos::PrescribedAtmosphere) = atmos.boundary_layer_height +@inline boundary_layer_height(atmos::PrescribedAtmosphere) = atmos.boundary_layer_height """ PrescribedAtmosphere(grid, times; diff --git a/src/OceanSeaIceModels/ocean_sea_ice_model.jl b/src/OceanSeaIceModels/ocean_sea_ice_model.jl index 29846e4a3..b74875e38 100644 --- a/src/OceanSeaIceModels/ocean_sea_ice_model.jl +++ b/src/OceanSeaIceModels/ocean_sea_ice_model.jl @@ -76,12 +76,6 @@ initialize_jld2_file!(filepath, init, jld2_kw, including, outputs, model::OSIM) write_output!(c::Checkpointer, model::OSIM) = write_output!(c, model.ocean.model) -function set_clock!(sim::PrescribedAtmosphere, time, iter) - sim.clock.time = time - sim.clock.iteration = iter - return nothing -end - function set_clock!(sim::OSIMSIM, time, iter) sim.model.clock.time = time sim.model.clock.iteration = iter From f8c144412dd72ed8873c5a40235df377222fd0c1 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Fri, 14 Mar 2025 17:36:40 +1100 Subject: [PATCH 25/54] move set_clock! for PrescribedAtmosphere to where it belongs --- .../PrescribedAtmospheres.jl | 14 ++++++++--- src/OceanSeaIceModels/ocean_sea_ice_model.jl | 23 ++++++++++++++----- 2 files changed, 28 insertions(+), 9 deletions(-) diff --git a/src/OceanSeaIceModels/PrescribedAtmospheres.jl b/src/OceanSeaIceModels/PrescribedAtmospheres.jl index ef08c58eb..345d65f3b 100644 --- a/src/OceanSeaIceModels/PrescribedAtmospheres.jl +++ b/src/OceanSeaIceModels/PrescribedAtmospheres.jl @@ -318,9 +318,17 @@ function Base.show(io::IO, pa::PrescribedAtmosphere) print(io, "└── boundary_layer_height: ", prettysummary(pa.boundary_layer_height)) end -function set_clock!(sim::PrescribedAtmosphere, time, iter) - sim.clock.time = time - sim.clock.iteration = iter +""" + set_clock!(sim, clock) + +Set the clock of `sim`ulation to match the values of `clock`. +""" +function set_clock!(sim::PrescribedAtmosphere, clock) + sim.clock.time = clock.time + sim.clock.iteration = clock.iteration + sim.clock.last_Δt = clock.last_Δt + sim.clock.last_stage_Δt = clock.last_stage_Δt + sim.clock.stage = clock.stage return nothing end diff --git a/src/OceanSeaIceModels/ocean_sea_ice_model.jl b/src/OceanSeaIceModels/ocean_sea_ice_model.jl index b74875e38..ff149cf72 100644 --- a/src/OceanSeaIceModels/ocean_sea_ice_model.jl +++ b/src/OceanSeaIceModels/ocean_sea_ice_model.jl @@ -19,6 +19,8 @@ import Oceananigans.Simulations: reset!, initialize!, iteration, run! import Oceananigans.TimeSteppers: time_step!, update_state!, time import Oceananigans.Utils: prettytime +import .PrescribedAtmospheres: set_clock! + struct OceanSeaIceModel{I, A, O, F, C, Arch} <: AbstractModel{Nothing, Arch} architecture :: Arch clock :: C @@ -76,9 +78,17 @@ initialize_jld2_file!(filepath, init, jld2_kw, including, outputs, model::OSIM) write_output!(c::Checkpointer, model::OSIM) = write_output!(c, model.ocean.model) -function set_clock!(sim::OSIMSIM, time, iter) - sim.model.clock.time = time - sim.model.clock.iteration = iter +""" + set_clock!(sim, clock) + +Set the clock of `sim`ulation to match the values of `clock`. +""" +function set_clock!(sim::OSIMSIM, clock) + sim.model.clock.time = clock.time + sim.model.clock.iteration = clock.iteration + sim.model.clock.last_Δt = clock.last_Δt + sim.model.clock.last_stage_Δt = clock.last_stage_Δt + sim.model.clock.stage = clock.stage return nothing end @@ -87,9 +97,10 @@ function set!(sim::OSIMSIMPA, pickup::Union{Bool, Integer, String}) set!(sim.model.ocean.model, checkpoint_file_path) - iteration, time = sim.model.ocean.model.clock.iteration, sim.model.ocean.model.clock.time - set_clock!(sim, time, iteration) - set_clock!(sim.model.atmosphere, time, iteration) + clock = sim.model.ocean.model.clock + + set_clock!(sim, clock) + set_clock!(sim.model.atmosphere, clock) return nothing end From cfba574e5d9ae279806fd37a32e4700c950bf273 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Fri, 14 Mar 2025 17:37:26 +1100 Subject: [PATCH 26/54] only pickup in the second time for clarity --- test/test_ocean_sea_ice_model.jl | 25 +++++++++++++------------ 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/test/test_ocean_sea_ice_model.jl b/test/test_ocean_sea_ice_model.jl index f04aada65..53a19f5a4 100644 --- a/test/test_ocean_sea_ice_model.jl +++ b/test/test_ocean_sea_ice_model.jl @@ -41,15 +41,15 @@ using ClimaSeaIce.SeaIceThermodynamics: melting_temperature halo = (7, 7, 7), z = (-6000, 0)) - bottom_height = regrid_bathymetry(grid; + bottom_height = regrid_bathymetry(grid; minimum_depth = 10, interpolation_passes = 20, major_basins = 1) - + grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map=true) free_surface = SplitExplicitFreeSurface(grid; substeps=20) - ocean = ocean_simulation(grid; free_surface) + ocean = ocean_simulation(grid; free_surface) backend = JRA55NetCDFBackend(4) atmosphere = JRA55PrescribedAtmosphere(arch; backend) @@ -69,9 +69,9 @@ using ClimaSeaIce.SeaIceThermodynamics: melting_temperature sea_ice_grid = TripolarGrid(arch; size=(50, 50, 1), z = (-10, 0)) sea_ice_grid = ImmersedBoundaryGrid(sea_ice_grid, GridFittedBottom(bottom_height)) - sea_ice = sea_ice_simulation(sea_ice_grid) + sea_ice = sea_ice_simulation(sea_ice_grid) liquidus = sea_ice.model.ice_thermodynamics.phase_transitions.liquidus - + # Set the ocean temperature and salinity set!(ocean.model, T=temperature_metadata[1], S=salinity_metadata[1]) above_freezing_ocean_temperature!(ocean, sea_ice) @@ -100,13 +100,13 @@ Return a test-bed coupled simulation with a Checkpointer. """ function testbed_coupled_simulation(grid; stop_iteration=8) ocean = ocean_simulation(grid) - + radiation = Radiation(arch) - + atmosphere = JRA55PrescribedAtmosphere(arch; backend=JRA55NetCDFBackend(4)) - + coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation) - + simulation = Simulation(coupled_model; Δt=10, stop_iteration) simulation.output_writers[:checkpoint] = Checkpointer(ocean.model; @@ -115,6 +115,7 @@ function testbed_coupled_simulation(grid; stop_iteration=8) dir = ".", verbose = true, overwrite_existing = true) + return simulation end @@ -122,17 +123,17 @@ end for arch in test_architectures Nx, Ny, Nz = 40, 25, 10 - + grid = LatitudeLongitudeGrid(arch; size = (Nx, Ny, Nz), halo = (7, 7, 7), z = (-6000, 0), latitude = (-75, 75), longitude = (0, 360)) - + simulation = testbed_coupled_simulation(grid; stop_iteration=8) - run!(simulation, pickup=true) + run!(simulation) # create a new simulation and pick up from the last checkpointer new_simulation = testbed_coupled_simulation(grid; stop_iteration=13) From 4110aca3ce4c0f6adf767d9fdaebd0be825504ef Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Sun, 16 Mar 2025 15:55:58 +1100 Subject: [PATCH 27/54] extend set!(model::OSIM,...) instead of set!(sim:OSIMSIM,...) --- src/OceanSeaIceModels/ocean_sea_ice_model.jl | 33 +++++++++----------- 1 file changed, 14 insertions(+), 19 deletions(-) diff --git a/src/OceanSeaIceModels/ocean_sea_ice_model.jl b/src/OceanSeaIceModels/ocean_sea_ice_model.jl index ff149cf72..14cef812e 100644 --- a/src/OceanSeaIceModels/ocean_sea_ice_model.jl +++ b/src/OceanSeaIceModels/ocean_sea_ice_model.jl @@ -1,5 +1,4 @@ using Oceananigans -using Oceananigans.OutputWriters: checkpoint_path using Oceananigans.TimeSteppers: Clock using Oceananigans: SeawaterBuoyancy using ClimaSeaIce.SeaIceThermodynamics: melting_temperature @@ -78,29 +77,25 @@ initialize_jld2_file!(filepath, init, jld2_kw, including, outputs, model::OSIM) write_output!(c::Checkpointer, model::OSIM) = write_output!(c, model.ocean.model) -""" - set_clock!(sim, clock) - -Set the clock of `sim`ulation to match the values of `clock`. -""" -function set_clock!(sim::OSIMSIM, clock) - sim.model.clock.time = clock.time - sim.model.clock.iteration = clock.iteration - sim.model.clock.last_Δt = clock.last_Δt - sim.model.clock.last_stage_Δt = clock.last_stage_Δt - sim.model.clock.stage = clock.stage +function set_clock!(model::OSIM, clock) + model.clock.time = clock.time + model.clock.iteration = clock.iteration + model.clock.last_Δt = clock.last_Δt + model.clock.last_stage_Δt = clock.last_stage_Δt + model.clock.stage = clock.stage return nothing end -function set!(sim::OSIMSIMPA, pickup::Union{Bool, Integer, String}) - checkpoint_file_path = checkpoint_path(pickup, sim.output_writers) +function set!(model::OSIM, checkpoint_file_path::AbstractString) + atmosphere = model.atmosphere + ocean = model.ocean.model - set!(sim.model.ocean.model, checkpoint_file_path) + set!(ocean, checkpoint_file_path) - clock = sim.model.ocean.model.clock - - set_clock!(sim, clock) - set_clock!(sim.model.atmosphere, clock) + # ensure all clocks follow the ocean model clock + clock = ocean.clock + set_clock!(model, clock) + set_clock!(atmosphere, clock) return nothing end From a9be97c4dfd44b1ccf281a530799ad111fa8a8b7 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Sun, 16 Mar 2025 16:06:07 +1100 Subject: [PATCH 28/54] use default radiation --- test/test_ocean_sea_ice_model.jl | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/test/test_ocean_sea_ice_model.jl b/test/test_ocean_sea_ice_model.jl index de1c1fb8a..a65bef162 100644 --- a/test/test_ocean_sea_ice_model.jl +++ b/test/test_ocean_sea_ice_model.jl @@ -69,16 +69,16 @@ using ClimaSeaIce.Rheologies ##### # Adding a sea ice model to the coupled model - τua = Field{Face, Center, Nothing}(grid) + τua = Field{Face, Center, Nothing}(grid) τva = Field{Center, Face, Nothing}(grid) - dynamics = SeaIceMomentumEquation(grid; + dynamics = SeaIceMomentumEquation(grid; coriolis = ocean.model.coriolis, top_momentum_stress = (u=τua, v=τva), rheology = ElastoViscoPlasticRheology(), solver = SplitExplicitSolver(120)) - sea_ice = sea_ice_simulation(grid; dynamics, advection=WENO(order=7)) + sea_ice = sea_ice_simulation(grid; dynamics, advection=WENO(order=7)) liquidus = sea_ice.model.ice_thermodynamics.phase_transitions.liquidus # Set the ocean temperature and salinity @@ -110,11 +110,9 @@ Return a test-bed coupled simulation with a Checkpointer. function testbed_coupled_simulation(grid; stop_iteration=8) ocean = ocean_simulation(grid) - radiation = Radiation(arch) - atmosphere = JRA55PrescribedAtmosphere(arch; backend=JRA55NetCDFBackend(4)) - coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation) + coupled_model = OceanSeaIceModel(ocean; atmosphere) simulation = Simulation(coupled_model; Δt=10, stop_iteration) From afd7a660cdb0742337b1b228d8d92a65bebfe541 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Sun, 16 Mar 2025 18:05:34 +1100 Subject: [PATCH 29/54] add set!(::PrescribedAtmospher, checkopoint_file_path) --- .../PrescribedAtmospheres.jl | 28 +++++++++++++++---- 1 file changed, 23 insertions(+), 5 deletions(-) diff --git a/src/OceanSeaIceModels/PrescribedAtmospheres.jl b/src/OceanSeaIceModels/PrescribedAtmospheres.jl index 99e1ca6fb..91c17cc5b 100644 --- a/src/OceanSeaIceModels/PrescribedAtmospheres.jl +++ b/src/OceanSeaIceModels/PrescribedAtmospheres.jl @@ -7,8 +7,11 @@ using Oceananigans.OutputReaders: FieldTimeSeries, update_field_time_series!, ex using Oceananigans.TimeSteppers: Clock, tick! using Adapt +using JLD2 using Thermodynamics.Parameters: AbstractThermodynamicsParameters +import Oceananigans.Fields: set! +import Oceananigans.OutputWriters: checkpointer_address import Oceananigans.TimeSteppers: time_step! import Thermodynamics.Parameters: @@ -22,7 +25,7 @@ import Thermodynamics.Parameters: T_0, # Enthalpy reference temperature LH_v0, # Vaporization enthalpy at the reference temperature LH_s0, # Sublimation enthalpy at the reference temperature - LH_f0, # Fusionn enthalpy at the reference temperature + LH_f0, # Fusion enthalpy at the reference temperature cp_d, # Heat capacity of dry air at constant pressure cp_v, # Isobaric specific heat capacity of gaseous water vapor cp_l, # Isobaric specific heat capacity of liquid water @@ -30,7 +33,7 @@ import Thermodynamics.Parameters: cv_v, # Heat capacity of dry air at constant volume cv_l, # Isobaric specific heat capacity of liquid water cv_i, # Isobaric specific heat capacity of liquid water - e_int_v0, # what? someting about reference internal energy of water vapor + e_int_v0, # what? something about reference internal energy of water vapor T_freeze, # Freezing temperature of _pure_ water T_triple, # Triple point temperature of _pure_ water press_triple, # Triple point pressure of pure water @@ -75,7 +78,7 @@ Construct a set of parameters that define the density of moist air, ``` where ``p`` is pressure, ``T`` is temperature, ``q`` defines the partition -of total mass into vapor, liqiud, and ice mass fractions, and +of total mass into vapor, liquid, and ice mass fractions, and ``Rᵐ`` is the effective specific gas constant for the mixture, ```math @@ -318,6 +321,22 @@ function Base.show(io::IO, pa::PrescribedAtmosphere) print(io, "└── boundary_layer_height: ", prettysummary(pa.boundary_layer_height)) end +# set the clock to be the same as the ocean model +function set!(model::PrescribedAtmosphere, checkpoint_file_path) + addr = checkpointer_address(model) + + jldopen(checkpoint_file_path, "r") do file + checkpointed_clock = file["$addr/clock"] + + # Update model clock + set_clock!(model, checkpointed_clock) + end + + return nothing +end + +checkpointer_address(::PrescribedAtmosphere) = "HydrostaticFreeSurfaceModel" + """ set_clock!(sim, clock) @@ -434,7 +453,7 @@ struct TwoBandDownwellingRadiation{SW, LW} end """ - TwoBandDownwellingRadiation(shortwave=nothing, longwave=nothing) + TwoBandDownwellingRadiation(; shortwave=nothing, longwave=nothing) Return a two-band model for downwelling radiation (split in a shortwave band and a longwave band) that passes through the atmosphere and arrives at the surface of ocean @@ -448,4 +467,3 @@ Adapt.adapt_structure(to, tsdr::TwoBandDownwellingRadiation) = adapt(to, tsdr.longwave)) end # module - From 1fddbf21eeb6881e5dd025fb94517a1c84e22402 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Sun, 16 Mar 2025 18:06:20 +1100 Subject: [PATCH 30/54] use set!(::PrescribedAtmosphere,...) --- src/OceanSeaIceModels/ocean_sea_ice_model.jl | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/src/OceanSeaIceModels/ocean_sea_ice_model.jl b/src/OceanSeaIceModels/ocean_sea_ice_model.jl index 33d27a0fa..306f4bea0 100644 --- a/src/OceanSeaIceModels/ocean_sea_ice_model.jl +++ b/src/OceanSeaIceModels/ocean_sea_ice_model.jl @@ -102,12 +102,10 @@ function set!(model::OSIM, checkpoint_file_path::AbstractString) ocean = model.ocean.model set!(ocean, checkpoint_file_path) + set!(atmosphere, checkpoint_file_path) + + set_clock!(model, ocean.clock) - # ensure all clocks follow the ocean model clock - clock = ocean.clock - set_clock!(model, clock) - set_clock!(atmosphere, clock) - return nothing end @@ -227,4 +225,3 @@ function above_freezing_ocean_temperature!(ocean, sea_ice::SeaIceSimulation) return nothing end - From 8699207c9b422f6d3b41e2d565c54a6e7260653d Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Sun, 16 Mar 2025 18:06:39 +1100 Subject: [PATCH 31/54] check all clocks are aligned --- test/test_ocean_sea_ice_model.jl | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/test/test_ocean_sea_ice_model.jl b/test/test_ocean_sea_ice_model.jl index a65bef162..4f701ddfd 100644 --- a/test/test_ocean_sea_ice_model.jl +++ b/test/test_ocean_sea_ice_model.jl @@ -147,8 +147,11 @@ end run!(new_simulation, pickup=true) - # ensure the ocean and atmosphere time step and iteration are the same - @test new_simulation.model.atmosphere.clock.iteration ≈ new_simulation.model.ocean.model.clock.iteration - @test new_simulation.model.atmosphere.clock.time ≈ new_simulation.model.ocean.model.clock.time + # ensure the ocean, atmosphere, and coupled model are all at same time and iteration + clock = new_simulation.model.ocean.model.clock + @test new_simulation.model.atmosphere.clock.iteration ≈ clock.iteration + @test new_simulation.model.atmosphere.clock.time ≈ clock.time + @test new_simulation.model.clock.iteration ≈ clock.iteration + @test new_simulation.model.clock.time ≈ clock.time end end From 9f9ca89561211d535926dd4570d362b677eab6dc Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Sun, 16 Mar 2025 18:09:07 +1100 Subject: [PATCH 32/54] drop unused aliases --- src/OceanSeaIceModels/ocean_sea_ice_model.jl | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/OceanSeaIceModels/ocean_sea_ice_model.jl b/src/OceanSeaIceModels/ocean_sea_ice_model.jl index 306f4bea0..ed138ec21 100644 --- a/src/OceanSeaIceModels/ocean_sea_ice_model.jl +++ b/src/OceanSeaIceModels/ocean_sea_ice_model.jl @@ -30,8 +30,6 @@ mutable struct OceanSeaIceModel{I, A, O, F, C, Arch} <: AbstractModel{Nothing, A end const OSIM = OceanSeaIceModel -const OSIMSIM = Simulation{<:OceanSeaIceModel} -const OSIMSIMPA = Simulation{<:OceanSeaIceModel{<:Any, <:PrescribedAtmosphere}} function Base.summary(model::OSIM) A = nameof(typeof(architecture(model))) From ba80b569ace5a504efebf1a7509f6899e77b6c8f Mon Sep 17 00:00:00 2001 From: Taimoor Sohail Date: Wed, 19 Mar 2025 13:38:08 +1100 Subject: [PATCH 33/54] Merge branch 'ts/checkpointer-for-coupled-model' of github.com:CliMA/ClimaOcean.jl into ts/checkpointer-for-coupled-model --- examples/checkpointer_mwe.jl | 72 ++ examples/generate_atmos_dataset.jl | 17 + src/CoupledSimulation.jl | 39 + src/DataWrangling/JRA55.jl | 724 ++++++++++++++++++ src/DistributedUtils.jl | 163 ++++ ...est_ocean_sea_ice_model_parameter_space.jl | 49 ++ test/test_simulations.jl | 75 ++ 7 files changed, 1139 insertions(+) create mode 100755 examples/checkpointer_mwe.jl create mode 100755 examples/generate_atmos_dataset.jl create mode 100755 src/CoupledSimulation.jl create mode 100755 src/DataWrangling/JRA55.jl create mode 100755 src/DistributedUtils.jl create mode 100755 test/test_ocean_sea_ice_model_parameter_space.jl create mode 100755 test/test_simulations.jl diff --git a/examples/checkpointer_mwe.jl b/examples/checkpointer_mwe.jl new file mode 100755 index 000000000..a0007fa0b --- /dev/null +++ b/examples/checkpointer_mwe.jl @@ -0,0 +1,72 @@ +using ClimaOcean +using Oceananigans +using CFTime +using Dates +using Printf + +arch = CPU() + +Nx = 120 +Ny = 50 +Nz = 20 + +grid = LatitudeLongitudeGrid(arch; + size = (Nx, Ny, Nz), + halo = (7, 7, 7), + z = (-6000, 0), + latitude = (-75, 75), + longitude = (0, 360)) + + +function testbed_coupled_simulation(grid; stop_iteration=8) + ocean = ocean_simulation(grid) + + radiation = Radiation(arch) + + atmosphere = JRA55PrescribedAtmosphere(arch; backend=JRA55NetCDFBackend(4)) + + coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation) + + return Simulation(coupled_model; Δt=10, stop_iteration) +end + +simulation = testbed_coupled_simulation(grid; stop_iteration=8) + +wall_time = Ref(time_ns()) + +function progress(sim) + ocean = sim.model.ocean + atmosphere = sim.model.atmosphere + + step_time = 1e-9 * (time_ns() - wall_time[]) + + msg = @sprintf("iteration: %d, sim time: %s, atmos time: %s, ocean time: %s, wall time: %s", + iteration(sim), sim.model.clock.time, atmosphere.clock.time, ocean.model.clock.time, prettytime(step_time)) + + @info msg + + wall_time[] = time_ns() +end + +simulation.callbacks[:progress] = Callback(progress, IterationInterval(1)) + +simulation.output_writers[:checkpoint] = Checkpointer(ocean.model; + schedule = IterationInterval(3), + prefix = "checkpointer", + dir = ".", + verbose = true, + overwrite_existing = true) + +run!(simulation) + +new_simulation = testbed_coupled_simulation(grid; stop_iteration=14) + +new_simulation.output_writers[:checkpoint] = Checkpointer(ocean.model; + schedule = IterationInterval(3), + prefix = "checkpointer", + dir = ".", + verbose = true, + overwrite_existing = true) + + +run!(new_simulation, pickup=true) diff --git a/examples/generate_atmos_dataset.jl b/examples/generate_atmos_dataset.jl new file mode 100755 index 000000000..8ba1d3233 --- /dev/null +++ b/examples/generate_atmos_dataset.jl @@ -0,0 +1,17 @@ +using ClimaOcean +using Oceananigans +using JLD2 + +time_indices = 1:1 + +qt = ClimaOcean.JRA55.JRA55_field_time_series(:specific_humidity; time_indices) +Tt = ClimaOcean.JRA55.JRA55_field_time_series(:temperature; time_indices) +pt = ClimaOcean.JRA55.JRA55_field_time_series(:sea_level_pressure; time_indices) + +Nx, Ny, Nz = size(qt[1]) + +q = Array(interior(qt[1], 1:4:Nx, 1:4:Ny, 1)) +T = Array(interior(Tt[1], 1:4:Nx, 1:4:Ny, 1)) +p = Array(interior(pt[1], 1:4:Nx, 1:4:Ny, 1)) + +@save "atmospheric_state.jld2" q T p diff --git a/src/CoupledSimulation.jl b/src/CoupledSimulation.jl new file mode 100755 index 000000000..f654a5084 --- /dev/null +++ b/src/CoupledSimulation.jl @@ -0,0 +1,39 @@ +module Simulations + +import Oceananigans: run! + +function run!(coupled_model::OceanSeaIceModel; schedule, + dir = ".", + prefix = "checkpoint", + overwrite_existing = false, + verbose = false, + cleanup = false, + properties = default_checkpointed_properties(coupled_model.ocean.model)) + + function run!(sim; pickup=false) + + start_run = time_ns() + + if we_want_to_pickup(pickup) + set!(sim, pickup) + end + + sim.initialized = false + sim.running = true + sim.run_wall_time = 0.0 + + while sim.running + time_step!(sim) + end + + for callback in values(sim.callbacks) + finalize!(callback, sim) + end + + # Increment the wall clock + end_run = time_ns() + sim.run_wall_time += 1e-9 * (end_run - start_run) + + return nothing + end + diff --git a/src/DataWrangling/JRA55.jl b/src/DataWrangling/JRA55.jl new file mode 100755 index 000000000..1fa655b73 --- /dev/null +++ b/src/DataWrangling/JRA55.jl @@ -0,0 +1,724 @@ +module JRA55 + +using Oceananigans +using Oceananigans.Units + +using Oceananigans: location +using Oceananigans.Architectures: arch_array +using Oceananigans.DistributedComputations +using Oceananigans.DistributedComputations: child_architecture +using Oceananigans.BoundaryConditions: fill_halo_regions! +using Oceananigans.Grids: λnodes, φnodes, on_architecture +using Oceananigans.Fields: interpolate! +using Oceananigans.OutputReaders: Cyclical, TotallyInMemory, AbstractInMemoryBackend, FlavorOfFTS, time_indices +using Oceananigans.TimeSteppers: Clock + +using ClimaOcean +using ClimaOcean.DistributedUtils: @root +using ClimaOcean.DataWrangling: download_progress + +using ClimaOcean.OceanSeaIceModels: + PrescribedAtmosphere, + PrescribedAtmosphereThermodynamicsParameters, + TwoBandDownwellingRadiation + +using CUDA: @allowscalar + +using NCDatasets +using JLD2 +using Dates +using Scratch + +import Oceananigans.Fields: set! +import Oceananigans.OutputReaders: new_backend, update_field_time_series! +using Downloads: download + +download_jra55_cache::String = "" + +function __init__() + global download_jra55_cache = @get_scratch!("JRA55") +end + +struct JRA55Data end +const JRA55PrescribedAtmosphere = PrescribedAtmosphere{<:Any, <:JRA55Data} + +# A list of all variables provided in the JRA55 dataset: +JRA55_variable_names = (:river_freshwater_flux, + :rain_freshwater_flux, + :snow_freshwater_flux, + :iceberg_freshwater_flux, + :specific_humidity, + :sea_level_pressure, + :relative_humidity, + :downwelling_longwave_radiation, + :downwelling_shortwave_radiation, + :temperature, + :eastward_velocity, + :northward_velocity) + +filenames = Dict( + :river_freshwater_flux => "RYF.friver.1990_1991.nc", # Freshwater fluxes from rivers + :rain_freshwater_flux => "RYF.prra.1990_1991.nc", # Freshwater flux from rainfall + :snow_freshwater_flux => "RYF.prsn.1990_1991.nc", # Freshwater flux from snowfall + :iceberg_freshwater_flux => "RYF.licalvf.1990_1991.nc", # Freshwater flux from calving icebergs + :specific_humidity => "RYF.huss.1990_1991.nc", # Surface specific humidity + :sea_level_pressure => "RYF.psl.1990_1991.nc", # Sea level pressure + :relative_humidity => "RYF.rhuss.1990_1991.nc", # Surface relative humidity + :downwelling_longwave_radiation => "RYF.rlds.1990_1991.nc", # Downwelling longwave radiation + :downwelling_shortwave_radiation => "RYF.rsds.1990_1991.nc", # Downwelling shortwave radiation + :temperature => "RYF.tas.1990_1991.nc", # Near-surface air temperature + :eastward_velocity => "RYF.uas.1990_1991.nc", # Eastward near-surface wind + :northward_velocity => "RYF.vas.1990_1991.nc", # Northward near-surface wind +) + +jra55_short_names = Dict( + :river_freshwater_flux => "friver", # Freshwater fluxes from rivers + :rain_freshwater_flux => "prra", # Freshwater flux from rainfall + :snow_freshwater_flux => "prsn", # Freshwater flux from snowfall + :iceberg_freshwater_flux => "licalvf", # Freshwater flux from calving icebergs + :specific_humidity => "huss", # Surface specific humidity + :sea_level_pressure => "psl", # Sea level pressure + :relative_humidity => "rhuss", # Surface relative humidity + :downwelling_longwave_radiation => "rlds", # Downwelling longwave radiation + :downwelling_shortwave_radiation => "rsds", # Downwelling shortwave radiation + :temperature => "tas", # Near-surface air temperature + :eastward_velocity => "uas", # Eastward near-surface wind + :northward_velocity => "vas", # Northward near-surface wind +) + +field_time_series_short_names = Dict( + :river_freshwater_flux => "Fri", # Freshwater fluxes from rivers + :rain_freshwater_flux => "Fra", # Freshwater flux from rainfall + :snow_freshwater_flux => "Fsn", # Freshwater flux from snowfall + :iceberg_freshwater_flux => "Fic", # Freshwater flux from calving icebergs + :specific_humidity => "qa", # Surface specific humidity + :sea_level_pressure => "pa", # Sea level pressure + :relative_humidity => "rh", # Surface relative humidity + :downwelling_longwave_radiation => "Ql", # Downwelling longwave radiation + :downwelling_shortwave_radiation => "Qs", # Downwelling shortwave radiation + :temperature => "Ta", # Near-surface air temperature + :eastward_velocity => "ua", # Eastward near-surface wind + :northward_velocity => "va", # Northward near-surface wind +) + +urls = Dict( + :shortwave_radiation => "https://www.dropbox.com/scl/fi/z6fkvmd9oe3ycmaxta131/" * + "RYF.rsds.1990_1991.nc?rlkey=r7q6zcbj6a4fxsq0f8th7c4tc&dl=0", + + :river_freshwater_flux => "https://www.dropbox.com/scl/fi/21ggl4p74k4zvbf04nb67/" * + "RYF.friver.1990_1991.nc?rlkey=ny2qcjkk1cfijmwyqxsfm68fz&dl=0", + + :rain_freshwater_flux => "https://www.dropbox.com/scl/fi/5icl1gbd7f5hvyn656kjq/" * + "RYF.prra.1990_1991.nc?rlkey=iifyjm4ppwyd8ztcek4dtx0k8&dl=0", + + :snow_freshwater_flux => "https://www.dropbox.com/scl/fi/1r4ajjzb3643z93ads4x4/" * + "RYF.prsn.1990_1991.nc?rlkey=auyqpwn060cvy4w01a2yskfah&dl=0", + + :iceberg_freshwater_flux => "https://www.dropbox.com/scl/fi/44nc5y27ohvif7lkvpyv0/" * + "RYF.licalvf.1990_1991.nc?rlkey=w7rqu48y2baw1efmgrnmym0jk&dl=0", + + :specific_humidity => "https://www.dropbox.com/scl/fi/66z6ymfr4ghkynizydc29/" * + "RYF.huss.1990_1991.nc?rlkey=107yq04aew8lrmfyorj68v4td&dl=0", + + :sea_level_pressure => "https://www.dropbox.com/scl/fi/0fk332027oru1iiseykgp/" * + "RYF.psl.1990_1991.nc?rlkey=4xpr9uah741483aukok6d7ctt&dl=0", + + :relative_humidity => "https://www.dropbox.com/scl/fi/1agwsp0lzvntuyf8bm9la/" * + "RYF.rhuss.1990_1991.nc?rlkey=8cd0vs7iy1rw58b9pc9t68gtz&dl=0", + + :downwelling_longwave_radiation => "https://www.dropbox.com/scl/fi/y6r62szkirrivua5nqq61/" * + "RYF.rlds.1990_1991.nc?rlkey=wt9yq3cyrvs2rbowoirf4nkum&dl=0", + + :downwelling_shortwave_radiation => "https://www.dropbox.com/scl/fi/z6fkvmd9oe3ycmaxta131/" * + "RYF.rsds.1990_1991.nc?rlkey=r7q6zcbj6a4fxsq0f8th7c4tc&dl=0", + + :temperature => "https://www.dropbox.com/scl/fi/fpl0npwi476w635g6lke9/" * + "RYF.tas.1990_1991.nc?rlkey=0skb9pe6lgbfbiaoybe7m945s&dl=0", + + :eastward_velocity => "https://www.dropbox.com/scl/fi/86wetpqla2x97isp8092g/" * + "RYF.uas.1990_1991.nc?rlkey=rcaf18sh1yz0v9g4hjm1249j0&dl=0", + + :northward_velocity => "https://www.dropbox.com/scl/fi/d38sflo9ddljstd5jwgml/" * + "RYF.vas.1990_1991.nc?rlkey=f9y3e57kx8xrb40gbstarf0x6&dl=0", +) + +compute_bounding_nodes(::Nothing, ::Nothing, LH, hnodes) = nothing +compute_bounding_nodes(bounds, ::Nothing, LH, hnodes) = bounds + +function compute_bounding_nodes(x::Number, ::Nothing, LH, hnodes) + ϵ = convert(typeof(x), 0.001) # arbitrary? + return (x - ϵ, x + ϵ) +end + +# TODO: remove the allowscalar +function compute_bounding_nodes(::Nothing, grid, LH, hnodes) + hg = hnodes(grid, LH()) + h₁ = @allowscalar minimum(hg) + h₂ = @allowscalar maximum(hg) + return h₁, h₂ +end + +function compute_bounding_indices(::Nothing, hc) + Nh = length(hc) + return 1, Nh +end + +function compute_bounding_indices(bounds::Tuple, hc) + h₁, h₂ = bounds + Nh = length(hc) + + # The following should work. If ᵒ are the extrema of nodes we want to + # interpolate to, and the following is a sketch of the JRA55 native grid, + # + # 1 2 3 4 5 + # | | | | | | + # | x ᵒ | x | x | x ᵒ | x | + # | | | | | | + # 1 2 3 4 5 6 + # + # then for example, we should find that (iᵢ, i₂) = (1, 5). + # So we want to reduce the first index by one, and limit them + # both by the available data. There could be some mismatch due + # to the use of different coordinate systems (ie whether λ ∈ (0, 360) + # which we may also need to handle separately. + i₁ = searchsortedfirst(hc, h₁) + i₂ = searchsortedfirst(hc, h₂) + i₁ = max(1, i₁ - 1) + i₂ = min(Nh, i₂) + + return i₁, i₂ +end + +infer_longitudinal_topology(::Nothing) = Periodic + +function infer_longitudinal_topology(λbounds) + λ₁, λ₂ = λbounds + TX = λ₂ - λ₁ ≈ 360 ? Periodic : Bounded + return TX +end + +function compute_bounding_indices(longitude, latitude, grid, LX, LY, λc, φc) + λbounds = compute_bounding_nodes(longitude, grid, LX, λnodes) + φbounds = compute_bounding_nodes(latitude, grid, LY, φnodes) + + i₁, i₂ = compute_bounding_indices(λbounds, λc) + j₁, j₂ = compute_bounding_indices(φbounds, φc) + TX = infer_longitudinal_topology(λbounds) + + return i₁, i₂, j₁, j₂, TX +end + +# Convert dates to range until Oceananigans supports dates natively +function jra55_times(native_times, start_time=native_times[1]) + + times = [] + for native_time in native_times + time = native_time - start_time + time = Second(time).value + push!(times, time) + end + + return times +end + +struct JRA55NetCDFBackend <: AbstractInMemoryBackend{Int} + start :: Int + length :: Int +end + +""" + JRA55NetCDFBackend(length) + +Represents a JRA55 FieldTimeSeries backed by JRA55 native .nc files. +""" +JRA55NetCDFBackend(length) = JRA55NetCDFBackend(1, length) + +Base.length(backend::JRA55NetCDFBackend) = backend.length +Base.summary(backend::JRA55NetCDFBackend) = string("JRA55NetCDFBackend(", backend.start, ", ", backend.length, ")") + +const JRA55NetCDFFTS = FlavorOfFTS{<:Any, <:Any, <:Any, <:Any, <:JRA55NetCDFBackend} + +function set!(fts::JRA55NetCDFFTS, path::String=fts.path, name::String=fts.name) + + ds = Dataset(path) + + # Note that each file should have the variables + # - ds["time"]: time coordinate + # - ds["lon"]: longitude at the location of the variable + # - ds["lat"]: latitude at the location of the variable + # - ds["lon_bnds"]: bounding longitudes between which variables are averaged + # - ds["lat_bnds"]: bounding latitudes between which variables are averaged + # - ds[shortname]: the variable data + + # Nodes at the variable location + λc = ds["lon"][:] + φc = ds["lat"][:] + LX, LY, LZ = location(fts) + i₁, i₂, j₁, j₂, TX = compute_bounding_indices(nothing, nothing, fts.grid, LX, LY, λc, φc) + + nn = time_indices(fts) + nn = collect(nn) + + if issorted(nn) + data = ds[name][i₁:i₂, j₁:j₂, nn] + else + # The time indices may be cycling past 1; eg ti = [6, 7, 8, 1]. + # However, DiskArrays does not seem to support loading data with unsorted + # indices. So to handle this, we load the data in chunks, where each chunk's + # indices are sorted, and then glue the data together. + m = findfirst(n -> n == 1, nn) + n1 = nn[1:m-1] + n2 = nn[m:end] + + data1 = ds[name][i₁:i₂, j₁:j₂, n1] + data2 = ds[name][i₁:i₂, j₁:j₂, n2] + data = cat(data1, data2, dims=3) + end + + close(ds) + + copyto!(interior(fts, :, :, 1, :), data) + fill_halo_regions!(fts) + + return nothing +end + +new_backend(::JRA55NetCDFBackend, start, length) = JRA55NetCDFBackend(start, length) + +""" + JRA55_field_time_series(variable_name; + architecture = CPU(), + grid = nothing, + location = nothing, + url = nothing, + dir = download_jra55_cache, + filename = nothing, + shortname = nothing, + latitude = nothing, + longitude = nothing, + backend = InMemory(), + time_indexing = Cyclical(), + preprocess_chunk_size = 10, + preprocess_architecture = CPU(), + time_indices = nothing) + +Return a `FieldTimeSeries` containing atmospheric reanalysis data for `variable_name`, +which describes one of the variables in the "repeat year forcing" dataset derived +from the Japanese 55-year atmospheric reanalysis for driving ocean-sea ice models (JRA55-do). +For more information about the derivation of the repeat-year forcing dataset, see + +> Stewart et al. (2020). JRA55-do-based repeat year forcing datasets for driving ocean–sea-ice models, _Ocean Modelling_, **147**, 101557, https://doi.org/10.1016/j.ocemod.2019.101557. + +The `variable_name`s (and their `shortname`s used in NetCDF files) available from the JRA55-do are: +- `:river_freshwater_flux` ("friver") +- `:rain_freshwater_flux` ("prra") +- `:snow_freshwater_flux` ("prsn") +- `:iceberg_freshwater_flux` ("licalvf") +- `:specific_humidity` ("huss") +- `:sea_level_pressure` ("psl") +- `:relative_humidity` ("rhuss") +- `:downwelling_longwave_radiation` ("rlds") +- `:downwelling_shortwave_radiation` ("rsds") +- `:temperature` ("ras") +- `:eastward_velocity` ("uas") +- `:northward_velocity` ("vas") + +Keyword arguments +================= + +- `architecture`: Architecture for the `FieldTimeSeries`. + Default: CPU() + +- `time_indices`: Indices of the timeseries to extract from file. + For example, `time_indices=1:3` returns a + `FieldTimeSeries` with the first three time snapshots + of `variable_name`. + +- `latitude`: Guiding latitude bounds for the resulting grid. + Used to slice the data when loading into memory. + Default: nothing, which retains the latitude range of the native grid. + +- `longitude`: Guiding longitude bounds for the resulting grid. + Used to slice the data when loading into memory. + Default: nothing, which retains the longitude range of the native grid. + +- `url`: The url accessed to download the data for `variable_name`. + Default: `ClimaOcean.JRA55.urls[variable_name]`. + +- `filename`: The name of the downloaded file. + Default: `ClimaOcean.JRA55.filenames[variable_name]`. + +- `shortname`: The "short name" of `variable_name` inside its NetCDF file. + Default: `ClimaOcean.JRA55.jra55_short_names[variable_name]`. + +- `interpolated_file`: file holding an Oceananigans compatible version of the JRA55 data. + If it does not exist it will be generated. + +- `time_chunks_in_memory`: number of fields held in memory. If `nothing` then the whole timeseries + is loaded (not recommended). +""" +function JRA55_field_time_series(variable_name; + architecture = CPU(), + grid = nothing, + location = nothing, + url = nothing, + dir = download_jra55_cache, + filename = nothing, + shortname = nothing, + latitude = nothing, + longitude = nothing, + backend = InMemory(), + time_indexing = Cyclical(), + preprocess_chunk_size = 10, + preprocess_architecture = CPU(), + time_indices = nothing) + + # OnDisk backends do not support time interpolation! + # Disallow OnDisk for JRA55 dataset loading + if backend isa OnDisk + msg = string("We cannot load the JRA55 dataset with an `OnDisk` backend") + throw(ArgumentError(msg)) + end + + if isnothing(filename) && !(variable_name ∈ JRA55_variable_names) + variable_strs = Tuple(" - :$name \n" for name in JRA55_variable_names) + variables_msg = prod(variable_strs) + + msg = string("The variable :$variable_name is not provided by the JRA55-do dataset!", '\n', + "The variables provided by the JRA55-do dataset are:", '\n', + variables_msg) + + throw(ArgumentError(msg)) + end + + filepath = isnothing(filename) ? joinpath(dir, filenames[variable_name]) : joinpath(dir, filename) + + if !isnothing(filename) && !isfile(filepath) && isnothing(url) + throw(ArgumentError("A filename was provided without a url, but the file does not exist.\n \ + If intended, please provide both the filename and url that should be used \n \ + to download the new file.")) + end + + isnothing(filename) && (filename = filenames[variable_name]) + isnothing(shortname) && (shortname = jra55_short_names[variable_name]) + isnothing(url) && (url = urls[variable_name]) + + # Record some important user decisions + totally_in_memory = backend isa TotallyInMemory + on_native_grid = isnothing(grid) + !on_native_grid && backend isa JRA55NetCDFBackend && error("Can't use custom grid with JRA55NetCDFBackend.") + + jld2_filepath = joinpath(dir, string("JRA55_repeat_year_", variable_name, ".jld2")) + fts_name = field_time_series_short_names[variable_name] + + # Note, we don't re-use existing jld2 files. + @root begin + isfile(filepath) || download(url, filepath; progress=download_progress) + isfile(jld2_filepath) && rm(jld2_filepath) + end + + # Determine default time indices + if totally_in_memory + # In this case, the whole time series is in memory. + # Either the time series is short, or we are doing a limited-area + # simulation, like in a single column. So, we conservatively + # set a default `time_indices = 1:2`. + isnothing(time_indices) && (time_indices = 1:2) + time_indices_in_memory = time_indices + native_fts_architecture = architecture + else + # In this case, part or all of the time series will be stored in a file. + # Note: if the user has provided a grid, we will have to preprocess the + # .nc JRA55 data into a .jld2 file. In this case, `time_indices` refers + # to the time_indices that we will preprocess; + # by default we choose all of them. The architecture is only the + # architecture used for preprocessing, which typically will be CPU() + # even if we would like the final FieldTimeSeries on the GPU. + isnothing(time_indices) && (time_indices = :) + + if backend isa JRA55NetCDFBackend + time_indices_in_memory = 1:length(backend) + native_fts_architecture = architecture + else # then `time_indices_in_memory` refers to preprocessing + maximum_index = min(preprocess_chunk_size, length(time_indices)) + time_indices_in_memory = 1:maximum_index + native_fts_architecture = preprocess_architecture + end + end + + # Set a default location. + if isnothing(location) + LX = LY = Center + else + LX, LY = location + end + + ds = Dataset(filepath) + + # Note that each file should have the variables + # - ds["time"]: time coordinate + # - ds["lon"]: longitude at the location of the variable + # - ds["lat"]: latitude at the location of the variable + # - ds["lon_bnds"]: bounding longitudes between which variables are averaged + # - ds["lat_bnds"]: bounding latitudes between which variables are averaged + # - ds[shortname]: the variable data + + # Nodes at the variable location + λc = ds["lon"][:] + φc = ds["lat"][:] + + # Interfaces for the "native" JRA55 grid + λn = ds["lon_bnds"][1, :] + φn = ds["lat_bnds"][1, :] + + # The .nc coordinates lon_bnds and lat_bnds do not include + # the last interface, so we push them here. + push!(φn, 90) + push!(λn, λn[1] + 360) + + # TODO: support loading just part of the JRA55 data. + # Probably with arguments that take latitude, longitude bounds. + i₁, i₂, j₁, j₂, TX = compute_bounding_indices(longitude, latitude, grid, LX, LY, λc, φc) + + native_times = ds["time"][time_indices] + data = ds[shortname][i₁:i₂, j₁:j₂, time_indices_in_memory] + λr = λn[i₁:i₂+1] + φr = φn[j₁:j₂+1] + Nrx, Nry, Nt = size(data) + close(ds) + + N = (Nrx, Nry) + H = min.(N, (3, 3)) + + JRA55_native_grid = LatitudeLongitudeGrid(native_fts_architecture, Float32; + halo = H, + size = N, + longitude = λr, + latitude = φr, + topology = (TX, Bounded, Flat)) + + boundary_conditions = FieldBoundaryConditions(JRA55_native_grid, (Center, Center, Nothing)) + times = jra55_times(native_times) + + if backend isa JRA55NetCDFBackend + fts = FieldTimeSeries{Center, Center, Nothing}(JRA55_native_grid, times; + backend, + time_indexing, + boundary_conditions, + path = filepath, + name = shortname) + + # Fill the data in a GPU-friendly manner + copyto!(interior(fts, :, :, 1, :), data) + fill_halo_regions!(fts) + + return fts + else + # Make times into an array for later preprocessing + if !totally_in_memory + times = collect(times) + end + + native_fts = FieldTimeSeries{Center, Center, Nothing}(JRA55_native_grid, times; + time_indexing, + boundary_conditions) + + # Fill the data in a GPU-friendly manner + copyto!(interior(native_fts, :, :, 1, :), data) + fill_halo_regions!(native_fts) + + if on_native_grid && totally_in_memory + return native_fts + + elseif totally_in_memory # but not on the native grid! + boundary_conditions = FieldBoundaryConditions(grid, (LX, LY, Nothing)) + fts = FieldTimeSeries{LX, LY, Nothing}(grid, times; time_indexing, boundary_conditions) + interpolate!(fts, native_fts) + return fts + end + end + + @info "Pre-processing JRA55 $variable_name data into a JLD2 file..." + + preprocessing_grid = on_native_grid ? JRA55_native_grid : grid + + # Re-open the dataset! + ds = Dataset(filepath) + all_datetimes = ds["time"][time_indices] + all_Nt = length(all_datetimes) + + all_times = jra55_times(all_datetimes) + + on_disk_fts = FieldTimeSeries{LX, LY, Nothing}(preprocessing_grid, all_times; + boundary_conditions, + backend = OnDisk(), + path = jld2_filepath, + name = fts_name) + + # Save data to disk, one field at a time + start_clock = time_ns() + n = 1 # on disk + m = 0 # in memory + + times_in_memory = all_times[time_indices_in_memory] + + fts = FieldTimeSeries{LX, LY, Nothing}(preprocessing_grid, times_in_memory; + boundary_conditions, + backend = InMemory(), + path = jld2_filepath, + name = fts_name) + + # Re-compute data + new_data = ds[shortname][i₁:i₂, j₁:j₂, time_indices_in_memory] + + if !on_native_grid + copyto!(interior(native_fts, :, :, 1, :), new_data[:, :, :]) + fill_halo_regions!(native_fts) + interpolate!(fts, native_fts) + else + copyto!(interior(fts, :, :, 1, :), new_data[:, :, :]) + end + + while n <= all_Nt + print(" ... processing time index $n of $all_Nt \r") + + if time_indices_in_memory isa Colon || n ∈ time_indices_in_memory + m += 1 + else # load new data + # Update time_indices + time_indices_in_memory = time_indices_in_memory .+ preprocess_chunk_size + n₁ = first(time_indices_in_memory) + + # Clip time_indices if they extend past the end of the dataset + if last(time_indices_in_memory) > all_Nt + time_indices_in_memory = UnitRange(n₁, all_Nt) + end + + # Re-compute times + new_times = jra55_times(all_times[time_indices_in_memory], all_times[n₁]) + native_fts.times = new_times + + # Re-compute data + new_data = ds[shortname][i₁:i₂, j₁:j₂, time_indices_in_memory] + fts.times = new_times + + if !on_native_grid + copyto!(interior(native_fts, :, :, 1, :), new_data[:, :, :]) + fill_halo_regions!(native_fts) + interpolate!(fts, native_fts) + else + copyto!(interior(fts, :, :, 1, :), new_data[:, :, :]) + end + + m = 1 # reset + end + + set!(on_disk_fts, fts[m], n, fts.times[m]) + + n += 1 + end + + elapsed = 1e-9 * (time_ns() - start_clock) + elapsed_str = prettytime(elapsed) + @info " ... done ($elapsed_str)" * repeat(" ", 20) + + close(ds) + + user_fts = FieldTimeSeries(jld2_filepath, fts_name; architecture, backend, time_indexing) + fill_halo_regions!(user_fts) + + return user_fts +end + +const AA = Oceananigans.Architectures.AbstractArchitecture + +JRA55PrescribedAtmosphere(time_indices=Colon(); kw...) = + JRA55PrescribedAtmosphere(CPU(), time_indices; kw...) + +JRA55PrescribedAtmosphere(arch::Distributed, time_indices=Colon(); kw...) = + JRA55PrescribedAtmosphere(child_architecture(arch), time_indices; kw...) + +# TODO: allow the user to pass dates +""" + JRA55PrescribedAtmosphere(architecture::AA, time_indices=Colon(); + backend = nothing, + time_indexing = Cyclical(), + reference_height = 10, # meters + include_rivers_and_icebergs = false, + other_kw...) + +Return a `PrescribedAtmosphere` representing JRA55 reanalysis data. +""" +function JRA55PrescribedAtmosphere(architecture::AA, time_indices=Colon(); + backend = nothing, + time_indexing = Cyclical(), + reference_height = 10, # meters + include_rivers_and_icebergs = false, + other_kw...) + + if isnothing(backend) # apply a default + Ni = try + length(time_indices) + catch + Inf + end + + # Manufacture a default for the number of fields to keep InMemory + Nf = min(24, Ni) + backend = JRA55NetCDFBackend(Nf) + end + + kw = (; time_indices, time_indexing, backend, architecture) + kw = merge(kw, other_kw) + + ua = JRA55_field_time_series(:eastward_velocity; kw...) + va = JRA55_field_time_series(:northward_velocity; kw...) + Ta = JRA55_field_time_series(:temperature; kw...) + qa = JRA55_field_time_series(:specific_humidity; kw...) + pa = JRA55_field_time_series(:sea_level_pressure; kw...) + Fra = JRA55_field_time_series(:rain_freshwater_flux; kw...) + Fsn = JRA55_field_time_series(:snow_freshwater_flux; kw...) + Ql = JRA55_field_time_series(:downwelling_longwave_radiation; kw...) + Qs = JRA55_field_time_series(:downwelling_shortwave_radiation; kw...) + + # In JRA55, rivers and icebergs are on a different grid and stored with + # a different frequency than the rest of the data. Here, we use the + # `PrescribedAtmosphere.auxiliary_freshwater_flux` feature to represent them. + if include_rivers_and_icebergs + Fri = JRA55_field_time_series(:river_freshwater_flux; kw...) + Fic = JRA55_field_time_series(:iceberg_freshwater_flux; kw...) + auxiliary_freshwater_flux = (rivers=Fri, icebergs=Fic) + else + auxiliary_freshwater_flux = nothing + end + + times = ua.times + freshwater_flux = (rain=Fra, snow=Fsn) + velocities = (u=ua, v=va) + tracers = (T=Ta, q=qa) + pressure = pa + downwelling_radiation = TwoBandDownwellingRadiation(shortwave=Qs, longwave=Ql) + + FT = eltype(ua) + boundary_layer_height = convert(FT, 600) + reference_height = convert(FT, reference_height) + thermodynamics_parameters = PrescribedAtmosphereThermodynamicsParameters(FT) + grid = ua.grid + metadata = JRA55Data() + + return PrescribedAtmosphere(grid, + Clock{Float64}(time = 0), + metadata, + velocities, + pressure, + tracers, + freshwater_flux, + auxiliary_freshwater_flux, + downwelling_radiation, + thermodynamics_parameters, + times, + reference_height, + boundary_layer_height) +end + +end # module diff --git a/src/DistributedUtils.jl b/src/DistributedUtils.jl new file mode 100755 index 000000000..ecc49e2c6 --- /dev/null +++ b/src/DistributedUtils.jl @@ -0,0 +1,163 @@ +module DistributedUtils + +export mpi_initialized, mpi_rank, mpi_size, global_barrier, global_communicator +export @root, @onrank, @distribute, @handshake + +using MPI + +##### +##### Handle commands, typically downloading files +##### which should be executed by only one rank or distributed among ranks +##### + +# Utilities to make the macro work importing only ClimaOcean and not MPI +mpi_initialized() = MPI.Initialized() +mpi_rank(comm) = MPI.Comm_rank(comm) +mpi_size(comm) = MPI.Comm_size(comm) +global_barrier(comm) = MPI.Barrier(comm) +global_communicator() = MPI.COMM_WORLD + +""" + @root communicator exs... + +Perform `exs` only on rank 0 in communicator, otherwise known as the "root" rank. +Other ranks will wait for the root rank to finish before continuing. +If `communicator` is not provided, `MPI.COMM_WORLD` is used. +""" +macro root(communicator, exp) + command = quote + if ClimaOcean.mpi_initialized() + rank = ClimaOcean.mpi_rank($communicator) + if rank == 0 + $exp + end + ClimaOcean.global_barrier($communicator) + else + $exp + end + end + return esc(command) +end + +macro root(exp) + command = quote + @root ClimaOcean.global_communicator() $exp + end + return esc(command) +end + +""" + @onrank communicator rank exs... + +Perform `exp` only on rank `rank` (0-based index) in `communicator`. +Other ranks will wait for rank `rank` to finish before continuing. +The expression is run anyways if MPI in not initialized. +If `communicator` is not provided, `MPI.COMM_WORLD` is used. +""" +macro onrank(communicator, on_rank, exp) + command = quote + mpi_initialized = ClimaOcean.mpi_initialized() + if !mpi_initialized + $exp + else + rank = ClimaOcean.mpi_rank($communicator) + if rank == $on_rank + $exp + end + ClimaOcean.global_barrier($communicator) + end + end + + return esc(command) +end + +macro onrank(rank, exp) + command = quote + @onrank ClimaOcean.global_communicator() $rank $exp + end + return esc(command) +end + +""" + @distribute communicator for i in iterable + ... + end + +Distribute a `for` loop among different ranks in `communicator`. +If `communicator` is not provided, `MPI.COMM_WORLD` is used. +""" +macro distribute(communicator, exp) + if exp.head != :for + error("The `@distribute` macro expects a `for` loop") + end + + iterable = exp.args[1].args[2] + variable = exp.args[1].args[1] + forbody = exp.args[2] + + # Safety net if the iterable variable has the same name as the + # reserved variable names (nprocs, counter, rank) + nprocs = ifelse(variable == :nprocs, :othernprocs, :nprocs) + counter = ifelse(variable == :counter, :othercounter, :counter) + rank = ifelse(variable == :rank, :otherrank, :rank) + + new_loop = quote + mpi_initialized = ClimaOcean.mpi_initialized() + if !mpi_initialized + $exp + else + $rank = ClimaOcean.mpi_rank($communicator) + $nprocs = ClimaOcean.mpi_size($communicator) + for ($counter, $variable) in enumerate($iterable) + if ($counter - 1) % $nprocs == $rank + $forbody + end + end + ClimaOcean.global_barrier($communicator) + end + end + + return esc(new_loop) +end + +macro distribute(exp) + command = quote + @distribute ClimaOcean.global_communicator() $exp + end + return esc(command) +end + +""" + @handshake communicator exs... + +perform `exs` on all ranks in `communicator`, but only one rank at a time, where +ranks `r2 > r1` wait for rank `r1` to finish before executing `exs`. +If `communicator` is not provided, `MPI.COMM_WORLD` is used. +""" +macro handshake(communicator, exp) + command = quote + mpi_initialized = ClimaOcean.mpi_initialized() + if !mpi_initialized + $exp + else + rank = ClimaOcean.mpi_rank($communicator) + nprocs = ClimaOcean.mpi_size($communicator) + for r in 0 : nprocs -1 + if rank == r + $exp + end + ClimaOcean.global_barrier($communicator) + end + end + end + return esc(command) +end + +macro handshake(exp) + command = quote + @handshake ClimaOcean.global_communicator() $exp + end + return esc(command) +end + +end # module diff --git a/test/test_ocean_sea_ice_model_parameter_space.jl b/test/test_ocean_sea_ice_model_parameter_space.jl new file mode 100755 index 000000000..2655dedd6 --- /dev/null +++ b/test/test_ocean_sea_ice_model_parameter_space.jl @@ -0,0 +1,49 @@ +using Oceananigans +using ClimaOcean +using Oceananigans.OrthogonalSphericalShellGrids + +start_time = time_ns() +arch = GPU() +grid = TripolarGrid(arch; + size = (50, 50, 10), + halo = (7, 7, 7), + z = (-6000, 0), + first_pole_longitude = 75, + north_poles_latitude = 55) + +bottom_height = retrieve_bathymetry(grid; + minimum_depth = 10, + dir = "./", + interpolation_passes = 20, + connected_regions_allowed = 0) + +grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map = true) + +elapsed = 1e-9 * (time_ns() - start_time) +@info "Grid / bathymetry construction time: " * prettytime(elapsed) + +start_time = time_ns() +free_surface = SplitExplicitFreeSurface(grid; substeps = 20) +ocean = ocean_simulation(grid; free_surface) +model = ocean.model +@info "Ocean simulation construction time: " * prettytime(elapsed) + +start_time = time_ns() +backend = JRA55NetCDFBackend(4) +atmosphere = JRA55PrescribedAtmosphere(arch; backend) +radiation = Radiation(arch) + +elapsed = 1e-9 * (time_ns() - start_time) +@info "Atmosphere construction time: " * prettytime(elapsed) + +# Fluxes are computed when the model is constructed, so we just test that this works. +start_time = time_ns() +coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation) + +elapsed = 1e-9 * (time_ns() - start_time) +@info "Coupled model construction time: " * prettytime(elapsed) + +start_time = time_ns() +time_step!(coupled_model, 1) +elapsed = 1e-9 * (time_ns() - start_time) +@info "One time step time: " * prettytime(elapsed) diff --git a/test/test_simulations.jl b/test/test_simulations.jl new file mode 100755 index 000000000..a892b427b --- /dev/null +++ b/test/test_simulations.jl @@ -0,0 +1,75 @@ +include("runtests_setup.jl") + +using CUDA +using Oceananigans.OrthogonalSphericalShellGrids +using ClimaOcean.OceanSeaIceModels: above_freezing_ocean_temperature! +using ClimaSeaIce.SeaIceThermodynamics: melting_temperature + +@inline kernel_melting_temperature(i, j, k, grid, liquidus, S) = @inbounds melting_temperature(liquidus, S[i, j, k]) + +@testset "Time stepping test" begin + + for arch in test_architectures + + ##### + ##### Ocean and prescribed atmosphere + ##### + + grid = TripolarGrid(arch; + size = (50, 50, 10), + halo = (7, 7, 7), + z = (-6000, 0)) + + bottom_height = retrieve_bathymetry(grid; + minimum_depth = 10, + dir = "./", + interpolation_passes = 20, + major_basins = 1) + + grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map=true) + + free_surface = SplitExplicitFreeSurface(grid; substeps=20) + ocean = ocean_simulation(grid; free_surface) + + backend = JRA55NetCDFBackend(4) + atmosphere = JRA55PrescribedAtmosphere(arch; backend) + radiation = Radiation(arch) + + # Fluxes are computed when the model is constructed, so we just test that this works. + @test begin + coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation) + time_step!(coupled_model, 1) + true + end + + ##### + ##### Coupled ocean-sea ice and prescribed atmosphere + ##### + + sea_ice_grid = TripolarGrid(arch; size=(50, 50, 1), z = (-10, 0)) + sea_ice_grid = ImmersedBoundaryGrid(sea_ice_grid, GridFittedBottom(bottom_height)) + + sea_ice = sea_ice_simulation(sea_ice_grid) + liquidus = sea_ice.model.ice_thermodynamics.phase_transitions.liquidus + + # Set the ocean temperature and salinity + set!(ocean.model, T=temperature_metadata[1], S=salinity_metadata[1]) + above_freezing_ocean_temperature!(ocean, sea_ice) + + # Test that ocean temperatures are above freezing + T = on_architecture(CPU(), ocean.model.tracers.T) + S = on_architecture(CPU(), ocean.model.tracers.S) + + Tm = KernelFunctionOperation{Center, Center, Center}(kernel_melting_temperature, grid, liquidus, S) + @test all(T .>= Tm) + + # Fluxes are computed when the model is constructed, so we just test that this works. + # And that we can time step with sea ice + @test begin + coupled_model = OceanSeaIceModel(ocean, sea_ice; atmosphere, radiation) + time_step!(coupled_model, 1) + true + end + end +end + From 90abea4d90d48c5d54832e10408785389f14c7a5 Mon Sep 17 00:00:00 2001 From: Taimoor Sohail Date: Wed, 19 Mar 2025 13:41:25 +1100 Subject: [PATCH 34/54] Removed bottom line --- src/CoupledSimulation.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/src/CoupledSimulation.jl b/src/CoupledSimulation.jl index f654a5084..abd52ad6c 100755 --- a/src/CoupledSimulation.jl +++ b/src/CoupledSimulation.jl @@ -36,4 +36,3 @@ function run!(coupled_model::OceanSeaIceModel; schedule, return nothing end - From 9efd6b83c29e2d9f2a2b74492b9578c477ac568f Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 20 Mar 2025 10:21:25 +1100 Subject: [PATCH 35/54] try to generalize --- .../PrescribedAtmospheres.jl | 39 +++++------ src/OceanSeaIceModels/ocean_sea_ice_model.jl | 66 ++++++++++++------- 2 files changed, 60 insertions(+), 45 deletions(-) diff --git a/src/OceanSeaIceModels/PrescribedAtmospheres.jl b/src/OceanSeaIceModels/PrescribedAtmospheres.jl index 91c17cc5b..daa6b83e7 100644 --- a/src/OceanSeaIceModels/PrescribedAtmospheres.jl +++ b/src/OceanSeaIceModels/PrescribedAtmospheres.jl @@ -1,17 +1,19 @@ module PrescribedAtmospheres using Oceananigans.Grids: grid_name -using Oceananigans.Utils: prettysummary, Time +using Oceananigans.Utils: prettysummary, pretty_filesize, prettytime, Time using Oceananigans.Fields: Center using Oceananigans.OutputReaders: FieldTimeSeries, update_field_time_series!, extract_field_time_series +using Oceananigans.OutputWriters: Checkpointer, checkpoint_path, serializeproperties! using Oceananigans.TimeSteppers: Clock, tick! using Adapt using JLD2 using Thermodynamics.Parameters: AbstractThermodynamicsParameters +import Oceananigans: prognostic_fields import Oceananigans.Fields: set! -import Oceananigans.OutputWriters: checkpointer_address +import Oceananigans.OutputWriters: checkpointer_address, write_output!, initialize_jld2_file! import Oceananigans.TimeSteppers: time_step! import Thermodynamics.Parameters: @@ -321,33 +323,26 @@ function Base.show(io::IO, pa::PrescribedAtmosphere) print(io, "└── boundary_layer_height: ", prettysummary(pa.boundary_layer_height)) end -# set the clock to be the same as the ocean model -function set!(model::PrescribedAtmosphere, checkpoint_file_path) - addr = checkpointer_address(model) +checkpointer_address(::PrescribedAtmosphere) = "PrescribedAtmosphere" - jldopen(checkpoint_file_path, "r") do file - checkpointed_clock = file["$addr/clock"] - - # Update model clock - set_clock!(model, checkpointed_clock) - end - - return nothing +function default_checkpointed_properties(model::PrescribedAtmosphere) + properties = [:clock,] + return properties end -checkpointer_address(::PrescribedAtmosphere) = "HydrostaticFreeSurfaceModel" +prognostic_fields(model::PrescribedAtmosphere) = () """ - set_clock!(sim, clock) + set_clock!(model, clock) -Set the clock of `sim`ulation to match the values of `clock`. +Set the clock of a `model` to match the values of `clock`. """ -function set_clock!(sim::PrescribedAtmosphere, clock) - sim.clock.time = clock.time - sim.clock.iteration = clock.iteration - sim.clock.last_Δt = clock.last_Δt - sim.clock.last_stage_Δt = clock.last_stage_Δt - sim.clock.stage = clock.stage +function set_clock!(model::PrescribedAtmosphere, clock) + model.clock.time = clock.time + model.clock.iteration = clock.iteration + model.clock.last_Δt = clock.last_Δt + model.clock.last_stage_Δt = clock.last_stage_Δt + model.clock.stage = clock.stage return nothing end diff --git a/src/OceanSeaIceModels/ocean_sea_ice_model.jl b/src/OceanSeaIceModels/ocean_sea_ice_model.jl index ed138ec21..7d64ddf5a 100644 --- a/src/OceanSeaIceModels/ocean_sea_ice_model.jl +++ b/src/OceanSeaIceModels/ocean_sea_ice_model.jl @@ -13,7 +13,8 @@ import Oceananigans.Architectures: architecture import Oceananigans.Fields: set! import Oceananigans.Models: timestepper, NaNChecker, default_nan_checker, initialization_update_state! import Oceananigans.OutputWriters: default_included_properties, checkpointer_address, - write_output!, initialize_jld2_file! + write_output!, initialize_jld2_file!, + required_checkpointed_properties, default_checkpointed_properties import Oceananigans.Simulations: reset!, initialize!, iteration, run! import Oceananigans.TimeSteppers: time_step!, update_state!, time import Oceananigans.Utils: prettytime @@ -30,6 +31,7 @@ mutable struct OceanSeaIceModel{I, A, O, F, C, Arch} <: AbstractModel{Nothing, A end const OSIM = OceanSeaIceModel +const OSIMPA = OceanSeaIceModel{<:Any, <:PrescribedAtmosphere} function Base.summary(model::OSIM) A = nameof(typeof(architecture(model))) @@ -54,17 +56,18 @@ function Base.show(io::IO, cm::OSIM) end # Assumption: We have an ocean! -architecture(model::OSIM) = architecture(model.ocean.model) -Base.eltype(model::OSIM) = Base.eltype(model.ocean.model) -prettytime(model::OSIM) = prettytime(model.clock.time) -iteration(model::OSIM) = model.clock.iteration -timestepper(::OSIM) = nothing -default_included_properties(::OSIM) = tuple() -prognostic_fields(::OSIM) = nothing -fields(::OSIM) = NamedTuple() -default_clock(TT) = Oceananigans.TimeSteppers.Clock{TT}(0, 0, 1) -time(model::OSIM) = model.clock.time -checkpointer_address(::OSIM) = "HydrostaticFreeSurfaceModel" +architecture(model::OSIM) = architecture(model.ocean.model) +Base.eltype(model::OSIM) = Base.eltype(model.ocean.model) +prettytime(model::OSIM) = prettytime(model.clock.time) +iteration(model::OSIM) = model.clock.iteration +timestepper(::OSIM) = nothing +default_included_properties(::OSIM) = tuple() +prognostic_fields(::OSIM) = nothing +fields(::OSIM) = NamedTuple() +default_clock(TT) = Oceananigans.TimeSteppers.Clock{TT}(0, 0, 1) +time(model::OSIM) = model.clock.time +required_checkpointed_properties(::OSIM) = [:clock] +default_checkpointed_properties(::OSIM) = [:clock] reset!(model::OSIM) = reset!(model.ocean) @@ -81,10 +84,26 @@ function initialize!(model::OSIM) return nothing end -initialize_jld2_file!(filepath, init, jld2_kw, including, outputs, model::OSIM) = - initialize_jld2_file!(filepath, init, jld2_kw, including, outputs, model.ocean.model) +checkpointer_address(::OceanSeaIceModel) = "OceanSeaIceModel" -write_output!(c::Checkpointer, model::OSIM) = write_output!(c, model.ocean.model) +# initialize_jld2_file!(filepath, init, jld2_kw, including, outputs, model::OSIM) = +# initialize_jld2_file!(filepath, init, jld2_kw, including, outputs, model.ocean.model) + +# for prescribed atmosphere just checkpoint the ocean model +# write_output!(c::Checkpointer, model::OSIMPA) = write_output!(c, model.ocean.model) + +# function write_output!(c::Checkpointer, model::OSIMPA) +# atmosphere = model.atmosphere +# ocean = model.ocean.model + +# write_output!(c, model) # just saves the clock + +# # deals with model components +# write_output!(c, atmosphere) +# write_output!(c, ocean) + +# return nothing +# end function set_clock!(model::OSIM, clock) model.clock.time = clock.time @@ -95,17 +114,18 @@ function set_clock!(model::OSIM, clock) return nothing end -function set!(model::OSIM, checkpoint_file_path::AbstractString) - atmosphere = model.atmosphere - ocean = model.ocean.model +# function set!(model::OSIM, checkpoint_file_path::AbstractString) +# atmosphere = model.atmosphere +# ocean = model.ocean.model - set!(ocean, checkpoint_file_path) - set!(atmosphere, checkpoint_file_path) +# set!(model, checkpoint_file_path) - set_clock!(model, ocean.clock) +# # deals with model components +# set!(ocean, checkpoint_file_path) +# set!(atmosphere, checkpoint_file_path) - return nothing -end +# return nothing +# end reference_density(unsupported) = throw(ArgumentError("Cannot extract reference density from $(typeof(unsupported))")) From 1e8d638bcbc51e8d8d423b8aa26f2e4ea5eccdb6 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 20 Mar 2025 19:24:19 +1100 Subject: [PATCH 36/54] don't assume ocean component is special --- src/OceanSeaIceModels/OceanSeaIceModels.jl | 6 ++ .../PrescribedAtmospheres.jl | 25 +++++--- src/OceanSeaIceModels/ocean_sea_ice_model.jl | 15 ++--- src/OceanSeaIceModels/output_writers.jl | 61 +++++++++++++++++++ 4 files changed, 93 insertions(+), 14 deletions(-) create mode 100644 src/OceanSeaIceModels/output_writers.jl diff --git a/src/OceanSeaIceModels/OceanSeaIceModels.jl b/src/OceanSeaIceModels/OceanSeaIceModels.jl index 54d43ead0..170b805d7 100644 --- a/src/OceanSeaIceModels/OceanSeaIceModels.jl +++ b/src/OceanSeaIceModels/OceanSeaIceModels.jl @@ -26,6 +26,11 @@ using ClimaOcean: stateindex using KernelAbstractions: @kernel, @index using KernelAbstractions.Extras.LoopInfo: @unroll +import Oceananigans.OutputWriters: checkpointer_address, + required_checkpointed_properties, + default_checkpointed_properties, + prognostic_fields + function downwelling_radiation end function freshwater_flux end function reference_density end @@ -69,6 +74,7 @@ import .InterfaceComputations: compute_sea_ice_ocean_fluxes! include("ocean_sea_ice_model.jl") +include("output_writers.jl") include("freezing_limited_ocean_temperature.jl") include("time_step_ocean_sea_ice_model.jl") diff --git a/src/OceanSeaIceModels/PrescribedAtmospheres.jl b/src/OceanSeaIceModels/PrescribedAtmospheres.jl index daa6b83e7..2986a8a90 100644 --- a/src/OceanSeaIceModels/PrescribedAtmospheres.jl +++ b/src/OceanSeaIceModels/PrescribedAtmospheres.jl @@ -4,7 +4,6 @@ using Oceananigans.Grids: grid_name using Oceananigans.Utils: prettysummary, pretty_filesize, prettytime, Time using Oceananigans.Fields: Center using Oceananigans.OutputReaders: FieldTimeSeries, update_field_time_series!, extract_field_time_series -using Oceananigans.OutputWriters: Checkpointer, checkpoint_path, serializeproperties! using Oceananigans.TimeSteppers: Clock, tick! using Adapt @@ -13,8 +12,10 @@ using Thermodynamics.Parameters: AbstractThermodynamicsParameters import Oceananigans: prognostic_fields import Oceananigans.Fields: set! -import Oceananigans.OutputWriters: checkpointer_address, write_output!, initialize_jld2_file! import Oceananigans.TimeSteppers: time_step! +import Oceananigans.OutputWriters: checkpointer_address, + required_checkpointed_properties, + default_checkpointed_properties import Thermodynamics.Parameters: gas_constant, # @@ -323,14 +324,24 @@ function Base.show(io::IO, pa::PrescribedAtmosphere) print(io, "└── boundary_layer_height: ", prettysummary(pa.boundary_layer_height)) end -checkpointer_address(::PrescribedAtmosphere) = "PrescribedAtmosphere" +checkpointer_address(::PrescribedAtmosphere) = "PrescribedAtmosphere" +default_checkpointed_properties(::PrescribedAtmosphere) = [:clock] +required_checkpointed_properties(::PrescribedAtmosphere) = [:clock] +prognostic_fields(::PrescribedAtmosphere) = () -function default_checkpointed_properties(model::PrescribedAtmosphere) - properties = [:clock,] - return properties +function set!(model::PrescribedAtmosphere, checkpoint_file_path::AbstractString) + addr = checkpointer_address(model) + + jldopen(checkpoint_file_path, "r") do file + checkpointed_clock = file["$addr/clock"] + + # Update model clock + set_clock!(model, checkpointed_clock) + end + + return nothing end -prognostic_fields(model::PrescribedAtmosphere) = () """ set_clock!(model, clock) diff --git a/src/OceanSeaIceModels/ocean_sea_ice_model.jl b/src/OceanSeaIceModels/ocean_sea_ice_model.jl index 7d64ddf5a..d558911ac 100644 --- a/src/OceanSeaIceModels/ocean_sea_ice_model.jl +++ b/src/OceanSeaIceModels/ocean_sea_ice_model.jl @@ -8,13 +8,15 @@ using SeawaterPolynomials: TEOS10EquationOfState import Thermodynamics as AtmosphericThermodynamics # Simulations interface -import Oceananigans: fields, prognostic_fields +import Oceananigans: fields, + prognostic_fields import Oceananigans.Architectures: architecture import Oceananigans.Fields: set! import Oceananigans.Models: timestepper, NaNChecker, default_nan_checker, initialization_update_state! -import Oceananigans.OutputWriters: default_included_properties, checkpointer_address, - write_output!, initialize_jld2_file!, - required_checkpointed_properties, default_checkpointed_properties +import Oceananigans.OutputWriters: checkpointer_address, + required_checkpointed_properties, + default_checkpointed_properties + import Oceananigans.Simulations: reset!, initialize!, iteration, run! import Oceananigans.TimeSteppers: time_step!, update_state!, time import Oceananigans.Utils: prettytime @@ -62,12 +64,13 @@ prettytime(model::OSIM) = prettytime(model.clock.time) iteration(model::OSIM) = model.clock.iteration timestepper(::OSIM) = nothing default_included_properties(::OSIM) = tuple() -prognostic_fields(::OSIM) = nothing +prognostic_fields(::OSIM) = tuple() fields(::OSIM) = NamedTuple() default_clock(TT) = Oceananigans.TimeSteppers.Clock{TT}(0, 0, 1) time(model::OSIM) = model.clock.time required_checkpointed_properties(::OSIM) = [:clock] default_checkpointed_properties(::OSIM) = [:clock] +checkpointer_address(::OSIM) = "OceanSeaIceModel" reset!(model::OSIM) = reset!(model.ocean) @@ -84,8 +87,6 @@ function initialize!(model::OSIM) return nothing end -checkpointer_address(::OceanSeaIceModel) = "OceanSeaIceModel" - # initialize_jld2_file!(filepath, init, jld2_kw, including, outputs, model::OSIM) = # initialize_jld2_file!(filepath, init, jld2_kw, including, outputs, model.ocean.model) diff --git a/src/OceanSeaIceModels/output_writers.jl b/src/OceanSeaIceModels/output_writers.jl new file mode 100644 index 000000000..196023bb3 --- /dev/null +++ b/src/OceanSeaIceModels/output_writers.jl @@ -0,0 +1,61 @@ +using JLD2 + +using Oceananigans.Utils: pretty_filesize +using Oceananigans.OutputWriters: checkpoint_path, + cleanup_checkpoints, + validate_properties + +import Oceananigans.Fields: set! +import Oceananigans.OutputWriters: write_output! + +wall_time = Ref(time_ns()) + +function set!(model::OSIMPA, checkpoint_file_path::AbstractString) + addr = checkpointer_address(model) + + jldopen(checkpoint_file_path, "r") do file + checkpointed_clock = file["$addr/clock"] + + # Update model clock + set_clock!(model, checkpointed_clock) + end + + # deal with model components + + for component in (model.ocean.model, model.atmosphere) + set!(component, checkpoint_file_path) + end + + return nothing +end + + +function write_output!(c::Checkpointer, model::OSIM) + + filepath = checkpoint_path(model.clock.iteration, c) + c.verbose && @info "Checkpointing to file $filepath..." + + t1 = time_ns() + + # write OceanSeaIceModel + mode = "w" + properties = validate_properties(model, default_checkpointed_properties(model)) + write_output!(c, model, filepath, mode, properties) + + # write the OceanSeaIceModel components + components = (model.atmosphere, model.ocean.model) + + mode = "a" # to append in the file already written above + for component in components + component_properties = validate_properties(component, default_checkpointed_properties(component)) + write_output!(c, component, filepath, mode, component_properties) + end + + 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) + + return nothing +end From 5a69a5269e15909ff2bb8c04e8a3a4ff1831d214 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 20 Mar 2025 19:33:11 +1100 Subject: [PATCH 37/54] bump Oceanigans compat --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 4df20cbeb..509e981c3 100644 --- a/Project.toml +++ b/Project.toml @@ -42,7 +42,7 @@ JLD2 = "0.4, 0.5" KernelAbstractions = "0.9" MPI = "0.20" NCDatasets = "0.12, 0.13, 0.14" -Oceananigans = "0.96 - 0.99" +Oceananigans = "0.96.1 - 0.99" OffsetArrays = "1.14" PrecompileTools = "1" Scratch = "1" From a74499fd0927c628cfc88bef4e36bca4bbaa69e6 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 20 Mar 2025 19:33:23 +1100 Subject: [PATCH 38/54] remove commented code --- src/OceanSeaIceModels/ocean_sea_ice_model.jl | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/src/OceanSeaIceModels/ocean_sea_ice_model.jl b/src/OceanSeaIceModels/ocean_sea_ice_model.jl index d558911ac..d490b1890 100644 --- a/src/OceanSeaIceModels/ocean_sea_ice_model.jl +++ b/src/OceanSeaIceModels/ocean_sea_ice_model.jl @@ -87,25 +87,6 @@ function initialize!(model::OSIM) return nothing end -# initialize_jld2_file!(filepath, init, jld2_kw, including, outputs, model::OSIM) = -# initialize_jld2_file!(filepath, init, jld2_kw, including, outputs, model.ocean.model) - -# for prescribed atmosphere just checkpoint the ocean model -# write_output!(c::Checkpointer, model::OSIMPA) = write_output!(c, model.ocean.model) - -# function write_output!(c::Checkpointer, model::OSIMPA) -# atmosphere = model.atmosphere -# ocean = model.ocean.model - -# write_output!(c, model) # just saves the clock - -# # deals with model components -# write_output!(c, atmosphere) -# write_output!(c, ocean) - -# return nothing -# end - function set_clock!(model::OSIM, clock) model.clock.time = clock.time model.clock.iteration = clock.iteration From 0a46067aef520cb66d931b991bfc4da5c26ba4aa Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 20 Mar 2025 19:33:51 +1100 Subject: [PATCH 39/54] properties in write_output! is kwarg --- src/OceanSeaIceModels/output_writers.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/OceanSeaIceModels/output_writers.jl b/src/OceanSeaIceModels/output_writers.jl index 196023bb3..b9ed746a7 100644 --- a/src/OceanSeaIceModels/output_writers.jl +++ b/src/OceanSeaIceModels/output_writers.jl @@ -40,7 +40,7 @@ function write_output!(c::Checkpointer, model::OSIM) # write OceanSeaIceModel mode = "w" properties = validate_properties(model, default_checkpointed_properties(model)) - write_output!(c, model, filepath, mode, properties) + write_output!(c, model, filepath, mode; properties) # write the OceanSeaIceModel components components = (model.atmosphere, model.ocean.model) @@ -48,7 +48,7 @@ function write_output!(c::Checkpointer, model::OSIM) mode = "a" # to append in the file already written above for component in components component_properties = validate_properties(component, default_checkpointed_properties(component)) - write_output!(c, component, filepath, mode, component_properties) + write_output!(c, component, filepath, mode, properties = component_properties) end t2, sz = time_ns(), filesize(filepath) From 98ec18b5fb74d43d99ba0e15f6e212f041ff2b75 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 20 Mar 2025 19:35:56 +1100 Subject: [PATCH 40/54] undo changes --- src/DataWrangling/JRA55.jl | 697 +------------------------------------ 1 file changed, 5 insertions(+), 692 deletions(-) diff --git a/src/DataWrangling/JRA55.jl b/src/DataWrangling/JRA55.jl index 1fa655b73..2e631fe1a 100755 --- a/src/DataWrangling/JRA55.jl +++ b/src/DataWrangling/JRA55.jl @@ -1,9 +1,10 @@ module JRA55 +export JRA55FieldTimeSeries, JRA55PrescribedAtmosphere, JRA55RepeatYear + using Oceananigans using Oceananigans.Units -using Oceananigans: location using Oceananigans.Architectures: arch_array using Oceananigans.DistributedComputations using Oceananigans.DistributedComputations: child_architecture @@ -11,15 +12,11 @@ using Oceananigans.BoundaryConditions: fill_halo_regions! using Oceananigans.Grids: λnodes, φnodes, on_architecture using Oceananigans.Fields: interpolate! using Oceananigans.OutputReaders: Cyclical, TotallyInMemory, AbstractInMemoryBackend, FlavorOfFTS, time_indices -using Oceananigans.TimeSteppers: Clock using ClimaOcean -using ClimaOcean.DistributedUtils: @root -using ClimaOcean.DataWrangling: download_progress using ClimaOcean.OceanSeaIceModels: PrescribedAtmosphere, - PrescribedAtmosphereThermodynamicsParameters, TwoBandDownwellingRadiation using CUDA: @allowscalar @@ -33,692 +30,8 @@ import Oceananigans.Fields: set! import Oceananigans.OutputReaders: new_backend, update_field_time_series! using Downloads: download -download_jra55_cache::String = "" - -function __init__() - global download_jra55_cache = @get_scratch!("JRA55") -end - -struct JRA55Data end -const JRA55PrescribedAtmosphere = PrescribedAtmosphere{<:Any, <:JRA55Data} - -# A list of all variables provided in the JRA55 dataset: -JRA55_variable_names = (:river_freshwater_flux, - :rain_freshwater_flux, - :snow_freshwater_flux, - :iceberg_freshwater_flux, - :specific_humidity, - :sea_level_pressure, - :relative_humidity, - :downwelling_longwave_radiation, - :downwelling_shortwave_radiation, - :temperature, - :eastward_velocity, - :northward_velocity) - -filenames = Dict( - :river_freshwater_flux => "RYF.friver.1990_1991.nc", # Freshwater fluxes from rivers - :rain_freshwater_flux => "RYF.prra.1990_1991.nc", # Freshwater flux from rainfall - :snow_freshwater_flux => "RYF.prsn.1990_1991.nc", # Freshwater flux from snowfall - :iceberg_freshwater_flux => "RYF.licalvf.1990_1991.nc", # Freshwater flux from calving icebergs - :specific_humidity => "RYF.huss.1990_1991.nc", # Surface specific humidity - :sea_level_pressure => "RYF.psl.1990_1991.nc", # Sea level pressure - :relative_humidity => "RYF.rhuss.1990_1991.nc", # Surface relative humidity - :downwelling_longwave_radiation => "RYF.rlds.1990_1991.nc", # Downwelling longwave radiation - :downwelling_shortwave_radiation => "RYF.rsds.1990_1991.nc", # Downwelling shortwave radiation - :temperature => "RYF.tas.1990_1991.nc", # Near-surface air temperature - :eastward_velocity => "RYF.uas.1990_1991.nc", # Eastward near-surface wind - :northward_velocity => "RYF.vas.1990_1991.nc", # Northward near-surface wind -) - -jra55_short_names = Dict( - :river_freshwater_flux => "friver", # Freshwater fluxes from rivers - :rain_freshwater_flux => "prra", # Freshwater flux from rainfall - :snow_freshwater_flux => "prsn", # Freshwater flux from snowfall - :iceberg_freshwater_flux => "licalvf", # Freshwater flux from calving icebergs - :specific_humidity => "huss", # Surface specific humidity - :sea_level_pressure => "psl", # Sea level pressure - :relative_humidity => "rhuss", # Surface relative humidity - :downwelling_longwave_radiation => "rlds", # Downwelling longwave radiation - :downwelling_shortwave_radiation => "rsds", # Downwelling shortwave radiation - :temperature => "tas", # Near-surface air temperature - :eastward_velocity => "uas", # Eastward near-surface wind - :northward_velocity => "vas", # Northward near-surface wind -) - -field_time_series_short_names = Dict( - :river_freshwater_flux => "Fri", # Freshwater fluxes from rivers - :rain_freshwater_flux => "Fra", # Freshwater flux from rainfall - :snow_freshwater_flux => "Fsn", # Freshwater flux from snowfall - :iceberg_freshwater_flux => "Fic", # Freshwater flux from calving icebergs - :specific_humidity => "qa", # Surface specific humidity - :sea_level_pressure => "pa", # Sea level pressure - :relative_humidity => "rh", # Surface relative humidity - :downwelling_longwave_radiation => "Ql", # Downwelling longwave radiation - :downwelling_shortwave_radiation => "Qs", # Downwelling shortwave radiation - :temperature => "Ta", # Near-surface air temperature - :eastward_velocity => "ua", # Eastward near-surface wind - :northward_velocity => "va", # Northward near-surface wind -) - -urls = Dict( - :shortwave_radiation => "https://www.dropbox.com/scl/fi/z6fkvmd9oe3ycmaxta131/" * - "RYF.rsds.1990_1991.nc?rlkey=r7q6zcbj6a4fxsq0f8th7c4tc&dl=0", - - :river_freshwater_flux => "https://www.dropbox.com/scl/fi/21ggl4p74k4zvbf04nb67/" * - "RYF.friver.1990_1991.nc?rlkey=ny2qcjkk1cfijmwyqxsfm68fz&dl=0", - - :rain_freshwater_flux => "https://www.dropbox.com/scl/fi/5icl1gbd7f5hvyn656kjq/" * - "RYF.prra.1990_1991.nc?rlkey=iifyjm4ppwyd8ztcek4dtx0k8&dl=0", - - :snow_freshwater_flux => "https://www.dropbox.com/scl/fi/1r4ajjzb3643z93ads4x4/" * - "RYF.prsn.1990_1991.nc?rlkey=auyqpwn060cvy4w01a2yskfah&dl=0", - - :iceberg_freshwater_flux => "https://www.dropbox.com/scl/fi/44nc5y27ohvif7lkvpyv0/" * - "RYF.licalvf.1990_1991.nc?rlkey=w7rqu48y2baw1efmgrnmym0jk&dl=0", - - :specific_humidity => "https://www.dropbox.com/scl/fi/66z6ymfr4ghkynizydc29/" * - "RYF.huss.1990_1991.nc?rlkey=107yq04aew8lrmfyorj68v4td&dl=0", - - :sea_level_pressure => "https://www.dropbox.com/scl/fi/0fk332027oru1iiseykgp/" * - "RYF.psl.1990_1991.nc?rlkey=4xpr9uah741483aukok6d7ctt&dl=0", - - :relative_humidity => "https://www.dropbox.com/scl/fi/1agwsp0lzvntuyf8bm9la/" * - "RYF.rhuss.1990_1991.nc?rlkey=8cd0vs7iy1rw58b9pc9t68gtz&dl=0", - - :downwelling_longwave_radiation => "https://www.dropbox.com/scl/fi/y6r62szkirrivua5nqq61/" * - "RYF.rlds.1990_1991.nc?rlkey=wt9yq3cyrvs2rbowoirf4nkum&dl=0", - - :downwelling_shortwave_radiation => "https://www.dropbox.com/scl/fi/z6fkvmd9oe3ycmaxta131/" * - "RYF.rsds.1990_1991.nc?rlkey=r7q6zcbj6a4fxsq0f8th7c4tc&dl=0", - - :temperature => "https://www.dropbox.com/scl/fi/fpl0npwi476w635g6lke9/" * - "RYF.tas.1990_1991.nc?rlkey=0skb9pe6lgbfbiaoybe7m945s&dl=0", - - :eastward_velocity => "https://www.dropbox.com/scl/fi/86wetpqla2x97isp8092g/" * - "RYF.uas.1990_1991.nc?rlkey=rcaf18sh1yz0v9g4hjm1249j0&dl=0", - - :northward_velocity => "https://www.dropbox.com/scl/fi/d38sflo9ddljstd5jwgml/" * - "RYF.vas.1990_1991.nc?rlkey=f9y3e57kx8xrb40gbstarf0x6&dl=0", -) - -compute_bounding_nodes(::Nothing, ::Nothing, LH, hnodes) = nothing -compute_bounding_nodes(bounds, ::Nothing, LH, hnodes) = bounds - -function compute_bounding_nodes(x::Number, ::Nothing, LH, hnodes) - ϵ = convert(typeof(x), 0.001) # arbitrary? - return (x - ϵ, x + ϵ) -end - -# TODO: remove the allowscalar -function compute_bounding_nodes(::Nothing, grid, LH, hnodes) - hg = hnodes(grid, LH()) - h₁ = @allowscalar minimum(hg) - h₂ = @allowscalar maximum(hg) - return h₁, h₂ -end - -function compute_bounding_indices(::Nothing, hc) - Nh = length(hc) - return 1, Nh -end - -function compute_bounding_indices(bounds::Tuple, hc) - h₁, h₂ = bounds - Nh = length(hc) - - # The following should work. If ᵒ are the extrema of nodes we want to - # interpolate to, and the following is a sketch of the JRA55 native grid, - # - # 1 2 3 4 5 - # | | | | | | - # | x ᵒ | x | x | x ᵒ | x | - # | | | | | | - # 1 2 3 4 5 6 - # - # then for example, we should find that (iᵢ, i₂) = (1, 5). - # So we want to reduce the first index by one, and limit them - # both by the available data. There could be some mismatch due - # to the use of different coordinate systems (ie whether λ ∈ (0, 360) - # which we may also need to handle separately. - i₁ = searchsortedfirst(hc, h₁) - i₂ = searchsortedfirst(hc, h₂) - i₁ = max(1, i₁ - 1) - i₂ = min(Nh, i₂) - - return i₁, i₂ -end - -infer_longitudinal_topology(::Nothing) = Periodic - -function infer_longitudinal_topology(λbounds) - λ₁, λ₂ = λbounds - TX = λ₂ - λ₁ ≈ 360 ? Periodic : Bounded - return TX -end - -function compute_bounding_indices(longitude, latitude, grid, LX, LY, λc, φc) - λbounds = compute_bounding_nodes(longitude, grid, LX, λnodes) - φbounds = compute_bounding_nodes(latitude, grid, LY, φnodes) - - i₁, i₂ = compute_bounding_indices(λbounds, λc) - j₁, j₂ = compute_bounding_indices(φbounds, φc) - TX = infer_longitudinal_topology(λbounds) - - return i₁, i₂, j₁, j₂, TX -end - -# Convert dates to range until Oceananigans supports dates natively -function jra55_times(native_times, start_time=native_times[1]) - - times = [] - for native_time in native_times - time = native_time - start_time - time = Second(time).value - push!(times, time) - end - - return times -end - -struct JRA55NetCDFBackend <: AbstractInMemoryBackend{Int} - start :: Int - length :: Int -end - -""" - JRA55NetCDFBackend(length) - -Represents a JRA55 FieldTimeSeries backed by JRA55 native .nc files. -""" -JRA55NetCDFBackend(length) = JRA55NetCDFBackend(1, length) - -Base.length(backend::JRA55NetCDFBackend) = backend.length -Base.summary(backend::JRA55NetCDFBackend) = string("JRA55NetCDFBackend(", backend.start, ", ", backend.length, ")") - -const JRA55NetCDFFTS = FlavorOfFTS{<:Any, <:Any, <:Any, <:Any, <:JRA55NetCDFBackend} - -function set!(fts::JRA55NetCDFFTS, path::String=fts.path, name::String=fts.name) - - ds = Dataset(path) - - # Note that each file should have the variables - # - ds["time"]: time coordinate - # - ds["lon"]: longitude at the location of the variable - # - ds["lat"]: latitude at the location of the variable - # - ds["lon_bnds"]: bounding longitudes between which variables are averaged - # - ds["lat_bnds"]: bounding latitudes between which variables are averaged - # - ds[shortname]: the variable data - - # Nodes at the variable location - λc = ds["lon"][:] - φc = ds["lat"][:] - LX, LY, LZ = location(fts) - i₁, i₂, j₁, j₂, TX = compute_bounding_indices(nothing, nothing, fts.grid, LX, LY, λc, φc) - - nn = time_indices(fts) - nn = collect(nn) - - if issorted(nn) - data = ds[name][i₁:i₂, j₁:j₂, nn] - else - # The time indices may be cycling past 1; eg ti = [6, 7, 8, 1]. - # However, DiskArrays does not seem to support loading data with unsorted - # indices. So to handle this, we load the data in chunks, where each chunk's - # indices are sorted, and then glue the data together. - m = findfirst(n -> n == 1, nn) - n1 = nn[1:m-1] - n2 = nn[m:end] - - data1 = ds[name][i₁:i₂, j₁:j₂, n1] - data2 = ds[name][i₁:i₂, j₁:j₂, n2] - data = cat(data1, data2, dims=3) - end - - close(ds) - - copyto!(interior(fts, :, :, 1, :), data) - fill_halo_regions!(fts) - - return nothing -end - -new_backend(::JRA55NetCDFBackend, start, length) = JRA55NetCDFBackend(start, length) - -""" - JRA55_field_time_series(variable_name; - architecture = CPU(), - grid = nothing, - location = nothing, - url = nothing, - dir = download_jra55_cache, - filename = nothing, - shortname = nothing, - latitude = nothing, - longitude = nothing, - backend = InMemory(), - time_indexing = Cyclical(), - preprocess_chunk_size = 10, - preprocess_architecture = CPU(), - time_indices = nothing) - -Return a `FieldTimeSeries` containing atmospheric reanalysis data for `variable_name`, -which describes one of the variables in the "repeat year forcing" dataset derived -from the Japanese 55-year atmospheric reanalysis for driving ocean-sea ice models (JRA55-do). -For more information about the derivation of the repeat-year forcing dataset, see - -> Stewart et al. (2020). JRA55-do-based repeat year forcing datasets for driving ocean–sea-ice models, _Ocean Modelling_, **147**, 101557, https://doi.org/10.1016/j.ocemod.2019.101557. - -The `variable_name`s (and their `shortname`s used in NetCDF files) available from the JRA55-do are: -- `:river_freshwater_flux` ("friver") -- `:rain_freshwater_flux` ("prra") -- `:snow_freshwater_flux` ("prsn") -- `:iceberg_freshwater_flux` ("licalvf") -- `:specific_humidity` ("huss") -- `:sea_level_pressure` ("psl") -- `:relative_humidity` ("rhuss") -- `:downwelling_longwave_radiation` ("rlds") -- `:downwelling_shortwave_radiation` ("rsds") -- `:temperature` ("ras") -- `:eastward_velocity` ("uas") -- `:northward_velocity` ("vas") - -Keyword arguments -================= - -- `architecture`: Architecture for the `FieldTimeSeries`. - Default: CPU() - -- `time_indices`: Indices of the timeseries to extract from file. - For example, `time_indices=1:3` returns a - `FieldTimeSeries` with the first three time snapshots - of `variable_name`. - -- `latitude`: Guiding latitude bounds for the resulting grid. - Used to slice the data when loading into memory. - Default: nothing, which retains the latitude range of the native grid. - -- `longitude`: Guiding longitude bounds for the resulting grid. - Used to slice the data when loading into memory. - Default: nothing, which retains the longitude range of the native grid. - -- `url`: The url accessed to download the data for `variable_name`. - Default: `ClimaOcean.JRA55.urls[variable_name]`. - -- `filename`: The name of the downloaded file. - Default: `ClimaOcean.JRA55.filenames[variable_name]`. - -- `shortname`: The "short name" of `variable_name` inside its NetCDF file. - Default: `ClimaOcean.JRA55.jra55_short_names[variable_name]`. - -- `interpolated_file`: file holding an Oceananigans compatible version of the JRA55 data. - If it does not exist it will be generated. - -- `time_chunks_in_memory`: number of fields held in memory. If `nothing` then the whole timeseries - is loaded (not recommended). -""" -function JRA55_field_time_series(variable_name; - architecture = CPU(), - grid = nothing, - location = nothing, - url = nothing, - dir = download_jra55_cache, - filename = nothing, - shortname = nothing, - latitude = nothing, - longitude = nothing, - backend = InMemory(), - time_indexing = Cyclical(), - preprocess_chunk_size = 10, - preprocess_architecture = CPU(), - time_indices = nothing) - - # OnDisk backends do not support time interpolation! - # Disallow OnDisk for JRA55 dataset loading - if backend isa OnDisk - msg = string("We cannot load the JRA55 dataset with an `OnDisk` backend") - throw(ArgumentError(msg)) - end - - if isnothing(filename) && !(variable_name ∈ JRA55_variable_names) - variable_strs = Tuple(" - :$name \n" for name in JRA55_variable_names) - variables_msg = prod(variable_strs) - - msg = string("The variable :$variable_name is not provided by the JRA55-do dataset!", '\n', - "The variables provided by the JRA55-do dataset are:", '\n', - variables_msg) - - throw(ArgumentError(msg)) - end - - filepath = isnothing(filename) ? joinpath(dir, filenames[variable_name]) : joinpath(dir, filename) - - if !isnothing(filename) && !isfile(filepath) && isnothing(url) - throw(ArgumentError("A filename was provided without a url, but the file does not exist.\n \ - If intended, please provide both the filename and url that should be used \n \ - to download the new file.")) - end - - isnothing(filename) && (filename = filenames[variable_name]) - isnothing(shortname) && (shortname = jra55_short_names[variable_name]) - isnothing(url) && (url = urls[variable_name]) - - # Record some important user decisions - totally_in_memory = backend isa TotallyInMemory - on_native_grid = isnothing(grid) - !on_native_grid && backend isa JRA55NetCDFBackend && error("Can't use custom grid with JRA55NetCDFBackend.") - - jld2_filepath = joinpath(dir, string("JRA55_repeat_year_", variable_name, ".jld2")) - fts_name = field_time_series_short_names[variable_name] - - # Note, we don't re-use existing jld2 files. - @root begin - isfile(filepath) || download(url, filepath; progress=download_progress) - isfile(jld2_filepath) && rm(jld2_filepath) - end - - # Determine default time indices - if totally_in_memory - # In this case, the whole time series is in memory. - # Either the time series is short, or we are doing a limited-area - # simulation, like in a single column. So, we conservatively - # set a default `time_indices = 1:2`. - isnothing(time_indices) && (time_indices = 1:2) - time_indices_in_memory = time_indices - native_fts_architecture = architecture - else - # In this case, part or all of the time series will be stored in a file. - # Note: if the user has provided a grid, we will have to preprocess the - # .nc JRA55 data into a .jld2 file. In this case, `time_indices` refers - # to the time_indices that we will preprocess; - # by default we choose all of them. The architecture is only the - # architecture used for preprocessing, which typically will be CPU() - # even if we would like the final FieldTimeSeries on the GPU. - isnothing(time_indices) && (time_indices = :) - - if backend isa JRA55NetCDFBackend - time_indices_in_memory = 1:length(backend) - native_fts_architecture = architecture - else # then `time_indices_in_memory` refers to preprocessing - maximum_index = min(preprocess_chunk_size, length(time_indices)) - time_indices_in_memory = 1:maximum_index - native_fts_architecture = preprocess_architecture - end - end - - # Set a default location. - if isnothing(location) - LX = LY = Center - else - LX, LY = location - end - - ds = Dataset(filepath) - - # Note that each file should have the variables - # - ds["time"]: time coordinate - # - ds["lon"]: longitude at the location of the variable - # - ds["lat"]: latitude at the location of the variable - # - ds["lon_bnds"]: bounding longitudes between which variables are averaged - # - ds["lat_bnds"]: bounding latitudes between which variables are averaged - # - ds[shortname]: the variable data - - # Nodes at the variable location - λc = ds["lon"][:] - φc = ds["lat"][:] - - # Interfaces for the "native" JRA55 grid - λn = ds["lon_bnds"][1, :] - φn = ds["lat_bnds"][1, :] - - # The .nc coordinates lon_bnds and lat_bnds do not include - # the last interface, so we push them here. - push!(φn, 90) - push!(λn, λn[1] + 360) - - # TODO: support loading just part of the JRA55 data. - # Probably with arguments that take latitude, longitude bounds. - i₁, i₂, j₁, j₂, TX = compute_bounding_indices(longitude, latitude, grid, LX, LY, λc, φc) - - native_times = ds["time"][time_indices] - data = ds[shortname][i₁:i₂, j₁:j₂, time_indices_in_memory] - λr = λn[i₁:i₂+1] - φr = φn[j₁:j₂+1] - Nrx, Nry, Nt = size(data) - close(ds) - - N = (Nrx, Nry) - H = min.(N, (3, 3)) - - JRA55_native_grid = LatitudeLongitudeGrid(native_fts_architecture, Float32; - halo = H, - size = N, - longitude = λr, - latitude = φr, - topology = (TX, Bounded, Flat)) - - boundary_conditions = FieldBoundaryConditions(JRA55_native_grid, (Center, Center, Nothing)) - times = jra55_times(native_times) - - if backend isa JRA55NetCDFBackend - fts = FieldTimeSeries{Center, Center, Nothing}(JRA55_native_grid, times; - backend, - time_indexing, - boundary_conditions, - path = filepath, - name = shortname) - - # Fill the data in a GPU-friendly manner - copyto!(interior(fts, :, :, 1, :), data) - fill_halo_regions!(fts) - - return fts - else - # Make times into an array for later preprocessing - if !totally_in_memory - times = collect(times) - end - - native_fts = FieldTimeSeries{Center, Center, Nothing}(JRA55_native_grid, times; - time_indexing, - boundary_conditions) - - # Fill the data in a GPU-friendly manner - copyto!(interior(native_fts, :, :, 1, :), data) - fill_halo_regions!(native_fts) - - if on_native_grid && totally_in_memory - return native_fts - - elseif totally_in_memory # but not on the native grid! - boundary_conditions = FieldBoundaryConditions(grid, (LX, LY, Nothing)) - fts = FieldTimeSeries{LX, LY, Nothing}(grid, times; time_indexing, boundary_conditions) - interpolate!(fts, native_fts) - return fts - end - end - - @info "Pre-processing JRA55 $variable_name data into a JLD2 file..." - - preprocessing_grid = on_native_grid ? JRA55_native_grid : grid - - # Re-open the dataset! - ds = Dataset(filepath) - all_datetimes = ds["time"][time_indices] - all_Nt = length(all_datetimes) - - all_times = jra55_times(all_datetimes) - - on_disk_fts = FieldTimeSeries{LX, LY, Nothing}(preprocessing_grid, all_times; - boundary_conditions, - backend = OnDisk(), - path = jld2_filepath, - name = fts_name) - - # Save data to disk, one field at a time - start_clock = time_ns() - n = 1 # on disk - m = 0 # in memory - - times_in_memory = all_times[time_indices_in_memory] - - fts = FieldTimeSeries{LX, LY, Nothing}(preprocessing_grid, times_in_memory; - boundary_conditions, - backend = InMemory(), - path = jld2_filepath, - name = fts_name) - - # Re-compute data - new_data = ds[shortname][i₁:i₂, j₁:j₂, time_indices_in_memory] - - if !on_native_grid - copyto!(interior(native_fts, :, :, 1, :), new_data[:, :, :]) - fill_halo_regions!(native_fts) - interpolate!(fts, native_fts) - else - copyto!(interior(fts, :, :, 1, :), new_data[:, :, :]) - end - - while n <= all_Nt - print(" ... processing time index $n of $all_Nt \r") - - if time_indices_in_memory isa Colon || n ∈ time_indices_in_memory - m += 1 - else # load new data - # Update time_indices - time_indices_in_memory = time_indices_in_memory .+ preprocess_chunk_size - n₁ = first(time_indices_in_memory) - - # Clip time_indices if they extend past the end of the dataset - if last(time_indices_in_memory) > all_Nt - time_indices_in_memory = UnitRange(n₁, all_Nt) - end - - # Re-compute times - new_times = jra55_times(all_times[time_indices_in_memory], all_times[n₁]) - native_fts.times = new_times - - # Re-compute data - new_data = ds[shortname][i₁:i₂, j₁:j₂, time_indices_in_memory] - fts.times = new_times - - if !on_native_grid - copyto!(interior(native_fts, :, :, 1, :), new_data[:, :, :]) - fill_halo_regions!(native_fts) - interpolate!(fts, native_fts) - else - copyto!(interior(fts, :, :, 1, :), new_data[:, :, :]) - end - - m = 1 # reset - end - - set!(on_disk_fts, fts[m], n, fts.times[m]) - - n += 1 - end - - elapsed = 1e-9 * (time_ns() - start_clock) - elapsed_str = prettytime(elapsed) - @info " ... done ($elapsed_str)" * repeat(" ", 20) - - close(ds) - - user_fts = FieldTimeSeries(jld2_filepath, fts_name; architecture, backend, time_indexing) - fill_halo_regions!(user_fts) - - return user_fts -end - -const AA = Oceananigans.Architectures.AbstractArchitecture - -JRA55PrescribedAtmosphere(time_indices=Colon(); kw...) = - JRA55PrescribedAtmosphere(CPU(), time_indices; kw...) - -JRA55PrescribedAtmosphere(arch::Distributed, time_indices=Colon(); kw...) = - JRA55PrescribedAtmosphere(child_architecture(arch), time_indices; kw...) - -# TODO: allow the user to pass dates -""" - JRA55PrescribedAtmosphere(architecture::AA, time_indices=Colon(); - backend = nothing, - time_indexing = Cyclical(), - reference_height = 10, # meters - include_rivers_and_icebergs = false, - other_kw...) - -Return a `PrescribedAtmosphere` representing JRA55 reanalysis data. -""" -function JRA55PrescribedAtmosphere(architecture::AA, time_indices=Colon(); - backend = nothing, - time_indexing = Cyclical(), - reference_height = 10, # meters - include_rivers_and_icebergs = false, - other_kw...) - - if isnothing(backend) # apply a default - Ni = try - length(time_indices) - catch - Inf - end - - # Manufacture a default for the number of fields to keep InMemory - Nf = min(24, Ni) - backend = JRA55NetCDFBackend(Nf) - end - - kw = (; time_indices, time_indexing, backend, architecture) - kw = merge(kw, other_kw) - - ua = JRA55_field_time_series(:eastward_velocity; kw...) - va = JRA55_field_time_series(:northward_velocity; kw...) - Ta = JRA55_field_time_series(:temperature; kw...) - qa = JRA55_field_time_series(:specific_humidity; kw...) - pa = JRA55_field_time_series(:sea_level_pressure; kw...) - Fra = JRA55_field_time_series(:rain_freshwater_flux; kw...) - Fsn = JRA55_field_time_series(:snow_freshwater_flux; kw...) - Ql = JRA55_field_time_series(:downwelling_longwave_radiation; kw...) - Qs = JRA55_field_time_series(:downwelling_shortwave_radiation; kw...) - - # In JRA55, rivers and icebergs are on a different grid and stored with - # a different frequency than the rest of the data. Here, we use the - # `PrescribedAtmosphere.auxiliary_freshwater_flux` feature to represent them. - if include_rivers_and_icebergs - Fri = JRA55_field_time_series(:river_freshwater_flux; kw...) - Fic = JRA55_field_time_series(:iceberg_freshwater_flux; kw...) - auxiliary_freshwater_flux = (rivers=Fri, icebergs=Fic) - else - auxiliary_freshwater_flux = nothing - end - - times = ua.times - freshwater_flux = (rain=Fra, snow=Fsn) - velocities = (u=ua, v=va) - tracers = (T=Ta, q=qa) - pressure = pa - downwelling_radiation = TwoBandDownwellingRadiation(shortwave=Qs, longwave=Ql) - - FT = eltype(ua) - boundary_layer_height = convert(FT, 600) - reference_height = convert(FT, reference_height) - thermodynamics_parameters = PrescribedAtmosphereThermodynamicsParameters(FT) - grid = ua.grid - metadata = JRA55Data() - - return PrescribedAtmosphere(grid, - Clock{Float64}(time = 0), - metadata, - velocities, - pressure, - tracers, - freshwater_flux, - auxiliary_freshwater_flux, - downwelling_radiation, - thermodynamics_parameters, - times, - reference_height, - boundary_layer_height) -end +include("JRA55_metadata.jl") +include("JRA55_field_time_series.jl") +include("JRA55_prescribed_atmosphere.jl") end # module From f9e9e6053bb9988fbca26f6b0b18f80ea7806772 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 20 Mar 2025 19:36:58 +1100 Subject: [PATCH 41/54] Delete examples/generate_atmos_dataset.jl --- examples/generate_atmos_dataset.jl | 17 ----------------- 1 file changed, 17 deletions(-) delete mode 100755 examples/generate_atmos_dataset.jl diff --git a/examples/generate_atmos_dataset.jl b/examples/generate_atmos_dataset.jl deleted file mode 100755 index 8ba1d3233..000000000 --- a/examples/generate_atmos_dataset.jl +++ /dev/null @@ -1,17 +0,0 @@ -using ClimaOcean -using Oceananigans -using JLD2 - -time_indices = 1:1 - -qt = ClimaOcean.JRA55.JRA55_field_time_series(:specific_humidity; time_indices) -Tt = ClimaOcean.JRA55.JRA55_field_time_series(:temperature; time_indices) -pt = ClimaOcean.JRA55.JRA55_field_time_series(:sea_level_pressure; time_indices) - -Nx, Ny, Nz = size(qt[1]) - -q = Array(interior(qt[1], 1:4:Nx, 1:4:Ny, 1)) -T = Array(interior(Tt[1], 1:4:Nx, 1:4:Ny, 1)) -p = Array(interior(pt[1], 1:4:Nx, 1:4:Ny, 1)) - -@save "atmospheric_state.jld2" q T p From 39c347a7312d174a6dcca31d2c64a698e9e47843 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 20 Mar 2025 19:39:37 +1100 Subject: [PATCH 42/54] Delete src/CoupledSimulation.jl --- src/CoupledSimulation.jl | 38 -------------------------------------- 1 file changed, 38 deletions(-) delete mode 100755 src/CoupledSimulation.jl diff --git a/src/CoupledSimulation.jl b/src/CoupledSimulation.jl deleted file mode 100755 index abd52ad6c..000000000 --- a/src/CoupledSimulation.jl +++ /dev/null @@ -1,38 +0,0 @@ -module Simulations - -import Oceananigans: run! - -function run!(coupled_model::OceanSeaIceModel; schedule, - dir = ".", - prefix = "checkpoint", - overwrite_existing = false, - verbose = false, - cleanup = false, - properties = default_checkpointed_properties(coupled_model.ocean.model)) - - function run!(sim; pickup=false) - - start_run = time_ns() - - if we_want_to_pickup(pickup) - set!(sim, pickup) - end - - sim.initialized = false - sim.running = true - sim.run_wall_time = 0.0 - - while sim.running - time_step!(sim) - end - - for callback in values(sim.callbacks) - finalize!(callback, sim) - end - - # Increment the wall clock - end_run = time_ns() - sim.run_wall_time += 1e-9 * (end_run - start_run) - - return nothing - end From 5281eab2b1de53c517355b87e437f27a637aa23a Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 20 Mar 2025 19:40:17 +1100 Subject: [PATCH 43/54] Delete test/test_ocean_sea_ice_model_parameter_space.jl --- ...est_ocean_sea_ice_model_parameter_space.jl | 49 ------------------- 1 file changed, 49 deletions(-) delete mode 100755 test/test_ocean_sea_ice_model_parameter_space.jl diff --git a/test/test_ocean_sea_ice_model_parameter_space.jl b/test/test_ocean_sea_ice_model_parameter_space.jl deleted file mode 100755 index 2655dedd6..000000000 --- a/test/test_ocean_sea_ice_model_parameter_space.jl +++ /dev/null @@ -1,49 +0,0 @@ -using Oceananigans -using ClimaOcean -using Oceananigans.OrthogonalSphericalShellGrids - -start_time = time_ns() -arch = GPU() -grid = TripolarGrid(arch; - size = (50, 50, 10), - halo = (7, 7, 7), - z = (-6000, 0), - first_pole_longitude = 75, - north_poles_latitude = 55) - -bottom_height = retrieve_bathymetry(grid; - minimum_depth = 10, - dir = "./", - interpolation_passes = 20, - connected_regions_allowed = 0) - -grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map = true) - -elapsed = 1e-9 * (time_ns() - start_time) -@info "Grid / bathymetry construction time: " * prettytime(elapsed) - -start_time = time_ns() -free_surface = SplitExplicitFreeSurface(grid; substeps = 20) -ocean = ocean_simulation(grid; free_surface) -model = ocean.model -@info "Ocean simulation construction time: " * prettytime(elapsed) - -start_time = time_ns() -backend = JRA55NetCDFBackend(4) -atmosphere = JRA55PrescribedAtmosphere(arch; backend) -radiation = Radiation(arch) - -elapsed = 1e-9 * (time_ns() - start_time) -@info "Atmosphere construction time: " * prettytime(elapsed) - -# Fluxes are computed when the model is constructed, so we just test that this works. -start_time = time_ns() -coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation) - -elapsed = 1e-9 * (time_ns() - start_time) -@info "Coupled model construction time: " * prettytime(elapsed) - -start_time = time_ns() -time_step!(coupled_model, 1) -elapsed = 1e-9 * (time_ns() - start_time) -@info "One time step time: " * prettytime(elapsed) From 368358e72fd296b9cedd681abb95bf2a5e83e6cf Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 20 Mar 2025 19:40:33 +1100 Subject: [PATCH 44/54] Delete test/test_simulations.jl --- test/test_simulations.jl | 75 ---------------------------------------- 1 file changed, 75 deletions(-) delete mode 100755 test/test_simulations.jl diff --git a/test/test_simulations.jl b/test/test_simulations.jl deleted file mode 100755 index a892b427b..000000000 --- a/test/test_simulations.jl +++ /dev/null @@ -1,75 +0,0 @@ -include("runtests_setup.jl") - -using CUDA -using Oceananigans.OrthogonalSphericalShellGrids -using ClimaOcean.OceanSeaIceModels: above_freezing_ocean_temperature! -using ClimaSeaIce.SeaIceThermodynamics: melting_temperature - -@inline kernel_melting_temperature(i, j, k, grid, liquidus, S) = @inbounds melting_temperature(liquidus, S[i, j, k]) - -@testset "Time stepping test" begin - - for arch in test_architectures - - ##### - ##### Ocean and prescribed atmosphere - ##### - - grid = TripolarGrid(arch; - size = (50, 50, 10), - halo = (7, 7, 7), - z = (-6000, 0)) - - bottom_height = retrieve_bathymetry(grid; - minimum_depth = 10, - dir = "./", - interpolation_passes = 20, - major_basins = 1) - - grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map=true) - - free_surface = SplitExplicitFreeSurface(grid; substeps=20) - ocean = ocean_simulation(grid; free_surface) - - backend = JRA55NetCDFBackend(4) - atmosphere = JRA55PrescribedAtmosphere(arch; backend) - radiation = Radiation(arch) - - # Fluxes are computed when the model is constructed, so we just test that this works. - @test begin - coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation) - time_step!(coupled_model, 1) - true - end - - ##### - ##### Coupled ocean-sea ice and prescribed atmosphere - ##### - - sea_ice_grid = TripolarGrid(arch; size=(50, 50, 1), z = (-10, 0)) - sea_ice_grid = ImmersedBoundaryGrid(sea_ice_grid, GridFittedBottom(bottom_height)) - - sea_ice = sea_ice_simulation(sea_ice_grid) - liquidus = sea_ice.model.ice_thermodynamics.phase_transitions.liquidus - - # Set the ocean temperature and salinity - set!(ocean.model, T=temperature_metadata[1], S=salinity_metadata[1]) - above_freezing_ocean_temperature!(ocean, sea_ice) - - # Test that ocean temperatures are above freezing - T = on_architecture(CPU(), ocean.model.tracers.T) - S = on_architecture(CPU(), ocean.model.tracers.S) - - Tm = KernelFunctionOperation{Center, Center, Center}(kernel_melting_temperature, grid, liquidus, S) - @test all(T .>= Tm) - - # Fluxes are computed when the model is constructed, so we just test that this works. - # And that we can time step with sea ice - @test begin - coupled_model = OceanSeaIceModel(ocean, sea_ice; atmosphere, radiation) - time_step!(coupled_model, 1) - true - end - end -end - From 021c306d8681fe9a42fdecd3b4289fedc022b179 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 20 Mar 2025 19:41:34 +1100 Subject: [PATCH 45/54] Delete src/DataWrangling/JRA55.jl --- src/DataWrangling/JRA55.jl | 37 ------------------------------------- 1 file changed, 37 deletions(-) delete mode 100755 src/DataWrangling/JRA55.jl diff --git a/src/DataWrangling/JRA55.jl b/src/DataWrangling/JRA55.jl deleted file mode 100755 index 2e631fe1a..000000000 --- a/src/DataWrangling/JRA55.jl +++ /dev/null @@ -1,37 +0,0 @@ -module JRA55 - -export JRA55FieldTimeSeries, JRA55PrescribedAtmosphere, JRA55RepeatYear - -using Oceananigans -using Oceananigans.Units - -using Oceananigans.Architectures: arch_array -using Oceananigans.DistributedComputations -using Oceananigans.DistributedComputations: child_architecture -using Oceananigans.BoundaryConditions: fill_halo_regions! -using Oceananigans.Grids: λnodes, φnodes, on_architecture -using Oceananigans.Fields: interpolate! -using Oceananigans.OutputReaders: Cyclical, TotallyInMemory, AbstractInMemoryBackend, FlavorOfFTS, time_indices - -using ClimaOcean - -using ClimaOcean.OceanSeaIceModels: - PrescribedAtmosphere, - TwoBandDownwellingRadiation - -using CUDA: @allowscalar - -using NCDatasets -using JLD2 -using Dates -using Scratch - -import Oceananigans.Fields: set! -import Oceananigans.OutputReaders: new_backend, update_field_time_series! -using Downloads: download - -include("JRA55_metadata.jl") -include("JRA55_field_time_series.jl") -include("JRA55_prescribed_atmosphere.jl") - -end # module From 48dd1b16606c1bdd8eae49f69e9b180654979ca6 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 20 Mar 2025 19:41:53 +1100 Subject: [PATCH 46/54] Delete src/DistributedUtils.jl --- src/DistributedUtils.jl | 163 ---------------------------------------- 1 file changed, 163 deletions(-) delete mode 100755 src/DistributedUtils.jl diff --git a/src/DistributedUtils.jl b/src/DistributedUtils.jl deleted file mode 100755 index ecc49e2c6..000000000 --- a/src/DistributedUtils.jl +++ /dev/null @@ -1,163 +0,0 @@ -module DistributedUtils - -export mpi_initialized, mpi_rank, mpi_size, global_barrier, global_communicator -export @root, @onrank, @distribute, @handshake - -using MPI - -##### -##### Handle commands, typically downloading files -##### which should be executed by only one rank or distributed among ranks -##### - -# Utilities to make the macro work importing only ClimaOcean and not MPI -mpi_initialized() = MPI.Initialized() -mpi_rank(comm) = MPI.Comm_rank(comm) -mpi_size(comm) = MPI.Comm_size(comm) -global_barrier(comm) = MPI.Barrier(comm) -global_communicator() = MPI.COMM_WORLD - -""" - @root communicator exs... - -Perform `exs` only on rank 0 in communicator, otherwise known as the "root" rank. -Other ranks will wait for the root rank to finish before continuing. -If `communicator` is not provided, `MPI.COMM_WORLD` is used. -""" -macro root(communicator, exp) - command = quote - if ClimaOcean.mpi_initialized() - rank = ClimaOcean.mpi_rank($communicator) - if rank == 0 - $exp - end - ClimaOcean.global_barrier($communicator) - else - $exp - end - end - return esc(command) -end - -macro root(exp) - command = quote - @root ClimaOcean.global_communicator() $exp - end - return esc(command) -end - -""" - @onrank communicator rank exs... - -Perform `exp` only on rank `rank` (0-based index) in `communicator`. -Other ranks will wait for rank `rank` to finish before continuing. -The expression is run anyways if MPI in not initialized. -If `communicator` is not provided, `MPI.COMM_WORLD` is used. -""" -macro onrank(communicator, on_rank, exp) - command = quote - mpi_initialized = ClimaOcean.mpi_initialized() - if !mpi_initialized - $exp - else - rank = ClimaOcean.mpi_rank($communicator) - if rank == $on_rank - $exp - end - ClimaOcean.global_barrier($communicator) - end - end - - return esc(command) -end - -macro onrank(rank, exp) - command = quote - @onrank ClimaOcean.global_communicator() $rank $exp - end - return esc(command) -end - -""" - @distribute communicator for i in iterable - ... - end - -Distribute a `for` loop among different ranks in `communicator`. -If `communicator` is not provided, `MPI.COMM_WORLD` is used. -""" -macro distribute(communicator, exp) - if exp.head != :for - error("The `@distribute` macro expects a `for` loop") - end - - iterable = exp.args[1].args[2] - variable = exp.args[1].args[1] - forbody = exp.args[2] - - # Safety net if the iterable variable has the same name as the - # reserved variable names (nprocs, counter, rank) - nprocs = ifelse(variable == :nprocs, :othernprocs, :nprocs) - counter = ifelse(variable == :counter, :othercounter, :counter) - rank = ifelse(variable == :rank, :otherrank, :rank) - - new_loop = quote - mpi_initialized = ClimaOcean.mpi_initialized() - if !mpi_initialized - $exp - else - $rank = ClimaOcean.mpi_rank($communicator) - $nprocs = ClimaOcean.mpi_size($communicator) - for ($counter, $variable) in enumerate($iterable) - if ($counter - 1) % $nprocs == $rank - $forbody - end - end - ClimaOcean.global_barrier($communicator) - end - end - - return esc(new_loop) -end - -macro distribute(exp) - command = quote - @distribute ClimaOcean.global_communicator() $exp - end - return esc(command) -end - -""" - @handshake communicator exs... - -perform `exs` on all ranks in `communicator`, but only one rank at a time, where -ranks `r2 > r1` wait for rank `r1` to finish before executing `exs`. -If `communicator` is not provided, `MPI.COMM_WORLD` is used. -""" -macro handshake(communicator, exp) - command = quote - mpi_initialized = ClimaOcean.mpi_initialized() - if !mpi_initialized - $exp - else - rank = ClimaOcean.mpi_rank($communicator) - nprocs = ClimaOcean.mpi_size($communicator) - for r in 0 : nprocs -1 - if rank == r - $exp - end - ClimaOcean.global_barrier($communicator) - end - end - end - return esc(command) -end - -macro handshake(exp) - command = quote - @handshake ClimaOcean.global_communicator() $exp - end - return esc(command) -end - -end # module From fdf45d77ff180d9a7534d449d30f57777f396c98 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 20 Mar 2025 19:44:26 +1100 Subject: [PATCH 47/54] don't import things we don't need --- src/OceanSeaIceModels/PrescribedAtmospheres.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/OceanSeaIceModels/PrescribedAtmospheres.jl b/src/OceanSeaIceModels/PrescribedAtmospheres.jl index 2986a8a90..ea0c3b824 100644 --- a/src/OceanSeaIceModels/PrescribedAtmospheres.jl +++ b/src/OceanSeaIceModels/PrescribedAtmospheres.jl @@ -1,7 +1,7 @@ module PrescribedAtmospheres using Oceananigans.Grids: grid_name -using Oceananigans.Utils: prettysummary, pretty_filesize, prettytime, Time +using Oceananigans.Utils: prettysummary, Time using Oceananigans.Fields: Center using Oceananigans.OutputReaders: FieldTimeSeries, update_field_time_series!, extract_field_time_series using Oceananigans.TimeSteppers: Clock, tick! From 875d04592740b4d41246f253f8eb5ad34cd20b98 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 20 Mar 2025 19:47:32 +1100 Subject: [PATCH 48/54] cleanup --- src/OceanSeaIceModels/ocean_sea_ice_model.jl | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/src/OceanSeaIceModels/ocean_sea_ice_model.jl b/src/OceanSeaIceModels/ocean_sea_ice_model.jl index d490b1890..aa9b9ec9d 100644 --- a/src/OceanSeaIceModels/ocean_sea_ice_model.jl +++ b/src/OceanSeaIceModels/ocean_sea_ice_model.jl @@ -96,19 +96,6 @@ function set_clock!(model::OSIM, clock) return nothing end -# function set!(model::OSIM, checkpoint_file_path::AbstractString) -# atmosphere = model.atmosphere -# ocean = model.ocean.model - -# set!(model, checkpoint_file_path) - -# # deals with model components -# set!(ocean, checkpoint_file_path) -# set!(atmosphere, checkpoint_file_path) - -# return nothing -# end - reference_density(unsupported) = throw(ArgumentError("Cannot extract reference density from $(typeof(unsupported))")) From 1e34bb526f60c6c0bd64847103119f628371b07d Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 20 Mar 2025 19:56:57 +1100 Subject: [PATCH 49/54] validate_properties -> validate_checkpointed_properties --- src/OceanSeaIceModels/output_writers.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/OceanSeaIceModels/output_writers.jl b/src/OceanSeaIceModels/output_writers.jl index b9ed746a7..d3836586f 100644 --- a/src/OceanSeaIceModels/output_writers.jl +++ b/src/OceanSeaIceModels/output_writers.jl @@ -3,7 +3,7 @@ using JLD2 using Oceananigans.Utils: pretty_filesize using Oceananigans.OutputWriters: checkpoint_path, cleanup_checkpoints, - validate_properties + validate_checkpointed_properties import Oceananigans.Fields: set! import Oceananigans.OutputWriters: write_output! @@ -39,7 +39,7 @@ function write_output!(c::Checkpointer, model::OSIM) # write OceanSeaIceModel mode = "w" - properties = validate_properties(model, default_checkpointed_properties(model)) + properties = validate_checkpointed_properties(model, default_checkpointed_properties(model)) write_output!(c, model, filepath, mode; properties) # write the OceanSeaIceModel components @@ -47,7 +47,7 @@ function write_output!(c::Checkpointer, model::OSIM) mode = "a" # to append in the file already written above for component in components - component_properties = validate_properties(component, default_checkpointed_properties(component)) + component_properties = validate_checkpointed_properties(component, default_checkpointed_properties(component)) write_output!(c, component, filepath, mode, properties = component_properties) end From 1f3a46be91ccbd1e198d445eeb790d1a744df774 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 20 Mar 2025 20:05:02 +1100 Subject: [PATCH 50/54] no need to duplicate validation --- src/OceanSeaIceModels/output_writers.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/OceanSeaIceModels/output_writers.jl b/src/OceanSeaIceModels/output_writers.jl index d3836586f..06e4b1550 100644 --- a/src/OceanSeaIceModels/output_writers.jl +++ b/src/OceanSeaIceModels/output_writers.jl @@ -39,16 +39,16 @@ function write_output!(c::Checkpointer, model::OSIM) # write OceanSeaIceModel mode = "w" - properties = validate_checkpointed_properties(model, default_checkpointed_properties(model)) - write_output!(c, model, filepath, mode; properties) + write_output!(c, model, filepath, mode, + properties = default_checkpointed_properties(model)) # write the OceanSeaIceModel components components = (model.atmosphere, model.ocean.model) mode = "a" # to append in the file already written above for component in components - component_properties = validate_checkpointed_properties(component, default_checkpointed_properties(component)) - write_output!(c, component, filepath, mode, properties = component_properties) + write_output!(c, component, filepath, mode, + properties = default_checkpointed_properties(component)) end t2, sz = time_ns(), filesize(filepath) From 4f9425a2e9539385421aec15916c6589cac9bca1 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Wed, 26 Mar 2025 10:12:13 +1100 Subject: [PATCH 51/54] fix initialize! and update_state! + add set_clock! --- .../PrescribedAtmospheres.jl | 39 +++++++++---------- 1 file changed, 18 insertions(+), 21 deletions(-) diff --git a/src/OceanSeaIceModels/PrescribedAtmospheres.jl b/src/OceanSeaIceModels/PrescribedAtmospheres.jl index ea0c3b824..405bd527b 100644 --- a/src/OceanSeaIceModels/PrescribedAtmospheres.jl +++ b/src/OceanSeaIceModels/PrescribedAtmospheres.jl @@ -10,9 +10,9 @@ using Adapt using JLD2 using Thermodynamics.Parameters: AbstractThermodynamicsParameters -import Oceananigans: prognostic_fields +import Oceananigans: initialize!, prognostic_fields import Oceananigans.Fields: set! -import Oceananigans.TimeSteppers: time_step! +import Oceananigans.TimeSteppers: time_step!, set_clock! import Oceananigans.OutputWriters: checkpointer_address, required_checkpointed_properties, default_checkpointed_properties @@ -332,31 +332,21 @@ prognostic_fields(::PrescribedAtmosphere) = () function set!(model::PrescribedAtmosphere, checkpoint_file_path::AbstractString) addr = checkpointer_address(model) + checkpointed_clock = nothing jldopen(checkpoint_file_path, "r") do file checkpointed_clock = file["$addr/clock"] - - # Update model clock - set_clock!(model, checkpointed_clock) end - return nothing -end - - -""" - set_clock!(model, clock) + # Update model clock + set_clock!(model, checkpointed_clock) -Set the clock of a `model` to match the values of `clock`. -""" -function set_clock!(model::PrescribedAtmosphere, clock) - model.clock.time = clock.time - model.clock.iteration = clock.iteration - model.clock.last_Δt = clock.last_Δt - model.clock.last_stage_Δt = clock.last_stage_Δt - model.clock.stage = clock.stage + @info "hi, I set the clock of the atmosphere" return nothing end +set_clock!(atmos::PrescribedAtmosphere, new_clock) = + set_clock!(atmos.clock, new_clock) + function default_atmosphere_velocities(grid, times) ua = FieldTimeSeries{Center, Center, Nothing}(grid, times) va = FieldTimeSeries{Center, Center, Nothing}(grid, times) @@ -366,7 +356,7 @@ end function default_atmosphere_tracers(grid, times) Ta = FieldTimeSeries{Center, Center, Nothing}(grid, times) qa = FieldTimeSeries{Center, Center, Nothing}(grid, times) - parent(Ta) .= 273.15 + 20 + parent(Ta) .= 273.15 + 20 # ᵒK return (T=Ta, q=qa) end @@ -384,13 +374,20 @@ end function default_atmosphere_pressure(grid, times) pa = FieldTimeSeries{Center, Center, Nothing}(grid, times) - parent(pa) .= 101325 + parent(pa) .= 101325 # Pa return pa end @inline function time_step!(atmos::PrescribedAtmosphere, Δt) tick!(atmos.clock, Δt) + update_state!(atmos) + + return nothing +end + +initialize!(atmos::PrescribedAtmosphere) = update_state!(atmos) +function update_state!(atmos::PrescribedAtmosphere) time = Time(atmos.clock.time) ftses = extract_field_time_series(atmos) From bf96083a4cf03575f615e30f60ebdb7281699bc1 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Wed, 26 Mar 2025 10:13:15 +1100 Subject: [PATCH 52/54] reorganize imports --- test/runtests_setup.jl | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/test/runtests_setup.jl b/test/runtests_setup.jl index b9f298124..2948b5313 100644 --- a/test/runtests_setup.jl +++ b/test/runtests_setup.jl @@ -1,8 +1,12 @@ using ClimaOcean using Oceananigans + using CUDA +using CFTime +using Dates using Test +using ClimaOcean.Bathymetry: download_bathymetry_cache using ClimaOcean.DataWrangling using ClimaOcean.ECCO using ClimaOcean.JRA55 @@ -10,11 +14,6 @@ using ClimaOcean.JRA55 using Oceananigans.Architectures: architecture, on_architecture using Oceananigans.OutputReaders: interpolate! -using ClimaOcean -using ClimaOcean.Bathymetry: download_bathymetry_cache -using CFTime -using Dates - using CUDA: @allowscalar gpu_test = parse(Bool, get(ENV, "GPU_TEST", "false")) @@ -30,11 +29,11 @@ temperature_metadata = Metadata(:temperature; dates, dataset=ECCO4Monthly()) salinity_metadata = Metadata(:salinity; dates, dataset=ECCO4Monthly()) # Fictitious grid that triggers bathymetry download -function download_bathymetry(; dir = download_bathymetry_cache, +function download_bathymetry(; dir = download_bathymetry_cache, filename = "ETOPO_2022_v1_60s_N90W180_surface.nc") - - grid = LatitudeLongitudeGrid(size = (10, 10, 1), - longitude = (0, 100), + + grid = LatitudeLongitudeGrid(size = (10, 10, 1), + longitude = (0, 100), latitude = (0, 50), z = (-6000, 0)) @@ -46,4 +45,3 @@ end # Trigger downloading JRA55 arch = first(test_architectures) atmosphere = JRA55PrescribedAtmosphere(arch; backend=JRA55NetCDFBackend(41)) - From cd7f0c5b48e260dee6a042c66981a8c48f131eb8 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Fri, 25 Apr 2025 09:02:16 +1000 Subject: [PATCH 53/54] wip --- src/ClimaOcean.jl | 11 +- .../component_interfaces.jl | 15 +- src/OceanSeaIceModels/ocean_sea_ice_model.jl | 20 +- src/OceanSeaIceModels/output_writers.jl | 17 +- src/OceanSimulations/ocean_simulation.jl | 23 +- test/runtests.jl | 12 +- test/test_ocean_sea_ice_model.jl | 211 +++++++++++++++--- 7 files changed, 236 insertions(+), 73 deletions(-) diff --git a/src/ClimaOcean.jl b/src/ClimaOcean.jl index a698920d2..a6c18f133 100644 --- a/src/ClimaOcean.jl +++ b/src/ClimaOcean.jl @@ -96,13 +96,12 @@ using PrecompileTools: @setup_workload, @compile_workload @setup_workload begin Nx, Ny, Nz = 32, 32, 10 @compile_workload begin - z = exponential_z_faces(Nz=Nz, depth=6000, h=34) - grid = Oceananigans.OrthogonalSphericalShellGrids.TripolarGrid(CPU(); size=(Nx, Ny, Nz), halo=(7, 7, 7), z) - grid = ImmersedBoundaryGrid(grid, GridFittedBottom((x, y) -> -5000)) - ocean = ocean_simulation(grid) - model = OceanSeaIceModel(ocean) + # z = exponential_z_faces(Nz=Nz, depth=6000, h=34) + # grid = Oceananigans.OrthogonalSphericalShellGrids.TripolarGrid(CPU(); size=(Nx, Ny, Nz), halo=(7, 7, 7), z) + # grid = ImmersedBoundaryGrid(grid, GridFittedBottom((x, y) -> -5000)) + # ocean = ocean_simulation(grid) + # model = OceanSeaIceModel(ocean) end end end # module - diff --git a/src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl b/src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl index ab77fc816..9d881c29c 100644 --- a/src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl +++ b/src/OceanSeaIceModels/InterfaceComputations/component_interfaces.jl @@ -169,10 +169,10 @@ Base.show(io::IO, crf::ComponentInterfaces) = print(io, summary(crf)) atmosphere_ocean_interface(::Nothing, args...) = nothing function atmosphere_ocean_interface(atmos, - ocean, - radiation, + ocean, + radiation, ao_flux_formulation, - temperature_formulation, + temperature_formulation, velocity_formulation, specific_humidity_formulation) @@ -204,8 +204,8 @@ atmosphere_sea_ice_interface(::Nothing, ::Nothing, args...) = nothing atmosphere_sea_ice_interface(::Nothing, ::SeaIceSimulation, args...) = nothing function atmosphere_sea_ice_interface(atmos, - sea_ice::SeaIceSimulation, - radiation, + sea_ice::SeaIceSimulation, + radiation, ai_flux_formulation, temperature_formulation, velocity_formulation) @@ -351,10 +351,10 @@ function ComponentInterfaces(atmosphere, ocean, sea_ice=nothing; liquidus = sea_ice.model.ice_thermodynamics.phase_transitions.liquidus, temperature_units = sea_ice_temperature_units) - net_momentum_fluxes = if sea_ice.model.dynamics isa Nothing + net_momentum_fluxes = if sea_ice.model.dynamics isa Nothing u = Field{Face, Center, Nothing}(sea_ice.model.grid) v = Field{Center, Face, Nothing}(sea_ice.model.grid) - (; u, v) + (; u, v) else u = sea_ice.model.dynamics.external_momentum_stresses.top.u v = sea_ice.model.dynamics.external_momentum_stresses.top.v @@ -408,4 +408,3 @@ function sea_ice_similarity_theory(sea_ice::SeaIceSimulation) interface_temperature_type = SkinTemperature(internal_flux) return SimilarityTheoryFluxes(; interface_temperature_type) end - diff --git a/src/OceanSeaIceModels/ocean_sea_ice_model.jl b/src/OceanSeaIceModels/ocean_sea_ice_model.jl index aa9b9ec9d..06e077ac9 100644 --- a/src/OceanSeaIceModels/ocean_sea_ice_model.jl +++ b/src/OceanSeaIceModels/ocean_sea_ice_model.jl @@ -18,7 +18,7 @@ import Oceananigans.OutputWriters: checkpointer_address, default_checkpointed_properties import Oceananigans.Simulations: reset!, initialize!, iteration, run! -import Oceananigans.TimeSteppers: time_step!, update_state!, time +import Oceananigans.TimeSteppers: time_step!, update_state!, time, set_clock! import Oceananigans.Utils: prettytime import .PrescribedAtmospheres: set_clock! @@ -74,6 +74,8 @@ checkpointer_address(::OSIM) = "OceanSeaIceModel" reset!(model::OSIM) = reset!(model.ocean) +components(model::OSIM) = (model.atmosphere, model.ocean.model) + # Make sure to initialize the exchanger here function initialization_update_state!(model::OSIM) initialize!(model.interfaces.exchanger, model.atmosphere) @@ -83,19 +85,25 @@ end function initialize!(model::OSIM) initialize!(model.ocean) + initialize!(model.atmosphere) initialize!(model.interfaces.exchanger, model.atmosphere) return nothing end function set_clock!(model::OSIM, clock) - model.clock.time = clock.time - model.clock.iteration = clock.iteration - model.clock.last_Δt = clock.last_Δt - model.clock.last_stage_Δt = clock.last_stage_Δt - model.clock.stage = clock.stage + + set_clock!(model.clock, clock) + + # set the component clocks + atmos, ocean = model.atmosphere, model.ocean + set_clock!(atmos, clock) + set_clock!(ocean, clock) + return nothing end +set_clock!(sim::Simulation, new_clock) = set_clock!(sim.model, new_clock) + reference_density(unsupported) = throw(ArgumentError("Cannot extract reference density from $(typeof(unsupported))")) diff --git a/src/OceanSeaIceModels/output_writers.jl b/src/OceanSeaIceModels/output_writers.jl index 06e4b1550..9144e2e39 100644 --- a/src/OceanSeaIceModels/output_writers.jl +++ b/src/OceanSeaIceModels/output_writers.jl @@ -10,19 +10,22 @@ import Oceananigans.OutputWriters: write_output! wall_time = Ref(time_ns()) -function set!(model::OSIMPA, checkpoint_file_path::AbstractString) +function set!(model::OSIM, checkpoint_file_path::AbstractString) addr = checkpointer_address(model) + checkpointed_clock = nothing jldopen(checkpoint_file_path, "r") do file checkpointed_clock = file["$addr/clock"] - - # Update model clock - set_clock!(model, checkpointed_clock) end + @show checkpointed_clock + + # Update model clock + set_clock!(model, checkpointed_clock) + # deal with model components - for component in (model.ocean.model, model.atmosphere) + for component in components(model) set!(component, checkpoint_file_path) end @@ -43,10 +46,8 @@ function write_output!(c::Checkpointer, model::OSIM) properties = default_checkpointed_properties(model)) # write the OceanSeaIceModel components - components = (model.atmosphere, model.ocean.model) - mode = "a" # to append in the file already written above - for component in components + for component in components(model) write_output!(c, component, filepath, mode, properties = default_checkpointed_properties(component)) end diff --git a/src/OceanSimulations/ocean_simulation.jl b/src/OceanSimulations/ocean_simulation.jl index 99c4adbc8..040b2ff82 100644 --- a/src/OceanSimulations/ocean_simulation.jl +++ b/src/OceanSimulations/ocean_simulation.jl @@ -33,23 +33,23 @@ function estimate_maximum_Δt(grid) # - 1.875 minutes for a 1/32 degree ocean Δt = 30minutes / Δθ - + return all_reduce(min, Δt, arch) end const TripolarOfSomeKind = Union{TripolarGrid, ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:TripolarGrid}} -function default_free_surface(grid::TripolarOfSomeKind; +function default_free_surface(grid::TripolarOfSomeKind; fixed_Δt = estimate_maximum_Δt(grid), - cfl = 0.7) + cfl = 0.7) free_surface = SplitExplicitFreeSurface(grid; cfl, fixed_Δt) return free_surface end -function default_free_surface(grid::DistributedGrid; +function default_free_surface(grid::DistributedGrid; fixed_Δt = compute_maximum_Δt(grid), - cfl = 0.7) - + cfl = 0.7) + free_surface = SplitExplicitFreeSurface(grid; cfl, fixed_Δt) substeps = length(free_surface.substepping.averaging_weights) substeps = all_reduce(max, substeps, architecture(grid)) @@ -84,8 +84,8 @@ default_tracer_advection() = FluxFormAdvection(WENO(order=7), @inline v_quadratic_bottom_drag(i, j, grid, c, Φ, μ) = @inbounds - μ * Φ.v[i, j, 1] * spᶜᶠᶜ(i, j, 1, grid, Φ) # Keep a constant linear drag parameter independent on vertical level -@inline u_immersed_bottom_drag(i, j, k, grid, clock, fields, μ) = @inbounds - μ * fields.u[i, j, k] * spᶠᶜᶜ(i, j, k, grid, fields) -@inline v_immersed_bottom_drag(i, j, k, grid, clock, fields, μ) = @inbounds - μ * fields.v[i, j, k] * spᶜᶠᶜ(i, j, k, grid, fields) +@inline u_immersed_bottom_drag(i, j, k, grid, clock, fields, μ) = @inbounds - μ * fields.u[i, j, k] * spᶠᶜᶜ(i, j, k, grid, fields) +@inline v_immersed_bottom_drag(i, j, k, grid, clock, fields, μ) = @inbounds - μ * fields.v[i, j, k] * spᶜᶠᶜ(i, j, k, grid, fields) # TODO: Specify the grid to a grid on the sphere; otherwise we can provide a different # function that requires latitude and longitude etc for computing coriolis=FPlane... @@ -145,10 +145,10 @@ function ocean_simulation(grid; end bottom_drag_coefficient = default_or_override(bottom_drag_coefficient) - + u_immersed_drag = FluxBoundaryCondition(u_immersed_bottom_drag, discrete_form=true, parameters=bottom_drag_coefficient) v_immersed_drag = FluxBoundaryCondition(v_immersed_bottom_drag, discrete_form=true, parameters=bottom_drag_coefficient) - + u_immersed_bc = ImmersedBoundaryCondition(bottom=u_immersed_drag) v_immersed_bc = ImmersedBoundaryCondition(bottom=v_immersed_drag) @@ -175,7 +175,7 @@ function ocean_simulation(grid; v_top_bc = FluxBoundaryCondition(τy) T_top_bc = FluxBoundaryCondition(Jᵀ) S_top_bc = FluxBoundaryCondition(Jˢ) - + u_bot_bc = FluxBoundaryCondition(u_quadratic_bottom_drag, discrete_form=true, parameters=bottom_drag_coefficient) v_bot_bc = FluxBoundaryCondition(v_quadratic_bottom_drag, discrete_form=true, parameters=bottom_drag_coefficient) @@ -228,4 +228,3 @@ end hasclosure(closure, ClosureType) = closure isa ClosureType hasclosure(closure_tuple::Tuple, ClosureType) = any(hasclosure(c, ClosureType) for c in closure_tuple) - diff --git a/test/runtests.jl b/test/runtests.jl index dcd162992..e25f5b458 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -5,6 +5,7 @@ using CUDA test_group = get(ENV, "TEST_GROUP", :all) test_group = Symbol(test_group) +test_group = :ocean_sea_ice_model using ClimaOcean.ECCO: download_dataset @@ -16,17 +17,17 @@ if test_group == :init || test_group == :all ### ### Download bathymetry data ### - - download_bathymetry() + + download_bathymetry() #### - #### Download JRA55 data + #### Download JRA55 data #### - + atmosphere = JRA55PrescribedAtmosphere() #### - #### Download ECCO data + #### Download ECCO data #### download_dataset(temperature_metadata) @@ -63,4 +64,3 @@ end if test_group == :distributed || test_group == :all include("test_distributed_utils.jl") end - diff --git a/test/test_ocean_sea_ice_model.jl b/test/test_ocean_sea_ice_model.jl index 4f701ddfd..535a98444 100644 --- a/test/test_ocean_sea_ice_model.jl +++ b/test/test_ocean_sea_ice_model.jl @@ -1,14 +1,17 @@ include("runtests_setup.jl") using CUDA +using Oceananigans.Units +using Oceananigans: prognostic_fields +using Oceananigans.TimeSteppers: QuasiAdamsBashforth2TimeStepper using Oceananigans.OrthogonalSphericalShellGrids -using ClimaOcean.OceanSeaIceModels: above_freezing_ocean_temperature! +using ClimaOcean.OceanSeaIceModels: above_freezing_ocean_temperature!, OceanSeaIceModel using ClimaSeaIce.SeaIceThermodynamics: melting_temperature using ClimaSeaIce.SeaIceMomentumEquations using ClimaSeaIce.Rheologies @inline kernel_melting_temperature(i, j, k, grid, liquidus, S) = @inbounds melting_temperature(liquidus, S[i, j, k]) - +#= @testset "Time stepping test" begin for arch in test_architectures @@ -101,34 +104,162 @@ using ClimaSeaIce.Rheologies end end end +=# """ - testbed_coupled_simulation(grid; stop_iteration=8) + test_ocean_sea_ice_model_checkpointer(grid) -Return a test-bed coupled simulation with a Checkpointer. +Return two ocean-sea ice coupled models to be used for testing purposes. """ -function testbed_coupled_simulation(grid; stop_iteration=8) - ocean = ocean_simulation(grid) - +function test_ocean_sea_ice_model_checkpointer(grid) + ocean = ocean_simulation(grid, closure=nothing, free_surface = ExplicitFreeSurface()) atmosphere = JRA55PrescribedAtmosphere(arch; backend=JRA55NetCDFBackend(4)) - coupled_model = OceanSeaIceModel(ocean; atmosphere) + true_coupled_model = OceanSeaIceModel(ocean; atmosphere) + # true_coupled_model = OceanSeaIceModel(ocean; atmosphere=nothing, radiation=nothing) + + test_coupled_model = deepcopy(true_coupled_model) + + return true_coupled_model, test_coupled_model +end + +function test_model_equality(test_model::OceanSeaIceModel, true_model::OceanSeaIceModel) + test_clock_equality(test_model.clock, true_model.clock) + + # test equality for components + test_model_equality(test_model.atmosphere, true_model.atmosphere) + test_model_equality(test_model.ocean.model, true_model.ocean.model) + + return nothing +end + +test_model_equality(test_model::PrescribedAtmosphere, true_model::PrescribedAtmosphere) = + test_clock_equality(test_model.clock, true_model.clock) + +function test_clock_equality(clock1::Clock, clock2::Clock) + @test clock1.iteration == clock2.iteration + @test clock1.time == clock2.time + @test clock1.last_Δt == clock2.last_Δt +return nothing +end + +function test_model_equality(test_model::HydrostaticFreeSurfaceModel, true_model::HydrostaticFreeSurfaceModel) + + test_clock_equality(test_model.clock, true_model.clock) + + CUDA.@allowscalar begin + test_model_fields = prognostic_fields(test_model) + true_model_fields = prognostic_fields(true_model) + field_names = keys(test_model_fields) - simulation = Simulation(coupled_model; Δt=10, stop_iteration) + for name in field_names + @show name + @show all(test_model_fields[name].data .≈ true_model_fields[name].data) - simulation.output_writers[:checkpoint] = Checkpointer(ocean.model; - schedule = IterationInterval(3), - prefix = "checkpointer", - dir = ".", - verbose = true, - overwrite_existing = true) + if test_model.timestepper isa QuasiAdamsBashforth2TimeStepper + if name ∈ keys(test_model.timestepper.Gⁿ) + @show name + # @show test_model.timestepper.Gⁿ[name] + # @show true_model.timestepper.Gⁿ[name] + @show all(test_model.timestepper.Gⁿ[name].data .≈ true_model.timestepper.Gⁿ[name].data) + @show all(test_model.timestepper.G⁻[name].data .≈ true_model.timestepper.G⁻[name].data) + end + end + end + end - return simulation + return nothing end -@testset "Checkpointer test" begin +function run_checkpointer_tests(true_model, test_model, Δt) + + true_simulation = Simulation(true_model; Δt, stop_iteration = 4) + true_simulation.output_writers[:checkpoint] = + Checkpointer(true_model, schedule=IterationInterval(4), overwrite_existing=true) + + run!(true_simulation) + + checkpointed_model = deepcopy(true_simulation.model) + + true_simulation.stop_iteration = 7 + run!(true_simulation) + + ##### + ##### Test `set!(model, checkpoint_file)` + ##### + + @info "Testing `set!(model, checkpoint_file)`" + + set!(test_model, "checkpoint_iteration4.jld2") + + test_model_equality(test_model, checkpointed_model) + + #= + ##### + ##### Test pickup from explicit checkpoint path + ##### + + @info "Testing pickup from explicit checkpoint path" + + test_simulation = Simulation(test_model; Δt, stop_iteration=7) + run!(test_simulation, pickup="checkpoint_iteration0.jld2") + + test_model_equality(test_model, true_model) + + run!(test_simulation, pickup="checkpoint_iteration4.jld2") + + @info "Testing model equality when running with pickup=checkpoint_iteration4.jld2." + + test_model_equality(test_model, true_model) + + @test test_simulation.model.clock.iteration == true_simulation.model.clock.iteration + @test test_simulation.model.clock.time == true_simulation.model.clock.time + test_model_equality(test_simulation.model, true_simulation.model) + + ##### + ##### Test `run!(sim, pickup=true) + ##### + + @info "Testing run!(sim, pickup=true)" + + # Pickup using existing checkpointer + test_simulation.output_writers[:checkpoint] = + Checkpointer(test_model, schedule=IterationInterval(4), overwrite_existing=true) + + run!(test_simulation, pickup=true) + @info " Testing model equality when running with pickup=true." + + @test test_simulation.model.clock.iteration == true_simulation.model.clock.iteration + @test test_simulation.model.clock.time == true_simulation.model.clock.time + test_model_equality(test_simulation.model, true_simulation.model) + + run!(test_simulation, pickup=0) + @info " Testing model equality when running with pickup=0." + + @test test_simulation.model.clock.iteration == true_simulation.model.clock.iteration + @test test_simulation.model.clock.time == true_simulation.model.clock.time + test_model_equality(test_simulation.model, true_simulation.model) + + run!(test_simulation, pickup=4) + @info " Testing model equality when running with pickup=4." + + @test test_simulation.model.clock.iteration == true_simulation.model.clock.iteration + @test test_simulation.model.clock.time == true_simulation.model.clock.time + test_model_equality(test_simulation.model, true_simulation.model) + + =# + + for file in Tuple("checkpoint_iteration$i.jld2" for i in 0:4:10) + rm(file, force=true) + end +end + +#= +@testset "Checkpointer tests" begin for arch in test_architectures + Δt = 1.2hours + Nx, Ny, Nz = 40, 25, 10 grid = LatitudeLongitudeGrid(arch; @@ -138,20 +269,46 @@ end latitude = (-75, 75), longitude = (0, 360)) - simulation = testbed_coupled_simulation(grid; stop_iteration=8) + true_coupled_model, test_coupled_model = test_ocean_sea_ice_model_checkpointer(grid) + + run_checkpointer_tests(true_coupled_model, test_coupled_model, Δt) + + #= + true_simulation = testbed_coupled_simulation(true_coupled_model; Δt, checkpoint_iteration_interval, + stop_iteration=intermediate_iteration) + + test_simulation = testbed_coupled_simulation(test_coupled_model; Δt, checkpoint_iteration_interval + stop_iteration = final_iteration) + + run!(true_simulation) # up to intermediate_iteration - run!(simulation) + checkpointed_model = deepcopy(true_simulation.model) - # create a new simulation and pick up from the last checkpointer - new_simulation = testbed_coupled_simulation(grid; stop_iteration=13) + set!(test_model, "checkpoint_iteration5.jld2") + + @test test_model.clock.iteration == checkpointed_model.clock.iteration + @test test_model.clock.time == checkpointed_model.clock.time run!(new_simulation, pickup=true) - # ensure the ocean, atmosphere, and coupled model are all at same time and iteration - clock = new_simulation.model.ocean.model.clock - @test new_simulation.model.atmosphere.clock.iteration ≈ clock.iteration - @test new_simulation.model.atmosphere.clock.time ≈ clock.time - @test new_simulation.model.clock.iteration ≈ clock.iteration - @test new_simulation.model.clock.time ≈ clock.time + # ensure all clocks all at same time and iteration + for clock in (simulation.model.clock, + simulation.model.atmosphere.clock, + simulation.model.ocean.model.clock) + + @test clock.iteration ≈ intermediate_iteration + @test clock.time ≈ intermediate_iteration * Δt + end + + for clock in (new_simulation.model.clock, + new_simulation.model.atmosphere.clock, + new_simulation.model.ocean.model.clock) + + @test clock.iteration ≈ final_iteration + @test clock.time ≈ final_iteration * Δt + end + =# end + end +=# From 8bd0d3608562225139e183bbfbe0e3101798e698 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Fri, 25 Apr 2025 11:31:15 +1000 Subject: [PATCH 54/54] Update ClimaOcean.jl --- src/ClimaOcean.jl | 6 ------ 1 file changed, 6 deletions(-) diff --git a/src/ClimaOcean.jl b/src/ClimaOcean.jl index 6edb0db65..d7c77ff4b 100644 --- a/src/ClimaOcean.jl +++ b/src/ClimaOcean.jl @@ -99,15 +99,9 @@ using PrecompileTools: @setup_workload, @compile_workload @setup_workload begin Nx, Ny, Nz = 32, 32, 10 @compile_workload begin -<<<<<<< HEAD - # z = exponential_z_faces(Nz=Nz, depth=6000, h=34) - # grid = Oceananigans.OrthogonalSphericalShellGrids.TripolarGrid(CPU(); size=(Nx, Ny, Nz), halo=(7, 7, 7), z) - # grid = ImmersedBoundaryGrid(grid, GridFittedBottom((x, y) -> -5000)) -======= z = exponential_z_faces(Nz=Nz, depth=6000, h=34) grid = Oceananigans.OrthogonalSphericalShellGrids.TripolarGrid(CPU(); size=(Nx, Ny, Nz), halo=(7, 7, 7), z) grid = ImmersedBoundaryGrid(grid, GridFittedBottom((x, y) -> -5000)) ->>>>>>> main # ocean = ocean_simulation(grid) # model = OceanSeaIceModel(ocean) end