Include background tracer and velocities in closure tendency terms

src/Fields/Fields.jl

@@ -3,7 +3,6 @@ module Fields
 export Face, Center
 export AbstractField, Field, Average, Integral, Reduction, Accumulation, field
 export CenterField, XFaceField, YFaceField, ZFaceField
-export BackgroundField
 export interior, data, xnode, ynode, znode, location
 export set!, compute!, @compute, regrid!
 export VelocityFields, TracerFields, TendencyFields, tracernames
@@ -24,7 +23,6 @@ include("field.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,
@@ -98,7 +98,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
 
@@ -40,6 +40,7 @@ PressureSolver(arch, grid) = error("None of the implemented pressure solvers for
 ##### 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,117 @@
+using Oceananigans.Fields: ZeroField, AbstractField, FunctionField, location
+using Oceananigans.Utils: prettysummary
+
+using Adapt
+
+# TODO: This code belongs in the Models module
+
+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=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: CenteredSecondOrder
 using Oceananigans.BuoyancyModels: 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}
 
 mutable struct NonhydrostaticModel{TS, E, A<:AbstractArchitecture, G, T, B, R, SD, U, C, Φ, F,
                                    V, S, K, BG, P, BGC, I, AF} <: AbstractModel{TS}
@@ -121,7 +122,7 @@ function NonhydrostaticModel(; grid,
           boundary_conditions::NamedTuple = NamedTuple(),
                                   tracers = (),
                               timestepper = :QuasiAdamsBashforth2,
-            background_fields::NamedTuple = NamedTuple(),
+        background_fields::BFOrNamedTuple = NamedTuple(),
             particles::ParticlesOrNothing = nothing,
     biogeochemistry::AbstractBGCOrNothing = nothing,
                                velocities = nothing,

src/Models/NonhydrostaticModels/nonhydrostatic_tendency_kernel_functions.jl

@@ -262,10 +262,16 @@ velocity components, tracer fields, and precalculated diffusivities where applic
 
     total_velocities = with_advective_forcing(forcing, total_velocities)
 
+    closure_c = if background_fields isa BackgroundFieldsWithClosureFluxes
+        SumOfArrays{2}(c, background_fields_c)
+    else
+        c
+    end
+
     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_tracer_index, closure_c, clock, model_fields, buoyancy)
+             - immersed_∇_dot_qᶜ(i, j, k, grid, closure_c, c_immersed_bc, closure, diffusivities, val_tracer_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