Include background tracer and velocities in closure tendency terms

src/Fields/Fields.jl

@@ -31,7 +31,6 @@ include("field_indices.jl")
 include("scans.jl")
 include("regridding_fields.jl")
 include("field_tuples.jl")
-include("background_fields.jl")
 include("interpolate.jl")
 include("show_fields.jl")
 include("broadcasting_abstract_fields.jl")

src/Models/Models.jl

@@ -1,7 +1,7 @@
 module Models
 
 export
-    NonhydrostaticModel,
+    NonhydrostaticModel, BackgroundField, BackgroundFields,
     ShallowWaterModel, ConservativeFormulation, VectorInvariantFormulation,
     HydrostaticFreeSurfaceModel,
     ExplicitFreeSurface, ImplicitFreeSurface, SplitExplicitFreeSurface,
@@ -100,7 +100,7 @@ include("HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl")
 include("ShallowWaterModels/ShallowWaterModels.jl")
 include("LagrangianParticleTracking/LagrangianParticleTracking.jl")
 
-using .NonhydrostaticModels: NonhydrostaticModel, PressureField
+using .NonhydrostaticModels: NonhydrostaticModel, PressureField, BackgroundField, BackgroundFields
 
 using .HydrostaticFreeSurfaceModels:
     HydrostaticFreeSurfaceModel,

src/Models/NonhydrostaticModels/NonhydrostaticModels.jl

@@ -1,6 +1,6 @@
 module NonhydrostaticModels
 
-export NonhydrostaticModel
+export NonhydrostaticModel, BackgroundField, BackgroundFields
 
 using DocStringExtensions
 
