From 6a87267ed923c0c291397f1d502a3f45d0743e09 Mon Sep 17 00:00:00 2001 From: "Gregory L. Wagner" Date: Mon, 16 Dec 2024 12:45:32 -0700 Subject: [PATCH 1/7] Updates diagnostic for computing mixed layer depth (#296) * Cleaner syntax for ECCOFieldTimeSeries * Fix indenting * Make grid positional argument * Eliminate architecture kwarg * Clean up ECCORestoring * More simplification * fix typo * fix typo * Bugfix * Fix test syntax * Bump project * Implement mixed layer depth diagnostic * Shuffle files * Change difference criterion * Deleted Diagnostics.jl * Update to latest Oceananigans * Add a test for MixedLayerDepthField * Add test for MixeDLayerDepthField * Add allowscalar * Hmm * Update src/Diagnostics/mixed_layer_depth.jl Co-authored-by: Simone Silvestri * Import static_column_depth --------- Co-authored-by: Navid C. Constantinou Co-authored-by: Simone Silvestri --- examples/ecco_inspect_temperature_salinity.jl | 87 ++++++++++++ examples/ecco_mixed_layer_depth.jl | 81 ++++++++++++ examples/inspect_ecco_data.jl | 97 -------------- mwe.jl | 11 ++ src/ClimaOcean.jl | 2 +- src/Diagnostics.jl | 124 ------------------ src/Diagnostics/Diagnostics.jl | 19 +++ src/Diagnostics/mixed_layer_depth.jl | 92 +++++++++++++ test/runtests.jl | 1 + test/runtests_setup.jl | 2 + test/test_diagnostics.jl | 52 ++++++++ test/test_jra55.jl | 2 +- 12 files changed, 347 insertions(+), 223 deletions(-) create mode 100644 examples/ecco_inspect_temperature_salinity.jl create mode 100644 examples/ecco_mixed_layer_depth.jl delete mode 100644 examples/inspect_ecco_data.jl create mode 100644 mwe.jl delete mode 100644 src/Diagnostics.jl create mode 100644 src/Diagnostics/Diagnostics.jl create mode 100644 src/Diagnostics/mixed_layer_depth.jl create mode 100644 test/test_diagnostics.jl diff --git a/examples/ecco_inspect_temperature_salinity.jl b/examples/ecco_inspect_temperature_salinity.jl new file mode 100644 index 00000000..4fcaf1ed --- /dev/null +++ b/examples/ecco_inspect_temperature_salinity.jl @@ -0,0 +1,87 @@ +using Oceananigans +using Oceananigans.ImmersedBoundaries: mask_immersed_field! + +using GLMakie +using Printf +using ClimaOcean +using ClimaOcean.DataWrangling.ECCO: ECCO_field, ECCOFieldTimeSeries +using CFTime +using Dates + +arch = CPU() +Nx = 360 ÷ 4 +Ny = 160 ÷ 4 + +z = ClimaOcean.DataWrangling.ECCO.ECCO_z +z = z[20:end] +Nz = length(z) - 1 + +grid = LatitudeLongitudeGrid(arch; z, + size = (Nx, Ny, Nz), + latitude = (-80, 80), + longitude = (0, 360)) + +bottom_height = regrid_bathymetry(grid; + minimum_depth = 10, + interpolation_passes = 5, + major_basins = 1) + +grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height)) + +T = CenterField(grid) +S = CenterField(grid) + +using SeawaterPolynomials: TEOS10EquationOfState +using Oceananigans.BuoyancyModels: buoyancy + +equation_of_state = TEOS10EquationOfState() +sb = SeawaterBuoyancy(; equation_of_state) +tracers = (T=T, S=S) +b = Field(buoyancy(sb, grid, tracers)) + +start = DateTimeProlepticGregorian(1993, 1, 1) +stop = DateTimeProlepticGregorian(1999, 1, 1) +dates = range(start; stop, step=Month(1)) + +Tmeta = ECCOMetadata(:temperature; dates) +Smeta = ECCOMetadata(:salinity; dates) + +Tt = ECCOFieldTimeSeries(Tmeta, grid; time_indices_in_memory=length(dates)) +St = ECCOFieldTimeSeries(Smeta, grid; time_indices_in_memory=length(dates)) + +fig = Figure(size=(900, 1050)) + +axT = Axis(fig[1, 1]) +axS = Axis(fig[2, 1]) +axb = Axis(fig[3, 1]) + +Nt = length(dates) +grid = T.grid +Nz = size(grid, 3) +kslider = Slider(fig[1:3, 0], range=1:Nz, startvalue=Nz, horizontal=false) +nslider = Slider(fig[4, 1:2], range=1:Nt, startvalue=1) +k = kslider.value +n = nslider.value + +Tk = @lift view(Tt[$n], :, :, $k) +Sk = @lift view(St[$n], :, :, $k) + +Δb = @lift begin + parent(T) .= parent(Tt[$n]) + parent(S) .= parent(St[$n]) + compute!(b) + mask_immersed_field!(b, NaN) + Δb = interior(b, :, :, Nz) .- interior(b, :, :, $k) + Δb +end + +hmT = heatmap!(axT, Tk, nan_color=:lightgray, colorrange=(-2, 30), colormap=:thermal) +hmS = heatmap!(axS, Sk, nan_color=:lightgray, colorrange=(31, 37), colormap=:haline) +hmb = heatmap!(axb, Δb, nan_color=:lightgray, colorrange=(0, 1e-3), colormap=:magma) + +Colorbar(fig[1, 2], hmT, label="Temperature (ᵒC)") +Colorbar(fig[2, 2], hmS, label="Salinity (g kg⁻¹)") +Colorbar(fig[3, 2], hmb, label="Buoyancy difference (m s⁻²)") + +display(fig) + diff --git a/examples/ecco_mixed_layer_depth.jl b/examples/ecco_mixed_layer_depth.jl new file mode 100644 index 00000000..ff3165e5 --- /dev/null +++ b/examples/ecco_mixed_layer_depth.jl @@ -0,0 +1,81 @@ +using ClimaOcean +using ClimaOcean.Diagnostics: MixedLayerDepthField +using ClimaOcean.DataWrangling.ECCO: ECCO_field, ECCOFieldTimeSeries +using Oceananigans +using GLMakie +using Printf +using CFTime +using Dates + +using SeawaterPolynomials: TEOS10EquationOfState +using Oceananigans.BuoyancyFormulations: buoyancy + +arch = CPU() +Nx = 360 ÷ 1 +Ny = 160 ÷ 1 + +z = ClimaOcean.DataWrangling.ECCO.ECCO_z +z = z[20:end] +Nz = length(z) - 1 + +grid = LatitudeLongitudeGrid(arch; z, + size = (Nx, Ny, Nz), + latitude = (-80, 80), + longitude = (0, 360)) + +bottom_height = regrid_bathymetry(grid; + minimum_depth = 10, + interpolation_passes = 5, + major_basins = 1) + +grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height)) + +start = DateTimeProlepticGregorian(1993, 1, 1) +stop = DateTimeProlepticGregorian(2003, 1, 1) +dates = range(start; stop, step=Month(1)) + +Tmeta = ECCOMetadata(:temperature; dates) +Smeta = ECCOMetadata(:salinity; dates) + +Tt = ECCOFieldTimeSeries(Tmeta, grid; time_indices_in_memory=2) +St = ECCOFieldTimeSeries(Smeta, grid; time_indices_in_memory=2) +ht = FieldTimeSeries{Center, Center, Nothing}(grid, Tt.times) + +equation_of_state = TEOS10EquationOfState() +sb = SeawaterBuoyancy(; equation_of_state) +tracers = (T=Tt[1], S=St[1]) +h = MixedLayerDepthField(sb, grid, tracers) + +Nt = length(ht) +for n = 1:Nt-1 + local tracers + tracers = (T=Tt[n], S=St[n]) + h.operand.buoyancy_perturbation = buoyancy(sb, grid, tracers) + @show n + @time compute!(h) + parent(ht[n]) .= parent(h) +end + +function titlestr(n) + d = dates[n] + yr = year(d) + mn = monthname(d) + return string("ECCO mixed layer depth on ", mn, " ", yr) +end + +fig = Figure(size=(1500, 800)) +axh = Axis(fig[2, 1], xlabel="Longitude", ylabel="Latitude") +n = Observable(1) + +str = @lift titlestr($n) +Label(fig[1, 1], str, tellwidth=false) + +hn = @lift ht[$n] +hm = heatmap!(axh, hn, colorrange=(0, 500), colormap=:magma, nan_color=:lightgray) +Colorbar(fig[2, 2], hm, label="Mixed layer depth (m)") +display(fig) + +record(fig, "ecco_mld.mp4", 1:Nt-1, framerate=4) do nn + @info "Drawing frame $nn of $Nt..." + n[] = nn +end diff --git a/examples/inspect_ecco_data.jl b/examples/inspect_ecco_data.jl deleted file mode 100644 index 3d260436..00000000 --- a/examples/inspect_ecco_data.jl +++ /dev/null @@ -1,97 +0,0 @@ -# # A quick look at ECCO data -# -# ClimaOcean can download and utilize data from the "ECCO" state estimate, -# which stands for "Estimating the Circulation and Climate of the Ocean" --- two! -# -# This example shows how to download three-dimensional temperature and salinity fields -# from ECCO, and makes a short animation to showcase the fields' content. -# -# For this example we need Oceananigans for Field utilities, CairoMakie for plotting, -# Printf for nice labeling, and of course ClimaOcean to actually download and construct -# the ECCO fields. - -using Oceananigans -using CairoMakie -using Printf - -using ClimaOcean.DataWrangling.ECCO: ECCO_field - -# The function `ECCO_field` provided by `ClimaOcean.DataWrangling.ECCO` automatically -# downloads ECCO data, if the data doesn't already exist at the default location. - -T = ECCO_field(:temperature) -S = ECCO_field(:salinity) - -# Next, we massage the ECCO data by inserting NaNs in "land cells", which -# are diagnosed by having an unphysically low temperature. - -Tp = parent(T) -Sp = parent(S) -Sp[Tp .< -10] .= NaN -Tp[Tp .< -10] .= NaN - -# # Plotting ECCO data -# -# We're ready to plot. We'll make an animation -# that depicts how the ECCO data changes with depth. - -fig = Figure(size=(900, 1050)) - -axT = Axis(fig[1, 1]) -axS = Axis(fig[2, 1]) - -# To make an animation that scrolls through the 3D temperature -# and salinity fields, we make an Observable for the vertical index, -# and then construct slices of T, S using the Observable, `k`. - -grid = T.grid -Nz = size(grid, 3) -k = Observable(Nz) - -Tk = @lift view(T, :, :, $k) -Sk = @lift view(S, :, :, $k) - -# Finally, we make a nice plot with a label that displays depth, colorbars, -# and light gray within land cells. - -hmT = heatmap!(axT, Tk, nan_color=:lightgray, colorrange=(-2, 30), colormap=:thermal) -hmS = heatmap!(axS, Sk, nan_color=:lightgray, colorrange=(31, 37), colormap=:haline) - -Colorbar(fig[1, 2], hmT, label="Temperature (ᵒC)") -Colorbar(fig[2, 2], hmS, label="Salinity (psu)") - -z = znodes(grid, Center()) -depth_str = @lift @sprintf("z = %d meters", z[$k]) -text!(axT, 50, 50, text=depth_str, color=:lemonchiffon, justification=:center, fontsize=20) -text!(axS, 50, 50, text=depth_str, color=:lemonchiffon, justification=:center, fontsize=20) - -# # Making the animation -# -# This animation is a little fancy. We start by displaying the surface -# field, then we scroll through depth to the bottom and pause there. -# Next, we scroll back to the surface and pause. - -stillframes = 10 -movingframes = Nz - -record(fig, "ECCO_temperature_salinity.mp4", framerate=4) do io - - [recordframe!(io) for _ = 1:stillframes] - - for kk in Nz:-2:1 - k[] = kk - recordframe!(io) - end - - [recordframe!(io) for _ = 1:stillframes] - - for kk in 1:2:Nz - k[] = kk - recordframe!(io) - end - - [recordframe!(io) for _ = 1:stillframes] -end -nothing #hide - -# ![](ECCO_temperature_salinity.mp4) diff --git a/mwe.jl b/mwe.jl new file mode 100644 index 00000000..f77545eb --- /dev/null +++ b/mwe.jl @@ -0,0 +1,11 @@ +using ClimaOcean +using NCDatasets + +cachepath = ClimaOcean.DataWrangling.JRA55.download_jra55_cache +filename = "RYF.tas.1990_1991.nc" +filepath = joinpath(cachepath, filename) + +ds = Dataset(filepath) +Nx, Ny, Nt = size(ds["tas"]) +ds["tas"][1, 1, [Nt, 1]] +close(ds) diff --git a/src/ClimaOcean.jl b/src/ClimaOcean.jl index 14b78efb..b64f6beb 100644 --- a/src/ClimaOcean.jl +++ b/src/ClimaOcean.jl @@ -78,7 +78,7 @@ include("VerticalGrids.jl") include("InitialConditions/InitialConditions.jl") include("DataWrangling/DataWrangling.jl") include("Bathymetry.jl") -include("Diagnostics.jl") +include("Diagnostics/Diagnostics.jl") include("OceanSimulations/OceanSimulations.jl") using .VerticalGrids diff --git a/src/Diagnostics.jl b/src/Diagnostics.jl deleted file mode 100644 index 0fad0b8e..00000000 --- a/src/Diagnostics.jl +++ /dev/null @@ -1,124 +0,0 @@ -module Diagnostics - -export MixedLayerDepthField - -using Oceananigans -using Oceananigans.BuoyancyFormulations: buoyancy -using Oceananigans.BoundaryConditions: fill_halo_regions! -using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid -using Oceananigans.Architectures: device, architecture -using Oceananigans.Utils: launch! -using Oceananigans.Grids: Center, Face, inactive_node, znode -using Oceananigans.Operators: Δzᶜᶜᶠ - -using KernelAbstractions: @index, @kernel -using KernelAbstractions.Extras.LoopInfo: @unroll - -import Oceananigans.Fields: compute! - -const c = Center() -const f = Face() - -@inline z_bottom(i, j, grid) = znode(c, c, f, i, j, 1, grid) -@inline z_bottom(i, j, grid::ImmersedBoundaryGrid) = @inbounds grid.immersed_boundary.bottom_height[i, j] - -@inline bottom(i, j, grid) = znode(c, c, f, i, j, 1, grid) -##### -##### MixedLayerDepthField -##### - -@kernel function compute_mld!(h, grid, b, Δb) - i, j = @index(Global, NTuple) - - Nz = grid.Nz - z_surface = znode(c, c, f, i, j, Nz+1, grid) - b_surface = @inbounds b[i, j, Nz] # buoyancy at surface - - # Start downwards look - z_ij = z_surface - found_mixed_layer_depth = false - - @unroll for k in Nz-1 : -1 : 1 # scroll from point just below surface - b⁺ = @inbounds b[i, j, k+1] - bᵏ = @inbounds b[i, j, k] - - # If buoyancy decreases downwards, both are > 0 - Δb⁺ = b_surface - b⁺ - Δbᵏ = b_surface - bᵏ - - zᵏ = znode(c, c, c, i, j, k, grid) - Δz⁺ = Δzᶜᶜᶠ(i, j, k+1, grid) - - # Assuming buoyancy decreases downwards - just_below_mixed_layer = (Δb⁺ < Δb) & (Δbᵏ >= Δb) - just_below_mixed_layer *= !found_mixed_layer_depth - - # Linearly interpolate to find mixed layer height - new_z_ij = zᵏ + (Δb - Δbᵏ) / (Δb⁺ - Δbᵏ) * Δz⁺ - - # Replace z_ij if we found a new mixed layer depth - z_ij = ifelse(!found_mixed_layer_depth | just_below_mixed_layer, new_z_ij, z_ij) - end - - # Note "-" since `h` is supposed to be "depth" rather than "height" - h_ij = z_surface - z_ij - @inbounds h[i, j, 1] = ifelse(inactive_node(i, j, Nz, grid, c, c, c), zero(grid), h_ij) -end - -struct MixedLayerDepthOperand{B, FT} - buoyancy_operation :: B - mixed_layer_buoyancy_differential :: FT -end - -Base.summary(op::MixedLayerDepthOperand) = "MixedLayerDepthOperand" - -const MixedLayerDepthField = Field{Center, Center, Nothing, <:MixedLayerDepthOperand} - -""" - MixedLayerDepthField(grid, tracers, buoyancy_model; Δb = 3e-4, field_kw...) - -Return a reduced `Field{Center, Center, Nothing}` that represents -mixed layer depth for `model`, based on a buoyancy differential criterion. -The mixed layer depth is defined as the depth ``h`` for which - -```math -b(z=0) - b(z=-h) = Δb -``` - -This criterion is solved by integrating downwards and linearly interpolating to find `h`, -assuming that ``b`` decreases with depth. - -Keyword arguments -================= - -* `Δb`: Buoyancy differential used to calculate mixed layer depth -* `field_kw`: Keyword arguments passed to `Field`. - -Example -======= - -```julia -h = MixedLayerDepth(model) -compute!(h) # compute mixed layer depth -``` -""" -function MixedLayerDepthField(grid, tracers, buoyancy_model; Δb = 3e-4, field_kw...) - b_op = buoyancy(buoyancy_model, grid, tracers) - operand = MixedLayerDepthOperand(b_op, Δb) - return Field{Center, Center, Nothing}(grid; operand, field_kw...) -end - -MixedLayerDepthField(model; kw...) = MixedLayerDepthField(model.grid, - model.tracers, - model.buoyancy; kw...) - -function compute!(h::MixedLayerDepthField, time=nothing) - arch = architecture(h) - b = h.operand.buoyancy_operation - Δb = h.operand.mixed_layer_buoyancy_differential - launch!(arch, h.grid, :xy, compute_mld!, h, h.grid, b, Δb) - fill_halo_regions!(h) - return h -end - -end # module diff --git a/src/Diagnostics/Diagnostics.jl b/src/Diagnostics/Diagnostics.jl new file mode 100644 index 00000000..7f3931f1 --- /dev/null +++ b/src/Diagnostics/Diagnostics.jl @@ -0,0 +1,19 @@ +module Diagnostics + +export MixedLayerDepthField, MixedLayerDepthOperand + +using Oceananigans +using Oceananigans.Architectures: architecture +using Oceananigans.BuoyancyFormulations: buoyancy +using Oceananigans.Grids: new_data, inactive_cell, znode +using Oceananigans.BoundaryConditions: FieldBoundaryConditions, fill_halo_regions! +using Oceananigans.Fields: FieldStatus +using Oceananigans.Utils: launch! + +import Oceananigans.Fields: compute! + +using KernelAbstractions: @index, @kernel + +include("mixed_layer_depth.jl") + +end # module diff --git a/src/Diagnostics/mixed_layer_depth.jl b/src/Diagnostics/mixed_layer_depth.jl new file mode 100644 index 00000000..8da401ea --- /dev/null +++ b/src/Diagnostics/mixed_layer_depth.jl @@ -0,0 +1,92 @@ +using Oceananigans.Grids: static_column_depthᶜᶜᵃ + +mutable struct MixedLayerDepthOperand{FT, B} + buoyancy_perturbation :: B + difference_criterion :: FT +end + +Base.summary(mldo::MixedLayerDepthOperand) = "MixedLayerDepthOperand" + +function MixedLayerDepthOperand(bm, grid, tracers; difference_criterion=1e-4) + buoyancy_perturbation = buoyancy(bm, grid, tracers) + difference_criterion = convert(eltype(grid), difference_criterion) + return MixedLayerDepthOperand(buoyancy_perturbation, difference_criterion) +end + +const MixedLayerDepthField = Field{<:Any, <:Any, <:Any, <:MixedLayerDepthOperand} + +""" + MixedLayerDepthField(bm, grid, tracers; difference_criterion=1e-4) + +""" +function MixedLayerDepthField(bm, grid, tracers; difference_criterion=3e-5) + operand = MixedLayerDepthOperand(bm, grid, tracers; difference_criterion) + loc = (Center, Center, Nothing) + indices = (:, :, :) + bcs = FieldBoundaryConditions(grid, loc) + data = new_data(grid, loc, indices) + recompute_safely = false + status = FieldStatus() + return Field(loc, grid, data, bcs, indices, operand, status) +end + +function compute!(mld::MixedLayerDepthField, time=nothing) + compute_mixed_layer_depth!(mld) + #@apply_regionally compute_mixed_layer_depth!(mld) + fill_halo_regions!(mld) + return mld +end + +function compute_mixed_layer_depth!(mld) + grid = mld.grid + arch = architecture(grid) + + launch!(arch, mld.grid, :xy, + _compute_mixed_layer_depth!, + mld, + grid, + mld.operand.buoyancy_perturbation, + mld.operand.difference_criterion) + + return mld +end + +const c = Center() +const f = Face() + +@kernel function _compute_mixed_layer_depth!(mld, grid, b, Δb★) + i, j = @index(Global, NTuple) + Nz = size(grid, 3) + + Δb = zero(grid) + bN = @inbounds b[i, j, Nz] + mixed = true + k = Nz - 1 + inactive = inactive_cell(i, j, k, grid) + + while !inactive & mixed & (k > 0) + Δb = @inbounds bN - b[i, j, k] + mixed = Δb < Δb★ + k -= 1 + inactive = inactive_cell(i, j, k, grid) + end + + # Linearly interpolate + # z★ = zk + Δz/Δb * (Δb★ - Δb) + zN = znode(i, j, Nz, grid, c, c, c) + zk = znode(i, j, k, grid, c, c, c) + Δz = zN - zk + z★ = zk - Δz/Δb * (Δb★ - Δb) + + # Special case when domain is one grid cell deep + z★ = ifelse(Δb == 0, zN, z★) + + # Apply various criterion + h = -z★ + h = max(h, zero(grid)) + H = static_column_depthᶜᶜᵃ(i, j, grid) + h = min(h, H) + + @inbounds mld[i, j, 1] = h +end + diff --git a/test/runtests.jl b/test/runtests.jl index 932f48ef..aca50ba5 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -58,6 +58,7 @@ end if test_group == :simulations || test_group == :all CUDA.set_runtime_version!(v"12.2", local_toolkit = true) # Seems to help in finding the correct CUDA version include("test_simulations.jl") + include("test_diagnostics.jl") end if test_group == :distributed || test_group == :all diff --git a/test/runtests_setup.jl b/test/runtests_setup.jl index 8dad9db8..ecd8c91b 100644 --- a/test/runtests_setup.jl +++ b/test/runtests_setup.jl @@ -16,6 +16,8 @@ using ClimaOcean.Bathymetry: download_bathymetry_cache using CFTime using Dates +using CUDA: @allowscalar + gpu_test = parse(Bool, get(ENV, "GPU_TEST", "false")) test_architectures = gpu_test ? [GPU()] : [CPU()] diff --git a/test/test_diagnostics.jl b/test/test_diagnostics.jl new file mode 100644 index 00000000..fc64742d --- /dev/null +++ b/test/test_diagnostics.jl @@ -0,0 +1,52 @@ +include("runtests_setup.jl") + +using SeawaterPolynomials: TEOS10EquationOfState +using Oceananigans.BuoyancyFormulations: buoyancy +using Oceananigans: location +using ClimaOcean.DataWrangling.ECCO: ECCO_field, ECCOFieldTimeSeries +using ClimaOcean.Diagnostics: MixedLayerDepthField, MixedLayerDepthOperand + +@testset "MixedLayerDepthField" begin + for arch in test_architectures + grid = LatitudeLongitudeGrid(arch; + size = (3, 3, 100), + latitude = (0, 30), + longitude = (150, 180), + z = (-1000, 0)) + + bottom_height = regrid_bathymetry(grid; + minimum_depth = 10, + interpolation_passes = 5, + major_basins = 1) + + grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height)) + + start = DateTimeProlepticGregorian(1993, 1, 1) + stop = DateTimeProlepticGregorian(1993, 2, 1) + dates = range(start; stop, step=Month(1)) + + Tmeta = ECCOMetadata(:temperature; dates) + Smeta = ECCOMetadata(:salinity; dates) + + Tt = ECCOFieldTimeSeries(Tmeta, grid; time_indices_in_memory=2) + St = ECCOFieldTimeSeries(Smeta, grid; time_indices_in_memory=2) + + equation_of_state = TEOS10EquationOfState() + sb = SeawaterBuoyancy(; equation_of_state) + tracers = (T=Tt[1], S=St[1]) + h = MixedLayerDepthField(sb, grid, tracers) + + @test h isa Field + @test location(h) == (Center, Center, Nothing) + @test h.operand isa MixedLayerDepthOperand + @test h.operand.buoyancy_perturbation isa KernelFunctionOperation + + compute!(h) + @test @allowscalar h[1, 1, 1] ≈ 16.255836 # m + + tracers = (T=Tt[2], S=St[2]) + h.operand.buoyancy_perturbation = buoyancy(sb, grid, tracers) + compute!(h) + @test @allowscalar h[1, 1, 1] ≈ 9.295890287 # m + end +end diff --git a/test/test_jra55.jl b/test/test_jra55.jl index 8f83aa68..fd938839 100644 --- a/test/test_jra55.jl +++ b/test/test_jra55.jl @@ -37,7 +37,7 @@ using ClimaOcean.OceanSeaIceModels: PrescribedAtmosphere @info "Testing loading preprocessed JRA55 data on $A..." in_memory_jra55_fts = JRA55_field_time_series(test_name; - time_indices = 1:3, + time_indices = 1:4, architecture = arch, backend = InMemory(2)) From 07d4ab844271d33a99a76a944c4ae0c352857d9b Mon Sep 17 00:00:00 2001 From: "Gregory L. Wagner" Date: Mon, 16 Dec 2024 13:33:21 -0700 Subject: [PATCH 2/7] Change default major basins to 1 (#300) --- src/Bathymetry.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Bathymetry.jl b/src/Bathymetry.jl index ea9a358a..6173a40b 100644 --- a/src/Bathymetry.jl +++ b/src/Bathymetry.jl @@ -90,7 +90,7 @@ function regrid_bathymetry(target_grid; url = "https://www.ngdc.noaa.gov/thredds/fileServer/global/ETOPO2022/60s/60s_surface_elev_netcdf", filename = "ETOPO_2022_v1_60s_N90W180_surface.nc", interpolation_passes = 1, - major_basins = Inf) # Allow an `Inf` number of ``lakes'' + major_basins = 1) # Allow an `Inf` number of ``lakes'' filepath = joinpath(dir, filename) fileurl = url * "/" * filename # joinpath on windows creates the wrong url From b6b87573e9e0c29578231e4097a24a4ce09078b9 Mon Sep 17 00:00:00 2001 From: "Gregory L. Wagner" Date: Mon, 16 Dec 2024 14:16:47 -0700 Subject: [PATCH 3/7] Update Project.toml (#301) --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 812d3c92..3aca4f89 100644 --- a/Project.toml +++ b/Project.toml @@ -2,7 +2,7 @@ name = "ClimaOcean" uuid = "0376089a-ecfe-4b0e-a64f-9c555d74d754" license = "MIT" authors = ["Climate Modeling Alliance and contributors"] -version = "0.3.0" +version = "0.3.1" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" From b3ae3fef0c5512213f63c8efe0bc532a2ed57481 Mon Sep 17 00:00:00 2001 From: "Gregory L. Wagner" Date: Tue, 17 Dec 2024 22:46:31 -0700 Subject: [PATCH 4/7] Add a `biogeochemistry kwarg to ocean_simulation (#304) * Add a bgc kwarg to ocean_simulation * Add three degree simulation; * Update 3 degree simulation * Add capability to specify user boundary conditions * Better comment * Pretty up 3 degree * Updates --- .../three_degree_simulation.jl | 72 +++++++++++++++++++ src/OceanSimulations/OceanSimulations.jl | 30 +++++--- 2 files changed, 92 insertions(+), 10 deletions(-) create mode 100644 experiments/three_degree_simulation/three_degree_simulation.jl diff --git a/experiments/three_degree_simulation/three_degree_simulation.jl b/experiments/three_degree_simulation/three_degree_simulation.jl new file mode 100644 index 00000000..549a44e7 --- /dev/null +++ b/experiments/three_degree_simulation/three_degree_simulation.jl @@ -0,0 +1,72 @@ +using ClimaOcean +using OrthogonalSphericalShellGrids +using Oceananigans +using Oceananigans.Units +using CFTime +using Dates +using Printf + +arch = CPU() +Nx = 120 +Ny = 60 +Nz = 50 + +z_faces = exponential_z_faces(; Nz, depth=6000, h=34) + +underlying_grid = TripolarGrid(arch; size=(Nx, Ny, Nz), z=z_faces) +bottom_height = regrid_bathymetry(underlying_grid) +grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom_height)) + +gm = Oceananigans.TurbulenceClosures.IsopycnalSkewSymmetricDiffusivity(κ_skew=4000, κ_symmetric=4000) +catke = ClimaOcean.OceanSimulations.default_ocean_closure() +viscous_closure = Oceananigans.TurbulenceClosures.HorizontalScalarDiffusivity(ν=4000) + +dates = DateTimeProlepticGregorian(1993, 1, 1) : Month(1) : DateTimeProlepticGregorian(1993, 12, 1) +temperature = ECCOMetadata(:temperature, dates, ECCO4Monthly()) +salinity = ECCOMetadata(:salinity, dates, ECCO4Monthly()) + +restoring_rate = 1/2days +mask = LinearlyTaperedPolarMask(southern=(-80, -70), northern=(70, 90)) +FT = ECCORestoring(temperature, grid; mask, rate=restoring_rate) +FS = ECCORestoring(salinity, grid; mask, rate=restoring_rate) + +ocean = ocean_simulation(grid; + momentum_advection = VectorInvariant(), + tracer_advection = Centered(order=2), + closure = (gm, catke, viscous_closure), + forcing = (T=FT, S=FT), + tracers = (:T, :S, :e)) + +set!(ocean.model, T=ECCOMetadata(:temperature; dates=first(dates)), + S=ECCOMetadata(:salinity; dates=first(dates))) + +radiation = Radiation(arch) +atmosphere = JRA55PrescribedAtmosphere(arch; backend=JRA55NetCDFBackend(20)) +coupled_model = OceanSeaIceModel(ocean; atmosphere, radiation) +simulation = Simulation(coupled_model; Δt=20minutes, stop_time=30days) + +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[]) + + @info @sprintf("Time: %s, n: %d, Δt: %s, max|u|: (%.2e, %.2e, %.2e) m s⁻¹, \ + extrema(T): (%.2f, %.2f) ᵒC, wall time: %s \n", + prettytime(sim), iteration(sim), prettytime(sim.Δt), + umax..., Tmax, Tmin, prettytime(step_time)) + + wall_time[] = time_ns() + + return nothing +end + +add_callback!(simulation, progress, IterationInterval(10)) + +run!(simulation) + diff --git a/src/OceanSimulations/OceanSimulations.jl b/src/OceanSimulations/OceanSimulations.jl index 964ed611..9e9f2d71 100644 --- a/src/OceanSimulations/OceanSimulations.jl +++ b/src/OceanSimulations/OceanSimulations.jl @@ -73,6 +73,12 @@ default_tracer_advection() = FluxFormAdvection(WENO(order=7), @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) +function add_required_boundary_conditions(user_boundary_conditions, grid, bottom_drag_coefficient) + + return boundary_conditions +end + + # 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... function ocean_simulation(grid; @@ -85,10 +91,12 @@ function ocean_simulation(grid; gravitational_acceleration = g_Earth, bottom_drag_coefficient = Default(0.003), forcing = NamedTuple(), + biogeochemistry = nothing, timestepper = :QuasiAdamsBashforth2, coriolis = Default(HydrostaticSphericalCoriolis(; rotation_rate)), momentum_advection = default_momentum_advection(), equation_of_state = TEOS10EquationOfState(; reference_density), + boundary_conditions::NamedTuple = NamedTuple(), tracer_advection = default_tracer_advection(), verbose = false) @@ -136,7 +144,7 @@ function ocean_simulation(grid; end bottom_drag_coefficient = convert(FT, bottom_drag_coefficient) - + # Set up boundary conditions using Field top_zonal_momentum_flux = τx = Field{Face, Center, Nothing}(grid) top_meridional_momentum_flux = τy = Field{Center, Face, Nothing}(grid) @@ -148,18 +156,19 @@ 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) - ocean_boundary_conditions = (u = FieldBoundaryConditions(top=u_top_bc, bottom=u_bot_bc, immersed=u_immersed_bc), - v = FieldBoundaryConditions(top=v_top_bc, bottom=v_bot_bc, immersed=v_immersed_bc), - T = FieldBoundaryConditions(top=T_top_bc), - S = FieldBoundaryConditions(top=S_top_bc)) + default_boundary_conditions = (u = FieldBoundaryConditions(top=u_top_bc, bottom=u_bot_bc, immersed=u_immersed_bc), + v = FieldBoundaryConditions(top=v_top_bc, bottom=v_bot_bc, immersed=v_immersed_bc), + T = FieldBoundaryConditions(top=T_top_bc), + S = FieldBoundaryConditions(top=S_top_bc)) - # Use the TEOS10 equation of state - teos10 = TEOS10EquationOfState(; reference_density) - buoyancy = SeawaterBuoyancy(; gravitational_acceleration, equation_of_state=teos10) + # Merge boundary conditions with preference to user + # TODO: support users specifying only _part_ of the bcs for u, v, T, S (ie adding the top and immersed + # conditions even when a user-bc is supplied). + boundary_conditions = merge(default_boundary_conditions, boundary_conditions) buoyancy = SeawaterBuoyancy(; gravitational_acceleration, equation_of_state) @@ -183,6 +192,7 @@ function ocean_simulation(grid; ocean_model = HydrostaticFreeSurfaceModel(; grid, buoyancy, closure, + biogeochemistry, tracer_advection, momentum_advection, tracers, @@ -190,7 +200,7 @@ function ocean_simulation(grid; free_surface, coriolis, forcing, - boundary_conditions = ocean_boundary_conditions) + boundary_conditions) ocean = Simulation(ocean_model; Δt, verbose) From 03fd900bd6bcadbc3cbc40ddfb0f8fec2f701105 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Thu, 19 Dec 2024 21:50:34 +0100 Subject: [PATCH 5/7] Remove grid from `OceanSeaIceModel` (#308) * remove grid from OceanSeaIceModel * change to nameof * correct summary * show method * remove grid --- .../ocean_sea_ice_surface_fluxes.jl | 7 +++++++ src/OceanSeaIceModels/ocean_sea_ice_model.jl | 16 +++++++++------- 2 files changed, 16 insertions(+), 7 deletions(-) diff --git a/src/OceanSeaIceModels/CrossRealmFluxes/ocean_sea_ice_surface_fluxes.jl b/src/OceanSeaIceModels/CrossRealmFluxes/ocean_sea_ice_surface_fluxes.jl index 3d8fdb3a..e6d856f1 100644 --- a/src/OceanSeaIceModels/CrossRealmFluxes/ocean_sea_ice_surface_fluxes.jl +++ b/src/OceanSeaIceModels/CrossRealmFluxes/ocean_sea_ice_surface_fluxes.jl @@ -56,6 +56,13 @@ const celsius_to_kelvin = 273.15 Base.summary(crf::OceanSeaIceSurfaceFluxes) = "OceanSeaIceSurfaceFluxes" Base.show(io::IO, crf::OceanSeaIceSurfaceFluxes) = print(io, summary(crf)) +function Base.show(io::IO, crf::OceanSeaIceSurfaceFluxes) + print(io, summary(crf), "\n") + print(io, "├── radiation: ", summary(crf.radiation), "\n") + print(io, "└── turbulent coefficients: ", summary(crf.turbulent), "\n") + return nothing +end + const SeaIceSimulation = Simulation{<:SeaIceModel} function OceanSeaIceSurfaceFluxes(ocean, sea_ice=nothing; diff --git a/src/OceanSeaIceModels/ocean_sea_ice_model.jl b/src/OceanSeaIceModels/ocean_sea_ice_model.jl index f41bb304..efc2ebe7 100644 --- a/src/OceanSeaIceModels/ocean_sea_ice_model.jl +++ b/src/OceanSeaIceModels/ocean_sea_ice_model.jl @@ -7,6 +7,7 @@ using SeawaterPolynomials: TEOS10EquationOfState # Simulations interface 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 @@ -15,9 +16,8 @@ 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, G} <: AbstractModel{Nothing} +struct OceanSeaIceModel{I, A, O, F, C} <: AbstractModel{Nothing} clock :: C - grid :: G # TODO: make it so Oceananigans.Simulation does not require this atmosphere :: A sea_ice :: I ocean :: O @@ -27,9 +27,8 @@ end const OSIM = OceanSeaIceModel function Base.summary(model::OSIM) - A = nameof(typeof(architecture(model.grid))) - G = nameof(typeof(model.grid)) - return string("OceanSeaIceModel{$A, $G}", + A = nameof(typeof(architecture(model))) + return string("OceanSeaIceModel{$A}", "(time = ", prettytime(model.clock.time), ", iteration = ", model.clock.iteration, ")") end @@ -37,10 +36,14 @@ function Base.show(io::IO, cm::OSIM) print(io, summary(cm), "\n") print(io, "├── ocean: ", summary(cm.ocean.model), "\n") print(io, "├── atmosphere: ", summary(cm.atmosphere), "\n") - print(io, "└── sea_ice: ", summary(cm.sea_ice), "\n") + print(io, "├── sea_ice: ", summary(cm.sea_ice), "\n") + print(io, "└── fluxes: ", summary(cm.fluxes)) return nothing end +# Assumption: We have an ocean! +architecture(model::OSIM) = architecture(model.ocean) + prettytime(model::OSIM) = prettytime(model.clock.time) iteration(model::OSIM) = model.clock.iteration timestepper(::OSIM) = nothing @@ -98,7 +101,6 @@ function OceanSeaIceModel(ocean, sea_ice=FreezingLimitedOceanTemperature(); radiation) ocean_sea_ice_model = OceanSeaIceModel(clock, - ocean.model.grid, atmosphere, sea_ice, ocean, From ff4b8831403d9ebde929e7d6952b6dff5c0da578 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Thu, 19 Dec 2024 21:55:08 +0100 Subject: [PATCH 6/7] adding a free surface (#307) --- src/DataWrangling/ECCO/ECCO_metadata.jl | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/DataWrangling/ECCO/ECCO_metadata.jl b/src/DataWrangling/ECCO/ECCO_metadata.jl index b9f78ddb..7e3a9d99 100644 --- a/src/DataWrangling/ECCO/ECCO_metadata.jl +++ b/src/DataWrangling/ECCO/ECCO_metadata.jl @@ -168,6 +168,7 @@ ECCO4_short_names = Dict( :salinity => "SALT", :u_velocity => "EVEL", :v_velocity => "NVEL", + :free_surface => "SSH", :sea_ice_thickness => "SIheff", :sea_ice_area_fraction => "SIarea", :net_heat_flux => "oceQnet" @@ -178,6 +179,7 @@ ECCO2_short_names = Dict( :salinity => "SALT", :u_velocity => "UVEL", :v_velocity => "VVEL", + :free_surface => "SSH", :sea_ice_thickness => "SIheff", :sea_ice_area_fraction => "SIarea", :net_heat_flux => "oceQnet" @@ -186,6 +188,7 @@ ECCO2_short_names = Dict( ECCO_location = Dict( :temperature => (Center, Center, Center), :salinity => (Center, Center, Center), + :free_surface => (Center, Center, Nothing), :sea_ice_thickness => (Center, Center, Nothing), :sea_ice_area_fraction => (Center, Center, Nothing), :net_heat_flux => (Center, Center, Nothing), From d8992d5e0ab8882b13ef7aaf39182d264568ff9f Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Sat, 21 Dec 2024 11:56:15 +0100 Subject: [PATCH 7/7] Distributed regridding of bathymetry (#309) * distributed regridding * another strategy * fix imports * cannot share subarrays apparently * remember not to redownload * fix test * add corect tests * test more architectures * test other architectures * fix distributed tests --- src/Bathymetry.jl | 44 ++++++++++++++++++--- test/test_distributed_utils.jl | 71 ++++++++++++++++++++++++++++++++++ 2 files changed, 110 insertions(+), 5 deletions(-) diff --git a/src/Bathymetry.jl b/src/Bathymetry.jl index 6173a40b..bf9ee264 100644 --- a/src/Bathymetry.jl +++ b/src/Bathymetry.jl @@ -7,7 +7,7 @@ using ..DataWrangling: download_progress using Oceananigans using Oceananigans.Architectures: architecture, on_architecture -using Oceananigans.DistributedComputations: child_architecture +using Oceananigans.DistributedComputations: DistributedGrid, reconstruct_global_grid, barrier!, all_reduce using Oceananigans.Grids: halo_size, λnodes, φnodes using Oceananigans.Grids: x_domain, y_domain using Oceananigans.Grids: topology @@ -95,11 +95,12 @@ function regrid_bathymetry(target_grid; filepath = joinpath(dir, filename) fileurl = url * "/" * filename # joinpath on windows creates the wrong url - @root if !isfile(filepath) # perform all this only on rank 0, aka the "root" rank + # No need for @root here, because only rank 0 accesses this function + if !isfile(filepath) Downloads.download(fileurl, filepath; progress=download_progress) end - - dataset = Dataset(filepath) + + dataset = Dataset(filepath, "r") FT = eltype(target_grid) @@ -115,7 +116,7 @@ function regrid_bathymetry(target_grid; close(dataset) # Diagnose target grid information - arch = child_architecture(architecture(target_grid)) + arch = architecture(target_grid) φ₁, φ₂ = y_domain(target_grid) λ₁, λ₂ = x_domain(target_grid) @@ -247,6 +248,39 @@ function interpolate_bathymetry_in_passes(native_z, target_grid; return target_z end +# Regridding bathymetry for distributed grids, we handle the whole process +# on just one rank, and share the results with the other processors. +function regrid_bathymetry(target_grid::DistributedGrid; kw...) + global_grid = reconstruct_global_grid(target_grid) + global_grid = on_architecture(CPU(), global_grid) + arch = architecture(target_grid) + Nx, Ny, _ = size(global_grid) + + # If all ranks open a gigantic bathymetry and the memory is + # shared, we could easily have OOM errors. + # We perform the reconstruction only on rank 0 and share the result. + bottom_height = if arch.local_rank == 0 + bottom_field = regrid_bathymetry(global_grid; kw...) + bottom_field.data[1:Nx, 1:Ny, 1] + else + zeros(Nx, Ny) + end + + # Synchronize + ClimaOcean.global_barrier(arch.communicator) + + # Share the result (can we share SubArrays?) + bottom_height = all_reduce(+, bottom_height, arch) + + # Partition the result + local_bottom_height = Field{Center, Center, Nothing}(target_grid) + set!(local_bottom_height, bottom_height) + fill_halo_regions!(local_bottom_height) + + return local_bottom_height +end + + """ remove_minor_basins!(z_data, keep_major_basins) diff --git a/test/test_distributed_utils.jl b/test/test_distributed_utils.jl index db23d14f..8c790690 100644 --- a/test/test_distributed_utils.jl +++ b/test/test_distributed_utils.jl @@ -3,7 +3,9 @@ include("runtests_setup.jl") using MPI MPI.Init() +using NCDatasets using ClimaOcean.ECCO: download_dataset, metadata_path +using Oceananigans.DistributedComputations: reconstruct_global_grid using CFTime using Dates @@ -74,3 +76,72 @@ end @test isfile(metadata_path(metadatum)) end end + +@testset "Distributed Bathymetry interpolation" begin + # We start by building a fake bathyemtry on rank 0 and save it to file + @root begin + λ = -180:0.1:180 + φ = 0:0.1:50 + + Nλ = length(λ) + Nφ = length(φ) + + ds = NCDataset("./trivial_bathymetry.nc", "c") + + # Define the dimension "lon" and "lat" with the size 361 and 51 resp. + defDim(ds, "lon", Nλ) + defDim(ds, "lat", Nφ) + defVar(ds, "lat", Float32, ("lat", )) + defVar(ds, "lon", Float32, ("lon", )) + + # Define the variables z + z = defVar(ds, "z", Float32, ("lon","lat")) + + # Generate some example data + data = [Float32(-i) for i = 1:Nλ, j = 1:Nφ] + + # write a the complete data set + ds["lon"][:] = λ + ds["lat"][:] = φ + z[:,:] = data + + close(ds) + end + + global_grid = LatitudeLongitudeGrid(CPU(); + size = (40, 40, 1), + longitude = (0, 100), + latitude = (0, 20), + z = (0, 1)) + + global_height = regrid_bathymetry(global_grid; + dir = "./", + filename = "trivial_bathymetry.nc", + interpolation_passes=10) + + arch_x = Distributed(CPU(), partition=Partition(4, 1)) + arch_y = Distributed(CPU(), partition=Partition(1, 4)) + arch_xy = Distributed(CPU(), partition=Partition(2, 2)) + + for arch in (arch_x, arch_y, arch_xy) + local_grid = LatitudeLongitudeGrid(arch; + size = (40, 40, 1), + longitude = (0, 100), + latitude = (0, 20), + z = (0, 1)) + + local_height = regrid_bathymetry(local_grid; + dir = "./", + filename = "trivial_bathymetry.nc", + interpolation_passes=10) + + Nx, Ny, _ = size(local_grid) + rx, ry, _ = arch.local_index + irange = (rx - 1) * Nx + 1 : rx * Nx + jrange = (ry - 1) * Ny + 1 : ry * Ny + + @handshake begin + @test interior(global_height, irange, jrange, 1) == interior(local_height, :, :, 1) + end + end +end \ No newline at end of file