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/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/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/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 diff --git a/src/ClimaOcean.jl b/src/ClimaOcean.jl index c7091a68..7daf1d6e 100644 --- a/src/ClimaOcean.jl +++ b/src/ClimaOcean.jl @@ -79,7 +79,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") include("SeaIceSimulations/SeaIceSimulations.jl") diff --git a/src/DataWrangling/ECCO/ECCO_metadata.jl b/src/DataWrangling/ECCO/ECCO_metadata.jl index 6595ff34..cdc458ba 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_concentration => "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_concentration => "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_concentration => (Center, Center, Nothing), :net_heat_flux => (Center, Center, Nothing), 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/src/OceanSeaIceModels/CrossRealmFluxes/cross_realm_surface_fluxes.jl b/src/OceanSeaIceModels/CrossRealmFluxes/cross_realm_surface_fluxes.jl index cdf43a4a..15037806 100644 --- a/src/OceanSeaIceModels/CrossRealmFluxes/cross_realm_surface_fluxes.jl +++ b/src/OceanSeaIceModels/CrossRealmFluxes/cross_realm_surface_fluxes.jl @@ -30,7 +30,10 @@ using KernelAbstractions: @kernel, @index ##### struct CrossRealmSurfaceFluxes{T, P, C, R, PI, PC, FT, UN, ATM} - turbulent :: T # the turbulent fluxes + turbulent :: T # turbulent atmopshere fluxes + # atmosphere_ocean + # atmosphere_sea_ice + # sea_ice_ocean prescribed :: P # Add `components` which will also store components of the total fluxes # (eg latent, sensible heat flux) diff --git a/src/OceanSeaIceModels/CrossRealmFluxes/ocean_sea_ice_surface_fluxes.jl b/src/OceanSeaIceModels/CrossRealmFluxes/ocean_sea_ice_surface_fluxes.jl new file mode 100644 index 00000000..e6d856f1 --- /dev/null +++ b/src/OceanSeaIceModels/CrossRealmFluxes/ocean_sea_ice_surface_fluxes.jl @@ -0,0 +1,199 @@ +using StaticArrays +using Thermodynamics +using SurfaceFluxes + +using ..OceanSeaIceModels: reference_density, + heat_capacity, + sea_ice_concentration, + downwelling_radiation, + freshwater_flux + +using ClimaSeaIce: SeaIceModel + +using Oceananigans: HydrostaticFreeSurfaceModel, architecture +using Oceananigans.Grids: inactive_node, node +using Oceananigans.BoundaryConditions: fill_halo_regions! +using Oceananigans.Fields: ConstantField, interpolate +using Oceananigans.Utils: launch!, Time, KernelParameters + +# using Oceananigans.OutputReaders: extract_field_time_series, update_field_time_series! + +using Oceananigans.Operators: ℑxᶜᵃᵃ, ℑyᵃᶜᵃ, ℑxᶠᵃᵃ, ℑyᵃᶠᵃ + +using KernelAbstractions: @kernel, @index + +##### +##### Container for organizing information related to fluxes +##### + +struct OceanSeaIceSurfaceFluxes{T, P, C, R, PI, PC, FT, UN, ATM} + turbulent :: T + prescribed :: P + # Add `components` which will also store components of the total fluxes + # (eg latent, sensible heat flux) + total :: C + radiation :: R + previous_ice_thickness :: PI + previous_ice_concentration :: PC + # The ocean is Boussinesq, so these are _only_ coupled properties: + ocean_reference_density :: FT + ocean_heat_capacity :: FT + freshwater_density :: FT + ocean_temperature_units :: UN + # Scratch space to store the atmosphere state at the surface + # interpolated to the ocean grid + surface_atmosphere_state :: ATM +end + +# Possible units for temperature and salinity +struct DegreesCelsius end +struct DegreesKelvin end + +const celsius_to_kelvin = 273.15 +@inline convert_to_kelvin(::DegreesCelsius, T::FT) where FT = T + convert(FT, celsius_to_kelvin) +@inline convert_to_kelvin(::DegreesKelvin, T) = T + +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; + atmosphere = nothing, + radiation = nothing, + freshwater_density = 1000, + ocean_temperature_units = DegreesCelsius(), + similarity_theory = nothing, + ocean_reference_density = reference_density(ocean), + ocean_heat_capacity = heat_capacity(ocean)) + + ocean_grid = ocean.model.grid + FT = eltype(ocean_grid) + + ocean_reference_density = convert(FT, ocean_reference_density) + ocean_heat_capacity = convert(FT, ocean_heat_capacity) + freshwater_density = convert(FT, freshwater_density) + + if !isnothing(atmosphere) + # It's the "thermodynamics gravitational acceleration" + # (as opposed to the one used for the free surface) + gravitational_acceleration = ocean.model.buoyancy.formulation.gravitational_acceleration + + if isnothing(similarity_theory) + similarity_theory = SimilarityTheoryTurbulentFluxes(ocean_grid; gravitational_acceleration) + end + end + + prescribed_fluxes = nothing + + if sea_ice isa SeaIceSimulation + previous_ice_thickness = deepcopy(sea_ice.model.ice_thickness) + previous_ice_concentration = deepcopy(sea_ice.model.ice_concentration) + else + previous_ice_thickness = nothing + previous_ice_concentration = nothing + end + + total_ocean_fluxes = ocean_model_fluxes(ocean.model, + ocean_reference_density, + ocean_heat_capacity) + + total_fluxes = (; ocean=total_ocean_fluxes) + + surface_atmosphere_state = interpolated_surface_atmosphere_state(ocean_grid) + + return OceanSeaIceSurfaceFluxes(similarity_theory, + prescribed_fluxes, + total_fluxes, + radiation, + previous_ice_thickness, + previous_ice_concentration, + ocean_reference_density, + ocean_heat_capacity, + freshwater_density, + ocean_temperature_units, + surface_atmosphere_state) +end + +function ocean_model_fluxes(model, ρₒ, cₚ) + grid = model.grid + τx = surface_flux(model.velocities.u) + τy = surface_flux(model.velocities.v) + τxᶜᶜᶜ = Field{Center, Center, Nothing}(grid) + τyᶜᶜᶜ = Field{Center, Center, Nothing}(grid) + + ocean_momentum_fluxes = (u = τx, # fluxes used in the model + v = τy, # + # Including these (which are only a user convenience, not needed for + # time-stepping) incurs about 100s in construction + # time for OceanSeaIceSurfaceFluxes: + # ρτx = ρₒ * τx, # momentum fluxes multiplied by reference density + # ρτy = ρₒ * τy, # user convenience + uᶜᶜᶜ = τxᶜᶜᶜ, # fluxes computed by bulk formula at cell centers + vᶜᶜᶜ = τyᶜᶜᶜ) + + tracers = model.tracers + ocean_tracer_fluxes = NamedTuple(name => surface_flux(tracers[name]) + for name in keys(tracers)) + + ocean_heat_flux = ρₒ * cₚ * ocean_tracer_fluxes.T + + fluxes = (momentum = ocean_momentum_fluxes, + tracers = ocean_tracer_fluxes, + heat = ocean_heat_flux) + + return fluxes +end + +function interpolated_surface_atmosphere_state(ocean_grid) + surface_atmosphere_state = (u = Field{Center, Center, Nothing}(ocean_grid), + v = Field{Center, Center, Nothing}(ocean_grid), + T = Field{Center, Center, Nothing}(ocean_grid), + q = Field{Center, Center, Nothing}(ocean_grid), + p = Field{Center, Center, Nothing}(ocean_grid), + Qs = Field{Center, Center, Nothing}(ocean_grid), + Qℓ = Field{Center, Center, Nothing}(ocean_grid), + Mp = Field{Center, Center, Nothing}(ocean_grid)) + + return surface_atmosphere_state +end + + +##### +##### Utility for interpolating tuples of fields +##### + +# Note: assumes loc = (c, c, nothing) (and the third location should +# not matter.) +@inline interp_atmos_time_series(J, X, time, grid, args...) = + interpolate(X, time, J, (c, c, nothing), grid, args...) + +@inline interp_atmos_time_series(ΣJ::NamedTuple, args...) = + interp_atmos_time_series(values(ΣJ), args...) + +@inline interp_atmos_time_series(ΣJ::Tuple{<:Any}, args...) = + interp_atmos_time_series(ΣJ[1], args...) + + interp_atmos_time_series(ΣJ[2], args...) + +@inline interp_atmos_time_series(ΣJ::Tuple{<:Any, <:Any}, args...) = + interp_atmos_time_series(ΣJ[1], args...) + + interp_atmos_time_series(ΣJ[2], args...) + +@inline interp_atmos_time_series(ΣJ::Tuple{<:Any, <:Any, <:Any}, args...) = + interp_atmos_time_series(ΣJ[1], args...) + + interp_atmos_time_series(ΣJ[2], args...) + + interp_atmos_time_series(ΣJ[3], args...) + +@inline interp_atmos_time_series(ΣJ::Tuple{<:Any, <:Any, <:Any, <:Any}, args...) = + interp_atmos_time_series(ΣJ[1], args...) + + interp_atmos_time_series(ΣJ[2], args...) + + interp_atmos_time_series(ΣJ[3], args...) + + interp_atmos_time_series(ΣJ[4], args...) + diff --git a/src/OceanSeaIceModels/ocean_sea_ice_model.jl b/src/OceanSeaIceModels/ocean_sea_ice_model.jl index 2730110f..9f79626f 100644 --- a/src/OceanSeaIceModels/ocean_sea_ice_model.jl +++ b/src/OceanSeaIceModels/ocean_sea_ice_model.jl @@ -11,6 +11,7 @@ using ClimaOcean.OceanSeaIceModels.CrossRealmFluxes: ClasiusClapyeronSaturation # 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 @@ -19,9 +20,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 @@ -31,9 +31,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 @@ -41,10 +40,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 @@ -130,7 +133,6 @@ function OceanSeaIceModel(ocean, sea_ice=FreezingLimitedOceanTemperature(); radiation) ocean_sea_ice_model = OceanSeaIceModel(clock, - ocean.model.grid, atmosphere, sea_ice, ocean, 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) 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))