Include background tracer and velocities in closure tendency terms

src/Fields/Fields.jl

@@ -30,7 +30,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, ZStar, ZCoordinate,
     ExplicitFreeSurface, ImplicitFreeSurface, SplitExplicitFreeSurface,
@@ -101,7 +101,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,121 @@
+using Oceananigans.Fields: ZeroField, ConstantField, 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)
+
+# Fallback, assuming a function
+regularize_background_field(LX, LY, LZ, f::BackgroundField, grid, clock) =
+    FunctionField{LX, LY, LZ}(f.func, grid; clock, parameters=f.parameters)
+
+regularize_background_field(LX, LY, LZ, f::BackgroundField{<:Number}, grid, clock) = ConstantField(f.func)
+
+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,