Change the user interface so that closure tracers are automatically added

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl

@@ -13,7 +13,7 @@ using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid
 using Oceananigans.Models: AbstractModel, validate_model_halo, validate_tracer_advection, extract_boundary_conditions, initialization_update_state!
 using Oceananigans.TimeSteppers: Clock, TimeStepper, update_state!, AbstractLagrangianParticles, SplitRungeKutta3TimeStepper
 using Oceananigans.TurbulenceClosures: validate_closure, with_tracers, build_diffusivity_fields, add_closure_specific_boundary_conditions
-using Oceananigans.TurbulenceClosures: time_discretization, implicit_diffusion_solver
+using Oceananigans.TurbulenceClosures: time_discretization, implicit_diffusion_solver, closure_required_tracers
 using Oceananigans.Utils: tupleit
 
 import Oceananigans: initialize!
@@ -141,14 +141,35 @@ function HydrostaticFreeSurfaceModel(; grid,
     @apply_regionally validate_model_halo(grid, momentum_advection, tracer_advection, closure)
 
     if !(grid isa MutableGridOfSomeKind) && (vertical_coordinate isa ZStarCoordinate)
-        error("The grid does not support ZStarCoordinate vertical coordinates. Use a `MutableVerticalDiscretization` to allow the use of ZStarCoordinate (see `MutableVerticalDiscretization`).")
+        msg = string("The grid ", summary(grid), " does not support ZStarCoordinate.", '\n',
+                     "z must be a MutableVerticalDiscretization to allow the use of ZStarCoordinate.")
+        throw(ArgumentError(msg))
     end
 
     # Validate biogeochemistry (add biogeochemical tracers automagically)
     tracers = tupleit(tracers) # supports tracers=:c keyword argument (for example)
     biogeochemical_fields = biogeochemical_auxiliary_fields(biogeochemistry)
     tracers, biogeochemical_fields = validate_biogeochemistry(tracers, biogeochemical_fields, biogeochemistry, grid, clock)
 
+    # Automatically append closure-required tracers and disallow users from specifying them explicitly
+    user_tracer_names = tracernames(tracers)
+    closure_tracer_names = closure_required_tracers(closure)
+
+    # Throw an error in case of a conflict between user-specified tracers and any other tracers.
+    if any(name ∈ user_tracer_names for name in closure_tracer_names)
+        msg = string("The tracer names $(user_tracer_names) overlap with the closure auxiliary", '\n',
+                     "tracer names $(closure_tracer_names) associated with $(summary(closure)).", '\n',
+                     "The names $(closure_tracer_names) cannot be specified explicitly", '\n',
+                     "or be the names of biogeochemical tracers.")
+        throw(ArgumentError(msg))
+    elseif tracers isa NamedTuple
+        closure_tracer_fields = Tuple(CenterField(grid) for _ in closure_tracer_names)
+        closure_tracers = NamedTuple{closure_tracer_names}(closure_tracer_fields)
+        tracers = merge(tracers, closure_tracers)
+    else
+        tracers = tuple(user_tracer_names..., closure_tracer_names...)
+    end
+
     # Reduce the advection order in directions that do not have enough grid points
     @apply_regionally momentum_advection = validate_momentum_advection(momentum_advection, grid)
     default_tracer_advection, tracer_advection = validate_tracer_advection(tracer_advection, grid)
@@ -259,4 +280,3 @@ initialize!(model::HydrostaticFreeSurfaceModel) = initialize_free_surface!(model
 
 # return the total advective velocities
 @inline total_velocities(model::HydrostaticFreeSurfaceModel) = model.velocities
-

src/TurbulenceClosures/TurbulenceClosures.jl

@@ -36,7 +36,8 @@ export
     ∂ⱼ_τ₂ⱼ,
     ∂ⱼ_τ₃ⱼ,
 
-    cell_diffusion_timescale
+    cell_diffusion_timescale,
+    closure_required_tracers
 
 using KernelAbstractions
 using Adapt
@@ -79,6 +80,23 @@ closure_summary(closure) = summary(closure)
 with_tracers(tracers, closure::AbstractTurbulenceClosure) = closure
 compute_diffusivities!(K, closure::AbstractTurbulenceClosure, args...; kwargs...) = nothing
 
+# Tracer names that a closure requires (eg TKE-based closures)
+# Fallbacks: by default closures do not require extra tracers.
+closure_required_tracers(::Nothing) = ()
+closure_required_tracers(::AbstractTurbulenceClosure) = ()
+closure_required_tracers(::AbstractArray{<:AbstractTurbulenceClosure}) = ()
+closure_required_tracers(closures::Tuple) = begin
+    names = Symbol[]
+    for c in closures
+        for n in closure_required_tracers(c)
+            if !(n in names)
+                push!(names, n)
+            end
+        end
+    end
+    return Tuple(names)
+end
+
 # The required halo size to calculate diffusivities. Take care that if the diffusivity can
 # be calculated from local information, still `B = 1`, because we need at least one additional
 # point at each side to calculate viscous fluxes at the edge of the domain.
@@ -210,4 +228,3 @@ include("turbulence_closure_diagnostics.jl")
 #####
 
 end # module
-

src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_vertical_diffusivity.jl

@@ -139,6 +139,9 @@ function with_tracers(tracer_names, closure::FlavorOfCATKE)
     return closure
 end
 
+# Required tracer names for CATKE
+closure_required_tracers(::FlavorOfCATKE) = tuple(:e)
+
 # For tuples of closures, we need to know _which_ closure is CATKE.
 # Here we take a "simple" approach that sorts the tuple so CATKE is first.
 # This is not sustainable though if multiple closures require this.

src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/tke_dissipation_vertical_diffusivity.jl

@@ -151,6 +151,9 @@ function with_tracers(tracer_names, closure::FlavorOfTD)
     return closure
 end
 
+# Required tracer names for TKEDissipation
+closure_required_tracers(::FlavorOfTD) = (:e, :ϵ)
+
 #####
 ##### Stratified displacement length scale limiter
 #####

test/benchmark_tests.jl

@@ -37,7 +37,7 @@ function ocean_benchmark(arch, Nx, Ny, Nz, topology, immersed, tracer_advection=
                                           buoyancy,
                                           closure,
                                           free_surface,
-                                          tracers = (:T, :S, :e))
+                                          tracers = (:T, :S))
 
     @info "Model is built"
 

test/test_hydrostatic_free_surface_models.jl

@@ -17,7 +17,6 @@ function time_step_hydrostatic_model_works(grid;
                                            velocities = nothing)
 
     buoyancy = BuoyancyTracer()
-    closure isa CATKEVerticalDiffusivity && push!(tracers, :e)
 
     model = HydrostaticFreeSurfaceModel(; grid, coriolis, tracers, velocities, buoyancy,
                                         momentum_advection, tracer_advection, free_surface, closure)
@@ -68,7 +67,7 @@ function time_step_hydrostatic_model_with_catke_works(arch, FT)
     model = HydrostaticFreeSurfaceModel(;
         grid,
         buoyancy = BuoyancyTracer(),
-        tracers = (:b, :e),
+        tracers = (:b,),
         closure = CATKEVerticalDiffusivity(FT)
     )
 

test/test_turbulence_closures.jl

@@ -175,7 +175,7 @@ function run_catke_tke_substepping_tests(arch, closure)
     grid = RectilinearGrid(arch, size=(2, 2, 2), extent=(100, 200, 300))
 
     model = HydrostaticFreeSurfaceModel(; grid, momentum_advection = nothing, tracer_advection = nothing,
-                                          closure, buoyancy=BuoyancyTracer(), tracers=(:b, :e))
+                                          closure, buoyancy=BuoyancyTracer(), tracers=(:b))
 
     # set random velocities
     Random.seed!(1234)
