From 8878541696aa0dadb587c2517b553b6447b60f7f Mon Sep 17 00:00:00 2001 From: "Gregory L. Wagner" Date: Mon, 16 Dec 2024 09:14:02 -0700 Subject: [PATCH] Add parameter to epsilon equation so buoyancy flux term may always be positive (#3886) * Add parameter to epsilon equation so buoyancy flux term may always be positive * Fix on_architecture for ensemble column grids * Bugfix in ColumnEnsembleSize in on_architecture * Fix formatting * Add test for ColumnEnsembleSize --- src/Grids/rectilinear_grid.jl | 30 ++++++++++++++++--- .../single_column_model_mode.jl | 14 +-------- src/OutputReaders/field_time_series.jl | 3 +- .../tke_dissipation_equations.jl | 9 ++++-- test/test_grids.jl | 17 ++++++++++- 5 files changed, 52 insertions(+), 21 deletions(-) diff --git a/src/Grids/rectilinear_grid.jl b/src/Grids/rectilinear_grid.jl index f012b9ab1f..9b1f9ad46a 100644 --- a/src/Grids/rectilinear_grid.jl +++ b/src/Grids/rectilinear_grid.jl @@ -280,6 +280,8 @@ function RectilinearGrid(architecture::AbstractArchitecture = CPU(), z) end + + """ Validate user input arguments to the `RectilinearGrid` constructor. """ function validate_rectilinear_grid_args(topology, size, halo, FT, extent, x, y, z) TX, TY, TZ = topology = validate_topology(topology) @@ -338,6 +340,20 @@ function Base.show(io::IO, grid::RectilinearGrid, withsummary=true) "└── ", z_summary) end +##### +##### For "column ensemble models" +##### + +struct ColumnEnsembleSize{C<:Tuple{Int, Int}} + ensemble :: C + Nz :: Int + Hz :: Int +end + +ColumnEnsembleSize(; Nz, ensemble=(0, 0), Hz=1) = ColumnEnsembleSize(ensemble, Nz, Hz) +validate_size(TX, TY, TZ, e::ColumnEnsembleSize) = tuple(e.ensemble[1], e.ensemble[2], e.Nz) +validate_halo(TX, TY, TZ, size, e::ColumnEnsembleSize) = tuple(0, 0, e.Hz) + ##### ##### Utilities ##### @@ -369,10 +385,16 @@ function constructor_arguments(grid::RectilinearGrid) # Kwargs topo = topology(grid) - size = (grid.Nx, grid.Ny, grid.Nz) - halo = (grid.Hx, grid.Hy, grid.Hz) - size = pop_flat_elements(size, topo) - halo = pop_flat_elements(halo, topo) + + if (topo[1] == Flat && grid.Nx > 1) || + (topo[2] == Flat && grid.Ny > 1) + size = halo = ColumnEnsembleSize(Nz=grid.Nz, Hz=grid.Hz, ensemble=(grid.Nx, grid.Ny)) + else + size = (grid.Nx, grid.Ny, grid.Nz) + halo = (grid.Hx, grid.Hy, grid.Hz) + size = pop_flat_elements(size, topo) + halo = pop_flat_elements(halo, topo) + end kwargs = Dict(:size => size, :halo => halo, diff --git a/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl b/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl index 7d32d872e3..5e3b765cb8 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/single_column_model_mode.jl @@ -2,13 +2,12 @@ using CUDA: @allowscalar using Oceananigans: UpdateStateCallsite using Oceananigans.Advection: AbstractAdvectionScheme -using Oceananigans.Grids: Flat, Bounded +using Oceananigans.Grids: Flat, Bounded, ColumnEnsembleSize using Oceananigans.Fields: ZeroField using Oceananigans.Coriolis: AbstractRotation using Oceananigans.TurbulenceClosures: AbstractTurbulenceClosure using Oceananigans.TurbulenceClosures.TKEBasedVerticalDiffusivities: CATKEVDArray -import Oceananigans.Grids: validate_size, validate_halo import Oceananigans.Models: validate_tracer_advection import Oceananigans.BoundaryConditions: fill_halo_regions! import Oceananigans.TurbulenceClosures: time_discretization, compute_diffusivities! @@ -100,17 +99,6 @@ end return ∇_dot_qᶜ(i, j, k, grid, closure, c, tracer_index, args...) end -struct ColumnEnsembleSize{C<:Tuple{Int, Int}} - ensemble :: C - Nz :: Int - Hz :: Int -end - -ColumnEnsembleSize(; Nz, ensemble=(0, 0), Hz=1) = ColumnEnsembleSize(ensemble, Nz, Hz) - -validate_size(TX, TY, TZ, e::ColumnEnsembleSize) = tuple(e.ensemble[1], e.ensemble[2], e.Nz) -validate_halo(TX, TY, TZ, size, e::ColumnEnsembleSize) = tuple(0, 0, e.Hz) - @inline function time_discretization(closure_array::AbstractArray) first_closure = @allowscalar first(closure_array) # assumes all closures have same time-discretization return time_discretization(first_closure) diff --git a/src/OutputReaders/field_time_series.jl b/src/OutputReaders/field_time_series.jl index afe9439e38..48ec492a8b 100644 --- a/src/OutputReaders/field_time_series.jl +++ b/src/OutputReaders/field_time_series.jl @@ -351,7 +351,8 @@ new_data(FT, grid, loc, indices, ::Nothing) = nothing # Apparently, not explicitly specifying Int64 in here makes this function # fail on x86 processors where `Int` is implied to be `Int32` -# see ClimaOcean commit 3c47d887659d81e0caed6c9df41b7438e1f1cd52 at https://github.com/CliMA/ClimaOcean.jl/actions/runs/8804916198/job/24166354095) +# see ClimaOcean commit 3c47d887659d81e0caed6c9df41b7438e1f1cd52 at +# https://github.com/CliMA/ClimaOcean.jl/actions/runs/8804916198/job/24166354095) function new_data(FT, grid, loc, indices, Nt::Union{Int, Int64}) space_size = total_size(grid, loc, indices) underlying_data = zeros(FT, architecture(grid), space_size..., Nt) diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/tke_dissipation_equations.jl b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/tke_dissipation_equations.jl index d24bb8d84b..ab10f29cbe 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/tke_dissipation_equations.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/tke_dissipation_equations.jl @@ -10,7 +10,8 @@ using CUDA Base.@kwdef struct TKEDissipationEquations{FT} Cᵋϵ :: FT = 1.92 Cᴾϵ :: FT = 1.44 - Cᵇϵ :: FT = -0.65 + Cᵇϵ⁺ :: FT = -0.65 + Cᵇϵ⁻ :: FT = -0.65 Cᵂu★ :: FT = 0.0 CᵂwΔ :: FT = 0.0 Cᵂα :: FT = 0.11 # Charnock parameter @@ -133,7 +134,11 @@ end # Patankar trick for ϵ-equation Cᵋϵ = closure_ij.tke_dissipation_equations.Cᵋϵ - Cᵇϵ = closure_ij.tke_dissipation_equations.Cᵇϵ + Cᵇϵ⁺ = closure_ij.tke_dissipation_equations.Cᵇϵ⁺ + Cᵇϵ⁻ = closure_ij.tke_dissipation_equations.Cᵇϵ⁻ + + N² = ℑzᵃᵃᶜ(i, j, k, grid, ∂z_b, buoyancy, tracers) + Cᵇϵ = ifelse(N² ≥ 0, Cᵇϵ⁺, Cᵇϵ⁻) Cᵇϵ_wb⁻ = min(Cᵇϵ * wb, zero(grid)) Cᵇϵ_wb⁺ = max(Cᵇϵ * wb, zero(grid)) diff --git a/test/test_grids.jl b/test/test_grids.jl index 000621f0f9..5ba9e5ca97 100644 --- a/test/test_grids.jl +++ b/test/test_grids.jl @@ -1,7 +1,7 @@ include("dependencies_for_runtests.jl") include("data_dependencies.jl") -using Oceananigans.Grids: total_extent, +using Oceananigans.Grids: total_extent, ColumnEnsembleSize, xspacings, yspacings, zspacings, xnode, ynode, znode, λnode, φnode, λspacings, φspacings @@ -997,6 +997,21 @@ end @test φ[1] isa FT @test λ[1] == λ₀ @test φ[1] == convert(FT, φ₀) + + halo_sz = ColumnEnsembleSize(; Nz=4, ensemble=(2, 3), Hz=5) + grid = RectilinearGrid(arch, FT; size=halo_sz, halo=halo_sz, + z=(-10, 0), topology=(Flat, Flat, Bounded)) + @test size(grid) == (2, 3, 4) + @test halo_size(grid) == (0, 0, 5) + + if arch isa GPU + cpu_grid = on_architecture(CPU(), grid) + @test size(cpu_grid) == (2, 3, 4) + @test halo_size(cpu_grid) == (0, 0, 5) + + gpu_grid = on_architecture(GPU(), grid) + @test gpu_grid == grid + end end end end