@@ -65,6 +65,7 @@ nonhydrostatic_pressure_solver(grid) = nonhydrostatic_pressure_solver(architectu
 ##### NonhydrostaticModel definition
 #####
 
+include("background_fields.jl")
 include("nonhydrostatic_model.jl")
 include("pressure_field.jl")
 include("show_nonhydrostatic_model.jl")

src/Models/NonhydrostaticModels/background_fields.jl

@@ -0,0 +1,115 @@
+using Oceananigans.Fields: ZeroField, AbstractField, FunctionField, location
+using Oceananigans.Utils: prettysummary
+
+using Adapt
+
+function background_velocity_fields(fields, grid, clock)
+    u = get(fields, :u, ZeroField())
+    v = get(fields, :v, ZeroField())
+    w = get(fields, :w, ZeroField())
+
+    u = regularize_background_field(Face, Center, Center, u, grid, clock)
+    v = regularize_background_field(Center, Face, Center, v, grid, clock)
+    w = regularize_background_field(Center, Center, Face, w, grid, clock)
+
+    return (; u, v, w)
+end
+
+function background_tracer_fields(bg, tracer_names, grid, clock)
+    tracer_fields =
+        Tuple(c ∈ keys(bg) ?
+              regularize_background_field(Center, Center, Center, getindex(bg, c), grid, clock) :
+              ZeroField()
+              for c in tracer_names)
+
+    return NamedTuple{tracer_names}(tracer_fields)
+end
+
+#####
+##### BackgroundFields (with option for including background closure fluxes)
+#####
+
+struct BackgroundFields{Q, U, C}
+    velocities :: U
+    tracers :: C
+    function BackgroundFields{Q}(velocities::U, tracers::C) where {Q, U, C}
+        return new{Q, U, C}(velocities, tracers)
+    end
+end
+
+Adapt.adapt_structure(to, bf::BackgroundFields{Q}) where Q =    
+    BackgroundFields{Q}(adapt(to, bf.velocities), adapt(to, bf.tracers))
+
+const BackgroundFieldsWithClosureFluxes = BackgroundFields{true}
+
+function BackgroundFields(; background_closure_fluxes=false, fields...)
+    u = get(fields, :u, ZeroField())
+    v = get(fields, :v, ZeroField())
+    w = get(fields, :w, ZeroField())
+    velocities = (; u, v, w)
+    tracers = NamedTuple(name => fields[name] for name in keys(fields) if !(name ∈ (:u, :v, :w)))
+    return BackgroundFields{background_closure_fluxes}(velocities, tracers)
+end
+
+function BackgroundFields(background_fields::BackgroundFields{Q}, tracer_names, grid, clock) where Q
+    velocities = background_velocity_fields(background_fields.velocities, grid, clock)
+    tracers = background_tracer_fields(background_fields.tracers, tracer_names, grid, clock)
+    return BackgroundFields{Q}(velocities, tracers)
+end
+
+function BackgroundFields(background_fields::NamedTuple, tracer_names, grid, clock)
+    velocities = background_velocity_fields(background_fields, grid, clock)
+    tracers = background_tracer_fields(background_fields, tracer_names, grid, clock)
+    return BackgroundFields{false}(velocities, tracers)
+end
+
+"""
+    BackgroundField{F, P}
+
+Temporary container for storing information about `BackgroundFields`.
+"""
+struct BackgroundField{F, P}
+    func:: F
+    parameters :: P
+end
+
+"""
+    BackgroundField(func; parameters=nothing)
+
+Returns a `BackgroundField` to be passed to `NonhydrostaticModel` for use
+as a background velocity or tracer field.
+
+If `parameters` is not provided, `func` must be callable with the signature
+
+```julia
+func(x, y, z, t)
+```
+
+If `parameters` is provided, `func` must be callable with the signature
+
+```julia
+func(x, y, z, t, parameters)
+```
+"""
+BackgroundField(func; parameters=nothing) = BackgroundField(func, parameters)
+
+regularize_background_field(LX, LY, LZ, f::BackgroundField{<:Function}, grid, clock) =
+    FunctionField{LX, LY, LZ}(f.func, grid; clock, parameters=f.parameters)
+
+regularize_background_field(LX, LY, LZ, func::Function, grid, clock) =
+    FunctionField{LX, LY, LZ}(func, grid; clock=clock)
+
+regularize_background_field(LX, LY, LZ, ::ZeroField, grid, clock) = ZeroField()
+
+function regularize_background_field(LX, LY, LZ, field::AbstractField, grid, clock)
+    if location(field) != (LX, LY, LZ)
+        throw(ArgumentError("Cannot use field at $(location(field)) as a background field at $((LX, LY, LZ))"))
+    end
+
+    return field
+end
+
+Base.show(io::IO, field::BackgroundField{F, P}) where {F, P} =
+    print(io, "BackgroundField{$F, $P}", "\n",
+          "├── func: $(prettysummary(field.func))", "\n",
+          "└── parameters: $(field.parameters)")

src/Models/NonhydrostaticModels/nonhydrostatic_model.jl

@@ -7,7 +7,7 @@ using Oceananigans.Advection: Centered, adapt_advection_order
 using Oceananigans.BuoyancyFormulations: validate_buoyancy, regularize_buoyancy, SeawaterBuoyancy
 using Oceananigans.Biogeochemistry: validate_biogeochemistry, AbstractBiogeochemistry, biogeochemical_auxiliary_fields
 using Oceananigans.BoundaryConditions: regularize_field_boundary_conditions
-using Oceananigans.Fields: BackgroundFields, Field, tracernames, VelocityFields, TracerFields, CenterField
+using Oceananigans.Fields: Field, tracernames, VelocityFields, TracerFields, CenterField
 using Oceananigans.Forcings: model_forcing
 using Oceananigans.Grids: inflate_halo_size, with_halo, architecture
 using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid
@@ -24,6 +24,7 @@ import Oceananigans.Models: total_velocities, default_nan_checker, timestepper
 
 const ParticlesOrNothing = Union{Nothing, AbstractLagrangianParticles}
 const AbstractBGCOrNothing = Union{Nothing, AbstractBiogeochemistry}
+const BFOrNamedTuple = Union{BackgroundFields, NamedTuple}
 
 # TODO: this concept may be more generally useful,
 # but for now we use it only for hydrostatic pressure anomalies for now.
@@ -122,7 +123,7 @@ function NonhydrostaticModel(; grid,
                              boundary_conditions::NamedTuple = NamedTuple(),
                              tracers = (),
                              timestepper = :RungeKutta3,
-                             background_fields::NamedTuple = NamedTuple(),
+                             background_fields::BFOrNamedTuple = NamedTuple(),
                              particles::ParticlesOrNothing = nothing,
                              biogeochemistry::AbstractBGCOrNothing = nothing,
                              velocities = nothing,

src/Models/NonhydrostaticModels/nonhydrostatic_tendency_kernel_functions.jl

@@ -8,6 +8,7 @@ using Oceananigans.Biogeochemistry: biogeochemical_transition, biogeochemical_dr
 using Oceananigans.TurbulenceClosures: ∂ⱼ_τ₁ⱼ, ∂ⱼ_τ₂ⱼ, ∂ⱼ_τ₃ⱼ, ∇_dot_qᶜ
 using Oceananigans.TurbulenceClosures: immersed_∂ⱼ_τ₁ⱼ, immersed_∂ⱼ_τ₂ⱼ, immersed_∂ⱼ_τ₃ⱼ, immersed_∇_dot_qᶜ
 using Oceananigans.Forcings: with_advective_forcing
+using Oceananigans.Fields: ZeroField
 
 "return the ``x``-gradient of hydrostatic pressure"
 hydrostatic_pressure_gradient_x(i, j, k, grid, hydrostatic_pressure) = ∂xᶠᶜᶜ(i, j, k, grid, hydrostatic_pressure)
@@ -17,6 +18,28 @@ hydrostatic_pressure_gradient_x(i, j, k, grid, ::Nothing) = zero(grid)
 hydrostatic_pressure_gradient_y(i, j, k, grid, hydrostatic_pressure) = ∂yᶜᶠᶜ(i, j, k, grid, hydrostatic_pressure)
 hydrostatic_pressure_gradient_y(i, j, k, grid, ::Nothing) = zero(grid)
 
+# Compiler shortcut, for the paranoid about ZeroField
+@inline sum_fields(a, b::ZeroField) = a
+@inline sum_fields(a, b)            = SumOfArrays{2}(a, b)
+
+@inline sum_fields(a, b::ZeroField, c::ZeroField) = a
+@inline sum_fields(a, b, c::ZeroField)            = SumOfArrays{2}(a, b)
+@inline sum_fields(a, b::ZeroField, c)            = SumOfArrays{2}(a, c)
+@inline sum_fields(a, b, c)                       = SumOfArrays{3}(a, b, c)
+
+@inline assemble_closure_velocities(velocities, background_fields) = velocities
+
+@inline function assemble_closure_velocities(velocities,
+                                             background_fields::BackgroundFieldsWithClosureFluxes)
+
+    u = sum_fields(velocities.u, background_fields.velocities.u)
+    v = sum_fields(velocities.v, background_fields.velocities.v)
+    w = sum_fields(velocities.w, background_fields.velocities.w)
+
+    return (; u, v, w)
+end
+
+
 """
     $(SIGNATURES)
 
@@ -60,17 +83,19 @@ pressure anomaly.
                                      clock,
                                      forcing)
 
-    model_fields = merge(velocities, tracers, auxiliary_fields)
-
     total_velocities = sum_of_velocities(velocities, background_fields.velocities)
     total_velocities = with_advective_forcing(forcing, total_velocities)
 
+    closure_velocities = assemble_closure_velocities(velocities, background_fields)    
+    closure_model_fields = merge(closure_velocities, tracers, auxiliary_fields)
+    model_fields = merge(velocities, tracers, auxiliary_fields)
+
     return ( - div_𝐯u(i, j, k, grid, advection, total_velocities, velocities.u)
              - div_𝐯u(i, j, k, grid, advection, velocities, background_fields.velocities.u)
              + x_dot_g_bᶠᶜᶜ(i, j, k, grid, buoyancy, tracers)
              - x_f_cross_U(i, j, k, grid, coriolis, velocities)
              - hydrostatic_pressure_gradient_x(i, j, k, grid, hydrostatic_pressure)
-             - ∂ⱼ_τ₁ⱼ(i, j, k, grid, closure, diffusivities, clock, model_fields, buoyancy)
+             - ∂ⱼ_τ₁ⱼ(i, j, k, grid, closure, diffusivities, clock, closure_model_fields, buoyancy)
              - immersed_∂ⱼ_τ₁ⱼ(i, j, k, grid, velocities, u_immersed_bc, closure, diffusivities, clock, model_fields)
              + x_curl_Uˢ_cross_U(i, j, k, grid, stokes_drift, velocities, clock.time)
              + ∂t_uˢ(i, j, k, grid, stokes_drift, clock.time)
@@ -120,17 +145,19 @@ pressure anomaly.
                                      clock,
                                      forcing)
 
-    model_fields = merge(velocities, tracers, auxiliary_fields)
-
     total_velocities = sum_of_velocities(velocities, background_fields.velocities)
     total_velocities = with_advective_forcing(forcing, total_velocities)
 
+    closure_velocities = assemble_closure_velocities(velocities, background_fields)    
+    closure_model_fields = merge(closure_velocities, tracers, auxiliary_fields)
+    model_fields = merge(velocities, tracers, auxiliary_fields)
+
     return ( - div_𝐯v(i, j, k, grid, advection, total_velocities, velocities.v)
              - div_𝐯v(i, j, k, grid, advection, velocities, background_fields.velocities.v)
              + y_dot_g_bᶜᶠᶜ(i, j, k, grid, buoyancy, tracers)
              - y_f_cross_U(i, j, k, grid, coriolis, velocities)
              - hydrostatic_pressure_gradient_y(i, j, k, grid, hydrostatic_pressure)
-             - ∂ⱼ_τ₂ⱼ(i, j, k, grid, closure, diffusivities, clock, model_fields, buoyancy)
+             - ∂ⱼ_τ₂ⱼ(i, j, k, grid, closure, diffusivities, clock, closure_model_fields, buoyancy)
              - immersed_∂ⱼ_τ₂ⱼ(i, j, k, grid, velocities, v_immersed_bc, closure, diffusivities, clock, model_fields)
              + y_curl_Uˢ_cross_U(i, j, k, grid, stokes_drift, velocities, clock.time)
              + ∂t_vˢ(i, j, k, grid, stokes_drift, clock.time)
@@ -183,16 +210,18 @@ velocity components, tracer fields, and precalculated diffusivities where applic
                                      clock,
                                      forcing)
 
-    model_fields = merge(velocities, tracers, auxiliary_fields)
-
     total_velocities = sum_of_velocities(velocities, background_fields.velocities)
     total_velocities = with_advective_forcing(forcing, total_velocities)
 
+    closure_velocities = assemble_closure_velocities(velocities, background_fields)    
+    closure_model_fields = merge(closure_velocities, tracers, auxiliary_fields)
+    model_fields = merge(velocities, tracers, auxiliary_fields)
+
     return ( - div_𝐯w(i, j, k, grid, advection, total_velocities, velocities.w)
              - div_𝐯w(i, j, k, grid, advection, velocities, background_fields.velocities.w)
              + maybe_z_dot_g_bᶜᶜᶠ(i, j, k, grid, hydrostatic_pressure, buoyancy, tracers)
              - z_f_cross_U(i, j, k, grid, coriolis, velocities)
-             - ∂ⱼ_τ₃ⱼ(i, j, k, grid, closure, diffusivities, clock, model_fields, buoyancy)
+             - ∂ⱼ_τ₃ⱼ(i, j, k, grid, closure, diffusivities, clock, closure_model_fields, buoyancy)
              - immersed_∂ⱼ_τ₃ⱼ(i, j, k, grid, velocities, w_immersed_bc, closure, diffusivities, clock, model_fields)
              + z_curl_Uˢ_cross_U(i, j, k, grid, stokes_drift, velocities, clock.time)
              + ∂t_wˢ(i, j, k, grid, stokes_drift, clock.time)
@@ -226,7 +255,7 @@ velocity components, tracer fields, and precalculated diffusivities where applic
 `clock` keeps track of `clock.time` and `clock.iteration`.
 """
 @inline function tracer_tendency(i, j, k, grid,
-                                 val_tracer_index::Val{tracer_index},
+                                 val_index::Val{tracer_index},
                                  val_tracer_name,
                                  advection,
                                  closure,
@@ -241,19 +270,30 @@ velocity components, tracer fields, and precalculated diffusivities where applic
                                  clock,
                                  forcing) where tracer_index
 
-    @inbounds c = tracers[tracer_index]
-    @inbounds background_fields_c = background_fields.tracers[tracer_index]
-    model_fields = merge(velocities, tracers, auxiliary_fields)
 
     biogeochemical_velocities = biogeochemical_drift_velocity(biogeochemistry, val_tracer_name)
 
     total_velocities = sum_of_velocities(velocities, background_fields.velocities, biogeochemical_velocities)
     total_velocities = with_advective_forcing(forcing, total_velocities)
 
+    @inbounds c = tracers[tracer_index]
+    @inbounds background_fields_c = background_fields.tracers[tracer_index]
+
+    closure_c = if background_fields isa BackgroundFieldsWithClosureFluxes
+        sum_fields(c, background_fields_c)
+    else
+        c
+    end
+
+    closure_velocities = assemble_closure_velocities(velocities, background_fields)
+    closure_model_fields = merge(closure_velocities, tracers, auxiliary_fields)
+    model_fields = merge(velocities, tracers, auxiliary_fields)
+
     return ( - div_Uc(i, j, k, grid, advection, total_velocities, c)
              - div_Uc(i, j, k, grid, advection, velocities, background_fields_c)
-             - ∇_dot_qᶜ(i, j, k, grid, closure, diffusivities, val_tracer_index, c, clock, model_fields, buoyancy)
-             - immersed_∇_dot_qᶜ(i, j, k, grid, c, c_immersed_bc, closure, diffusivities, val_tracer_index, clock, model_fields)
+             - ∇_dot_qᶜ(i, j, k, grid, closure, diffusivities, val_index, closure_c, clock, closure_model_fields, buoyancy)
+             - immersed_∇_dot_qᶜ(i, j, k, grid, closure_c, c_immersed_bc, closure, diffusivities, val_index, clock, model_fields)
              + biogeochemical_transition(i, j, k, grid, biogeochemistry, val_tracer_name, clock, model_fields)
              + forcing(i, j, k, grid, clock, model_fields))
 end
+

test/test_time_stepping.jl

@@ -2,7 +2,6 @@ include("dependencies_for_runtests.jl")
 
 using TimesDates: TimeDate
 using Oceananigans.Grids: topological_tuple_length, total_size
-using Oceananigans.Fields: BackgroundField
 using Oceananigans.TimeSteppers: Clock
 using Oceananigans.TurbulenceClosures: CATKEVerticalDiffusivity
 using Oceananigans.TurbulenceClosures.Smagorinskys: LagrangianAveraging, DynamicSmagorinsky