@@ -209,14 +209,17 @@ function run_time_step_with_catke_tests(arch, closure)
     grid = RectilinearGrid(arch, size=(2, 2, 2), extent=(1, 2, 3))
     buoyancy = BuoyancyTracer()
 
-    # These shouldn't work (need :e in tracers)
-    @test_throws ArgumentError HydrostaticFreeSurfaceModel(; grid, closure, buoyancy, tracers=:b)
-    @test_throws ArgumentError HydrostaticFreeSurfaceModel(; grid, closure, buoyancy, tracers=(:b, :E))
+    # These should work now (:e is auto-appended)
+    @test HydrostaticFreeSurfaceModel(; grid, closure, buoyancy, tracers=:b) isa HydrostaticFreeSurfaceModel
+    @test HydrostaticFreeSurfaceModel(; grid, closure, buoyancy, tracers=(:b, :E)) isa HydrostaticFreeSurfaceModel
 
     # CATKE isn't supported with NonhydrostaticModel (we don't diffuse vertical velocity)
     @test_throws ErrorException NonhydrostaticModel(; grid, closure, buoyancy, tracers=(:b, :c, :e))
 
-    model = HydrostaticFreeSurfaceModel(; grid, closure, buoyancy, tracers = (:b, :c, :e))
+    # Supplying closure tracers explicitly should error
+    @test_throws ArgumentError HydrostaticFreeSurfaceModel(; grid, closure, buoyancy, tracers = (:b, :c, :e))
+
+    model = HydrostaticFreeSurfaceModel(; grid, closure, buoyancy, tracers = (:b, :c))
 
     # Default boundary condition is Flux, Nothing... with CATKE this has to change.
     @test !(model.tracers.e.boundary_conditions.top.condition isa BoundaryCondition{Flux, Nothing})

