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

@@ -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,

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)

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)

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))

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