(0.91.1) CATKE numerics: discretely correct shear production, split-e…

…xplicit TKE substepping (#3585)

* Generalize definition of Ri

* Non-negative Ri

* Use quasi-conservative reconstruction of shear production and buoyancy flux

* Add convective mixing length for momentum

* Update minimum TKE

* Major reorganizations to get CATKE time-stepping to work

* Working on comparison

* Some updates

* Heroics to get things to wokrk

* Update tke time stepping

* Bit of cleanup

* Figure out the right temporal stencil

* Updated default convective buoyancy flux and delete ccc methods

* Add back convective mixing length for ccc

* Grr also add back stability function for ccc

* Better error message when setting fails for field time series

* Better notation for surface buoyancy flux;

* Implement split-explicit time-stepping for CATKE

* Optimize CATKE substepping a bit more

* Performance enhancements

* Add third method to CATKE comparison

* Minor improvements

* Clean up

* Update

* Move around some inbounds

* Delete unused code

* Clean up time stepping

* Little things

* Improve show and consistently use surface-ignoring reconstruction

* Update definition of Ri and try a new stability funciton

* Delete show

* Queue up change to sheared convection

* Update scale

* Another tweak

* Add four more parameters for shear turbulence in unstable stratification (#3600)

* Add special negative Ri branch to the stability function

* Update parameters

* Improve show for CATKE

* Skip test for tuples-with-CATKE for now

* Try to fix checkpointing / regression test issues

* Its a special test so we have to comment it out

* New parameters

* Fixing missing word

* Address review comments

* delete unused FT

* Update Project.toml

---------

Co-authored-by: Navid C. Constantinou <[email protected]>

Project.toml

@@ -1,7 +1,7 @@
 name = "Oceananigans"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 authors = ["Climate Modeling Alliance and contributors"]
-version = "0.91.0"
+version = "0.91.1"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/Grids/Grids.jl

@@ -155,6 +155,8 @@ const YZFlatGrid = AbstractGrid{<:Any, <:Any, Flat, Flat}
 const XYZFlatGrid = AbstractGrid{<:Any, Flat, Flat, Flat}
 
 isrectilinear(grid) = false
+@inline active_surface_map(::AbstractGrid) = nothing
+@inline active_interior_map(::AbstractGrid) = nothing
 
 include("grid_utils.jl")
 include("nodes_and_spacings.jl")

src/ImmersedBoundaries/active_cells_map.jl

@@ -4,6 +4,7 @@ using Oceananigans.Grids: AbstractGrid
 
 using KernelAbstractions: @kernel, @index
 
+import Oceananigans.Grids: active_surface_map, active_interior_map
 import Oceananigans.Utils: active_cells_work_layout
 
 using Oceananigans.Solvers: solve_batched_tridiagonal_system_z!, ZDirection
@@ -33,15 +34,13 @@ struct EastMap  end
 struct SouthMap end
 struct NorthMap end
 
-@inline active_surface_map(::AbstractGrid)      = nothing
 @inline active_surface_map(::ActiveZColumnsIBG) = ZColumnMap()
 
 @inline active_interior_map(::Val{:west})  = WestMap()
 @inline active_interior_map(::Val{:east})  = EastMap()
 @inline active_interior_map(::Val{:south}) = SouthMap()
 @inline active_interior_map(::Val{:north}) = NorthMap()
 
-@inline active_interior_map(::AbstractGrid)              = nothing
 @inline active_interior_map(::ActiveCellsIBG)            = InteriorMap()
 @inline active_interior_map(::DistributedActiveCellsIBG) = InteriorMap()
 

src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_boundary_tendencies.jl

@@ -1,16 +1,16 @@
 import Oceananigans.Models: compute_boundary_tendencies!
+import Oceananigans.Models: compute_boundary_tendencies!
+
 using Oceananigans.Grids: halo_size
-using Oceananigans.TurbulenceClosures: required_halo_size
+using Oceananigans.ImmersedBoundaries: active_interior_map, DistributedActiveCellsIBG
 using Oceananigans.Models.NonhydrostaticModels: boundary_tendency_kernel_parameters,
                                                 boundary_p_kernel_parameters, 
                                                 boundary_κ_kernel_parameters,
                                                 boundary_parameters
 
-import Oceananigans.Models: compute_boundary_tendencies!
-
-using Oceananigans.ImmersedBoundaries: active_interior_map, DistributedActiveCellsIBG
+using Oceananigans.TurbulenceClosures: required_halo_size
 
-# We assume here that top/bottom BC are always synched (no partitioning in z)
+# We assume here that top/bottom BC are always synchronized (no partitioning in z)
 function compute_boundary_tendencies!(model::HydrostaticFreeSurfaceModel)
     grid = model.grid
     arch = architecture(grid)
@@ -71,3 +71,4 @@ function boundary_w_kernel_parameters(grid, arch)
 
     return boundary_parameters(sizes, offs, grid, arch)
 end
+

src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl

@@ -1,17 +1,14 @@
-import Oceananigans.TimeSteppers: compute_tendencies!
 import Oceananigans: tracer_tendency_kernel_function
+import Oceananigans.TimeSteppers: compute_tendencies!
+import Oceananigans.Models: complete_communication_and_compute_boundary!
+import Oceananigans.Models: interior_tendency_kernel_parameters
 
+using Oceananigans: fields, prognostic_fields, TendencyCallsite, UpdateStateCallsite
 using Oceananigans.Utils: work_layout, KernelParameters
-using Oceananigans.Fields: immersed_boundary_condition
 using Oceananigans.Grids: halo_size
-using Oceananigans: fields, prognostic_fields, TendencyCallsite, UpdateStateCallsite
+using Oceananigans.Fields: immersed_boundary_condition
 using Oceananigans.Biogeochemistry: update_tendencies!
-
-import Oceananigans.TimeSteppers: compute_tendencies!
-import Oceananigans: tracer_tendency_kernel_function
-
-import Oceananigans.Models: complete_communication_and_compute_boundary!
-import Oceananigans.Models: interior_tendency_kernel_parameters
+using Oceananigans.TurbulenceClosures.CATKEVerticalDiffusivities: FlavorOfCATKE
 
 using Oceananigans.ImmersedBoundaries: active_interior_map, ActiveCellsIBG, 
                                        InteriorMap, active_linear_index_to_tuple
@@ -54,23 +51,6 @@ function compute_tendencies!(model::HydrostaticFreeSurfaceModel, callbacks)
     return nothing
 end
 
-using Oceananigans.TurbulenceClosures.CATKEVerticalDiffusivities: FlavorOfCATKE
-
-@inline tracer_tendency_kernel_function(model::HFSM, name, c, K)                     = compute_hydrostatic_free_surface_Gc!, c, K
-@inline tracer_tendency_kernel_function(model::HFSM, ::Val{:e}, c::FlavorOfCATKE, K) = compute_hydrostatic_free_surface_Ge!, c, K
-
-function tracer_tendency_kernel_function(model::HFSM, ::Val{:e}, closures::Tuple, diffusivity_fields::Tuple)
-    catke_index = findfirst(c -> c isa FlavorOfCATKE, closures)
-
-    if isnothing(catke_index)
-        return compute_hydrostatic_free_surface_Gc!, closures, diffusivity_fields
-    else
-        catke_closure = closures[catke_index]
-        catke_diffusivity_fields = diffusivity_fields[catke_index]
-        return compute_hydrostatic_free_surface_Ge!, catke_closure, catke_diffusivity_fields 
-    end
-end
-
 @inline function top_tracer_boundary_conditions(grid, tracers)
     names = propertynames(tracers)
     values = Tuple(tracers[c].boundary_conditions.top for c in names)
@@ -87,38 +67,31 @@ function compute_hydrostatic_free_surface_tendency_contributions!(model, kernel_
 
     compute_hydrostatic_momentum_tendencies!(model, model.velocities, kernel_parameters; active_cells_map)
 
-    top_tracer_bcs = top_tracer_boundary_conditions(grid, model.tracers)
-
     for (tracer_index, tracer_name) in enumerate(propertynames(model.tracers))
+
         @inbounds c_tendency    = model.timestepper.Gⁿ[tracer_name]
         @inbounds c_advection   = model.advection[tracer_name]
         @inbounds c_forcing     = model.forcing[tracer_name]
         @inbounds c_immersed_bc = immersed_boundary_condition(model.tracers[tracer_name])
 
-        tendency_kernel!, closure, diffusivity = tracer_tendency_kernel_function(model,
-                                                                                 Val(tracer_name),
-                                                                                 model.closure,
-                                                                                 model.diffusivity_fields)
-
         args = tuple(Val(tracer_index),
                      Val(tracer_name),
                      c_advection,
-                     closure,
+                     model.closure,
                      c_immersed_bc,
                      model.buoyancy,
                      model.biogeochemistry,
                      model.velocities,
                      model.free_surface,
                      model.tracers,
-                     top_tracer_bcs,
-                     diffusivity,
+                     model.diffusivity_fields,
                      model.auxiliary_fields,
                      c_forcing,
                      model.clock)
 
         for parameters in kernel_parameters
             launch!(arch, grid, parameters,
-                    tendency_kernel!,
+                    compute_hydrostatic_free_surface_Gc!,
                     c_tendency,
                     grid,
                     active_cells_map,
@@ -267,18 +240,6 @@ end
     @inbounds Gc[i, j, k] = hydrostatic_free_surface_tracer_tendency(i, j, k, grid, args...)
 end
 
-""" Calculate the right-hand-side of the subgrid scale energy equation. """
-@kernel function compute_hydrostatic_free_surface_Ge!(Ge, grid, map, args)
-    i, j, k = @index(Global, NTuple)
-    @inbounds Ge[i, j, k] = hydrostatic_turbulent_kinetic_energy_tendency(i, j, k, grid, args...)
-end
-
-@kernel function compute_hydrostatic_free_surface_Ge!(Ge, grid::ActiveCellsIBG, map, args)
-    idx = @index(Global, Linear)
-    i, j, k = active_linear_index_to_tuple(idx, map, grid)
-    @inbounds Ge[i, j, k] = hydrostatic_turbulent_kinetic_energy_tendency(i, j, k, grid, args...)
-end
-
 #####
 ##### Tendency calculators for an explicit free surface
 #####

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl

@@ -11,8 +11,9 @@ import Oceananigans.TimeSteppers: ab2_step!
 
 setup_free_surface!(model, free_surface, χ) = nothing
 
-function ab2_step!(model::HydrostaticFreeSurfaceModel, Δt, χ)
+function ab2_step!(model::HydrostaticFreeSurfaceModel, Δt)
 
+    χ = model.timestepper.χ
     setup_free_surface!(model, model.free_surface, χ)
 
     # Step locally velocity and tracers
@@ -69,24 +70,33 @@ ab2_step_tracers!(::EmptyNamedTuple, model, Δt, χ) = nothing
 
 function ab2_step_tracers!(tracers, model, Δt, χ)
 
+    closure = model.closure
+
     # Tracer update kernels
     for (tracer_index, tracer_name) in enumerate(propertynames(tracers))
-        Gⁿ = model.timestepper.Gⁿ[tracer_name]
-        G⁻ = model.timestepper.G⁻[tracer_name]
-        tracer_field = tracers[tracer_name]
-        closure = model.closure
-
-        launch!(model.architecture, model.grid, :xyz,
-                ab2_step_field!, tracer_field, Δt, χ, Gⁿ, G⁻)
-
-        implicit_step!(tracer_field,
-                       model.timestepper.implicit_solver,
-                       closure,
-                       model.diffusivity_fields,
-                       Val(tracer_index),
-                       model.clock,
-                       Δt)
+
+        # TODO: do better than this silly criteria, also need to check closure tuples
+        if closure isa FlavorOfCATKE && tracer_name == :e
+            @debug "Skipping AB2 step for e"
+        else
+            Gⁿ = model.timestepper.Gⁿ[tracer_name]
+            G⁻ = model.timestepper.G⁻[tracer_name]
+            tracer_field = tracers[tracer_name]
+            closure = model.closure
+
+            launch!(model.architecture, model.grid, :xyz,
+                    ab2_step_field!, tracer_field, Δt, χ, Gⁿ, G⁻)
+
+            implicit_step!(tracer_field,
+                           model.timestepper.implicit_solver,
+                           closure,
+                           model.diffusivity_fields,
+                           Val(tracer_index),
+                           model.clock,
+                           Δt)
+        end
     end
 
     return nothing
 end
+

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl

@@ -156,7 +156,7 @@ function HydrostaticFreeSurfaceModel(; grid,
     boundary_conditions = regularize_field_boundary_conditions(boundary_conditions, grid, prognostic_field_names)
 
     # Finally, we ensure that closure-specific boundary conditions, such as
-    # those required by TKEBasedVerticalDiffusivity, are enforced:
+    # those required by CATKEVerticalDiffusivity, are enforced:
     boundary_conditions = add_closure_specific_boundary_conditions(closure, boundary_conditions, grid, tracernames(tracers), buoyancy)
 
     # Ensure `closure` describes all tracers

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl

@@ -114,7 +114,6 @@ where `c = C[tracer_index]`.
                                                           velocities,
                                                           free_surface,
                                                           tracers,
-                                                          top_tracer_bcs,
                                                           diffusivities,
                                                           auxiliary_fields,
                                                           forcing,
@@ -162,31 +161,3 @@ The tendency is called ``G_η`` and defined via
                        + forcings.η(i, j, k_top, grid, clock, model_fields))
 end
 
-@inline function hydrostatic_turbulent_kinetic_energy_tendency(i, j, k, grid,
-                                                               val_tracer_index::Val{tracer_index},
-                                                               val_tracer_name,
-                                                               advection,
-                                                               closure,
-                                                               e_immersed_bc,
-                                                               buoyancy,
-                                                               biogeochemistry,
-                                                               velocities,
-                                                               free_surface,
-                                                               tracers,
-                                                               top_tracer_bcs,
-                                                               diffusivities,
-                                                               auxiliary_fields,
-                                                               forcing,
-                                                               clock) where tracer_index
-
-    @inbounds e = tracers[tracer_index]
-    model_fields = merge(hydrostatic_fields(velocities, free_surface, tracers), auxiliary_fields)
-
-    return ( - div_Uc(i, j, k, grid, advection, velocities, e)
-             - ∇_dot_qᶜ(i, j, k, grid, closure, diffusivities, val_tracer_index, e, clock, model_fields, buoyancy)
-             - immersed_∇_dot_qᶜ(i, j, k, grid, e, e_immersed_bc, closure, diffusivities, val_tracer_index, clock, model_fields)
-             + shear_production(i, j, k, grid, closure, velocities, tracers, buoyancy, diffusivities)
-             + buoyancy_flux(i, j, k, grid, closure, velocities, tracers, buoyancy, diffusivities)
-             - dissipation(i, j, k, grid, closure, velocities, tracers, buoyancy, diffusivities)
-             + forcing(i, j, k, grid, clock, model_fields))
-end

src/Models/HydrostaticFreeSurfaceModels/store_hydrostatic_free_surface_tendencies.jl

@@ -1,6 +1,6 @@
 using KernelAbstractions: @index, @kernel
 
-using Oceananigans.TimeSteppers:  store_field_tendencies!
+using Oceananigans.TimeSteppers: store_field_tendencies!
 
 using Oceananigans: prognostic_fields
 using Oceananigans.Grids: AbstractGrid
@@ -19,7 +19,6 @@ end
 store_free_surface_tendency!(free_surface, model) = nothing
 
 function store_free_surface_tendency!(::ExplicitFreeSurface, model)
-
     launch!(model.architecture, model.grid, :xy,
             _store_free_surface_tendency!,
             model.timestepper.G⁻.η,
@@ -32,15 +31,22 @@ function store_tendencies!(model::HydrostaticFreeSurfaceModel)
     prognostic_field_names = keys(prognostic_fields(model))
     three_dimensional_prognostic_field_names = filter(name -> name != :η, prognostic_field_names)
 
+    closure = model.closure
+
     for field_name in three_dimensional_prognostic_field_names
 
-        launch!(model.architecture, model.grid, :xyz,
-                store_field_tendencies!,
-                model.timestepper.G⁻[field_name],
-                model.timestepper.Gⁿ[field_name])
+        if closure isa FlavorOfCATKE && field_name == :e
+            @debug "Skipping store tendencies for e"
+        else
+            launch!(model.architecture, model.grid, :xyz,
+                    store_field_tendencies!,
+                    model.timestepper.G⁻[field_name],
+                    model.timestepper.Gⁿ[field_name])
+        end
     end
 
     store_free_surface_tendency!(model.free_surface, model)
 
     return nothing
 end
+

src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl

@@ -25,7 +25,7 @@ hydrostatic pressure) to the current model state. If `callbacks` are provided (i
 they are called in the end.
 """
 update_state!(model::HydrostaticFreeSurfaceModel, callbacks=[]; compute_tendencies = true) =
-         update_state!(model, model.grid, callbacks; compute_tendencies)
+    update_state!(model, model.grid, callbacks; compute_tendencies)
 
 function update_state!(model::HydrostaticFreeSurfaceModel, grid, callbacks; compute_tendencies = true)
     @apply_regionally mask_immersed_model_fields!(model, grid)

src/Models/interleave_communication_and_computation.jl

@@ -29,7 +29,7 @@ compute_boundary_tendencies!(model) = nothing
 interior_tendency_kernel_parameters(grid) = :xyz # fallback
 
 interior_tendency_kernel_parameters(grid::DistributedGrid) = 
-            interior_tendency_kernel_parameters(grid, architecture(grid))
+    interior_tendency_kernel_parameters(grid, architecture(grid))
 
 interior_tendency_kernel_parameters(grid, ::SynchronizedDistributed) = :xyz
 

src/Oceananigans.jl

@@ -203,6 +203,7 @@ include("Operators/Operators.jl")
 include("BoundaryConditions/BoundaryConditions.jl")
 include("Fields/Fields.jl")
 include("AbstractOperations/AbstractOperations.jl")
+include("TimeSteppers/TimeSteppers.jl")
 include("Advection/Advection.jl")
 include("Solvers/Solvers.jl")
 include("OutputReaders/OutputReaders.jl")
@@ -225,7 +226,6 @@ include("Biogeochemistry.jl")
 include("ImmersedBoundaries/ImmersedBoundaries.jl")
 # include("DistributedComputations/DistributedComputations.jl")
 
-include("TimeSteppers/TimeSteppers.jl")
 include("Models/Models.jl")
 
 # Output and Physics, time-stepping, and models

src/OutputReaders/set_field_time_series.jl

@@ -1,3 +1,5 @@
+using Printf
+
 #####
 ##### set!
 #####
@@ -19,6 +21,14 @@ function set!(fts::InMemoryFTS, path::String=fts.path, name::String=fts.name)
     for n in time_indices(fts)
         t = fts.times[n]
         file_index = find_time_index(t, file_times)
+
+        if isnothing(file_index)
+            msg = string("Error setting ", summary(fts), '\n')
+            msg *= @sprintf("Can't find data for time %.1e and time index %d\n", t, n)
+            msg *= @sprintf("for field %s at path %s", path, name)
+            error(msg)
+        end
+
         file_iter = file_iterations[file_index]
 
         # Note: use the CPU for this step

src/TimeSteppers/TimeSteppers.jl

@@ -52,6 +52,7 @@ abstract type AbstractLagrangianParticles end
 step_lagrangian_particles!(model, Δt) = nothing
 
 reset!(timestepper) = nothing
+implicit_step!(field, ::Nothing, args...; kwargs...) = nothing
 
 include("clock.jl")
 include("store_tendencies.jl")