validation/distributed_simulations/distributed_scaling/distributed_hydrostatic_simulation.jl

@@ -53,7 +53,7 @@ function run_hydrostatic_simulation!(grid_size, ranks, FT::DataType = Float64;
                                           coriolis,
                                           closure,
                                           free_surface,
-                                          tracers = (:T, :S, :e),
+                                          tracers = (:T, :S),
                                           buoyancy,
                                           timestepper)
 

validation/mesoscale_turbulence/eddying_channel.jl

@@ -180,7 +180,7 @@ model = HydrostaticFreeSurfaceModel(grid = grid,
                                     buoyancy = BuoyancyTracer(),
                                     coriolis = coriolis,
                                     closure = (horizontal_closure, vertical_closure, catke),
-                                    tracers = (:b, :e, :c),
+                                    tracers = (:b, :c),
                                     boundary_conditions = (b=b_bcs, u=u_bcs, v=v_bcs),
                                     forcing = (; b=Fb,)
                                     )

validation/mesoscale_turbulence/zonally_averaged_channel.jl

@@ -179,7 +179,7 @@ model = HydrostaticFreeSurfaceModel(grid = grid,
                                     coriolis = coriolis,
                                     #closure = (horizontal_diffusivity, convective_adjustment, gent_mcwilliams_diffusivity),
                                     closure = (catke, diffusive_closure..., gent_mcwilliams_diffusivity),
-                                    tracers = (:b, :e, :c),
+                                    tracers = (:b, :c),
                                     boundary_conditions = (b=b_bcs, u=u_bcs, v=v_bcs),
                                     forcing = (; b=Fb))
 

validation/vertical_mixing_closures/bottom_boundary_layer.jl

@@ -37,7 +37,7 @@ u_bcs = FieldBoundaryConditions(bottom = u_bottom_bc)
 v_bcs = FieldBoundaryConditions(bottom = v_bottom_bc)
 
 model = HydrostaticFreeSurfaceModel(; grid, closure, coriolis,
-                                    tracers = (:b, :e),
+                                    tracers = (:b,),
                                     buoyancy = BuoyancyTracer(),
                                     boundary_conditions = (; u=u_bcs, v=v_bcs))
 
@@ -113,4 +113,3 @@ record(fig, "bottom_boundary_layer.mp4", 1:Nt, framerate=12) do nn
     n[] = nn
 end
 
-

validation/vertical_mixing_closures/column_windy_convection.jl

@@ -50,7 +50,7 @@ end
 for closure in closures_to_run
 
     model = HydrostaticFreeSurfaceModel(; grid, closure, coriolis,
-                                        tracers = (:b, :e, :ϵ),
+                                        tracers = (:b,),
                                         buoyancy = BuoyancyTracer(),
                                         boundary_conditions = (; b=b_bcs, u=u_bcs))
 
@@ -165,4 +165,3 @@ display(fig)
 # record(fig, "windy_convection.mp4", 1:Nt, framerate=24) do nn
 #     n[] = nn
 # end
-

validation/vertical_mixing_closures/compare_catke_k_epsilon.jl

@@ -48,7 +48,7 @@ for closure in (k_epsilon, catke) #, k_epsilon_const_stability)
     global model
 
     model = HydrostaticFreeSurfaceModel(; grid, closure, coriolis,
-                                        tracers = (:b, :e, :ϵ),
+                                        tracers = (:b,),
                                         buoyancy = BuoyancyTracer(),
                                         boundary_conditions=(u=u_bcs, b=b_bcs))
 
@@ -111,4 +111,3 @@ lines!(axϵ, ϵn[3], z, color=colors[3])
 Legend(fig[0, 1:4], axe, nbanks=3, framevisible=false, tellheight=true)
 
 fig
-

validation/vertical_mixing_closures/gpu_tkevd_ensemble.jl

@@ -31,7 +31,7 @@ coriolis_ensemble = CuArray([FPlane(f=f_ij(i, j)) for i=1:Ex, j=1:Ey])
 
 model = HydrostaticFreeSurfaceModel(architecture = GPU(),
                                     grid = grid,
-                                    tracers = (:b, :e),
+                                    tracers = (:b,),
                                     buoyancy = BuoyancyTracer(),
                                     coriolis = coriolis_ensemble,
                                     boundary_conditions = (b=b_bcs, u=u_bcs, v=v_bcs),

validation/vertical_mixing_closures/heterogeneous_windy_convection.jl

@@ -66,7 +66,7 @@ model = HydrostaticFreeSurfaceModel(; grid, closure,
                                     momentum_advection = WENO(),
                                     tracer_advection = WENO(),
                                     coriolis = FPlane(f=1e-4),
-                                    tracers = (:b, :e),
+                                    tracers = (:b,),
                                     boundary_conditions = (; b=b_bcs, u=u_bcs),
                                     buoyancy = BuoyancyTracer())
 
@@ -249,4 +249,3 @@ record(fig, filename[1:end-5] * ".mp4", 1:Nt, framerate=24) do nn
     n[] = nn
 end
 =#
-

validation/vertical_mixing_closures/k_epsilon_boundary_layer.jl

@@ -35,7 +35,7 @@ closure = TKEDissipationVerticalDiffusivity()
 #closure = CATKEVerticalDiffusivity()
 
 model = HydrostaticFreeSurfaceModel(; grid, closure, coriolis,
-                                    tracers = (:b, :e, :ϵ),
+                                    tracers = (:b,),
                                     buoyancy = BuoyancyTracer(),
                                     boundary_conditions=(u=u_bcs, b=b_bcs))
 
@@ -120,4 +120,3 @@ lines!(axs, 𝕊ᶜn, zf, label="𝕊ᶜ")
 axislegend(axs, position=:rb)
 
 fig
-

validation/vertical_mixing_closures/test_single_column_model.jl

@@ -10,7 +10,7 @@ boundary_conditions = (b = FieldBoundaryConditions(top = FluxBoundaryCondition(1
                        u = FieldBoundaryConditions(top = FluxBoundaryCondition(-2e-4)))
 
 model = HydrostaticFreeSurfaceModel(; grid, closure, coriolis, boundary_conditions,
-                                    tracers = (:b, :e), buoyancy = BuoyancyTracer())
+                                    tracers = (:b,), buoyancy = BuoyancyTracer())
 
 bᵢ(z) = 1e-6 * z
 set!(model, b=bᵢ, e=1e-6)
@@ -22,4 +22,3 @@ time_step!(model, 600)
     time_step!(model, 600)
 end
 
-

validation/vertical_mixing_closures/tupled_vertical_diffusion.jl

@@ -8,12 +8,11 @@ closure = (VerticalScalarDiffusivity(VerticallyImplicitTimeDiscretization(), κ=
            CATKEVerticalDiffusivity())
 
 model = HydrostaticFreeSurfaceModel(; grid, closure,
-                                    tracers = (:b, :e),
+                                    tracers = (:b,),
                                     buoyancy = BuoyancyTracer())
 
 bᵢ(z) = 1e-5 * z
 set!(model, b = bᵢ)
 simulation = Simulation(model, Δt=1minute, stop_iteration=10)
 
 run!(simulation)
-