(0.95.7) Correctly account for field names in ContinuousBoundaryFunctions

Project.toml

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

src/BoundaryConditions/continuous_boundary_function.jl

@@ -72,7 +72,7 @@ The regularization of `bc.condition::ContinuousBoundaryFunction` requries
    of the boundary.
 """
 function regularize_boundary_condition(bc::BoundaryCondition{C, <:ContinuousBoundaryFunction},
-                                       grid, loc, dim, Side, prognostic_field_names) where C
+                                       grid, loc, dim, Side, field_names) where C
 
     boundary_func = bc.condition
 
@@ -81,7 +81,7 @@ function regularize_boundary_condition(bc::BoundaryCondition{C, <:ContinuousBoun
 
     indices, interps = index_and_interp_dependencies(LX, LY, LZ,
                                                      boundary_func.field_dependencies,
-                                                     prognostic_field_names)
+                                                     field_names)
 
     regularized_boundary_func = ContinuousBoundaryFunction{LX, LY, LZ, Side}(boundary_func.func,
                                                                              boundary_func.parameters,

src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl

@@ -75,9 +75,19 @@ Return a flattened `NamedTuple` of the fields in `model.velocities`, `model.free
 `model.tracers`, and any auxiliary fields for a `HydrostaticFreeSurfaceModel` model.
 """
 @inline fields(model::HydrostaticFreeSurfaceModel) = 
-    merge(hydrostatic_fields(model.velocities, model.free_surface, model.tracers),
-          model.auxiliary_fields,
-          biogeochemical_auxiliary_fields(model.biogeochemistry))
+    merge(hydrostatic_fields(model.velocities, model.free_surface, model.tracers), model.auxiliary_fields)
+
+"""
+    field_names(model)
+
+Return a tuple of the field names associated with `HydrostaticFreeSurfaceModel`. 
+Equivalent to doing `keys(fields(model))`.
+"""
+@inline field_names(model::HydrostaticFreeSurfaceModel) = field_names(model.free_surface, tracernames(model.tracers), model.auxiliary_fields)
+
+@inline field_names(free_surface, tracernames, auxiliary_fields) = (:u, :v, :w, tracernames..., :η, keys(auxiliary_fields)...)
+@inline field_names(::Nothing, tracernames, auxiliary_fields)    = (:u, :v, :w, tracernames..., keys(auxiliary_fields)...)
+@inline field_names(::SplitExplicitFreeSurface, tracernames, auxiliary_fields) = (:u, :v, :w, tracernames..., :η, :U, :V, keys(auxiliary_fields)...)
 
 """
     prognostic_fields(model::HydrostaticFreeSurfaceModel)
@@ -111,11 +121,11 @@ Return a flattened `NamedTuple` of the prognostic fields associated with `Hydros
 
 @inline hydrostatic_fields(velocities, free_surface::SplitExplicitFreeSurface, tracers) = merge((u = velocities.u,
                                                                                                  v = velocities.v,
-                                                                                                 w = velocities.w,
-                                                                                                 η = free_surface.η,
-                                                                                                 U = free_surface.barotropic_velocities.U,
-                                                                                                 V = free_surface.barotropic_velocities.V),
-                                                                                                 tracers)
+                                                                                                 w = velocities.w),
+                                                                                                 tracers,
+                                                                                                 (; η = free_surface.η,
+                                                                                                    U = free_surface.barotropic_velocities.U,
+                                                                                                    V = free_surface.barotropic_velocities.V))
 
 @inline hydrostatic_fields(velocities, ::Nothing, tracers) = merge((u = velocities.u,
                                                                     v = velocities.v,

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl

@@ -158,14 +158,14 @@ function HydrostaticFreeSurfaceModel(; grid,
                                          extract_boundary_conditions(diffusivity_fields))
 
     # Next, we form a list of default boundary conditions:
-    prognostic_field_names = (:u, :v, :w, tracernames(tracers)..., :η, keys(auxiliary_fields)...)
-    default_boundary_conditions = NamedTuple{prognostic_field_names}(Tuple(FieldBoundaryConditions()
-                                                                           for name in prognostic_field_names))
+    hydrostatic_field_names = field_names(free_surface, tracernames(tracers), auxiliary_fields)
+    default_boundary_conditions = NamedTuple{hydrostatic_field_names}(Tuple(FieldBoundaryConditions()
+                                                                      for name in hydrostatic_field_names))
 
     # Then we merge specified, embedded, and default boundary conditions. Specified boundary conditions
     # have precedence, followed by embedded, followed by default.
     boundary_conditions = merge(default_boundary_conditions, embedded_boundary_conditions, boundary_conditions)
-    boundary_conditions = regularize_field_boundary_conditions(boundary_conditions, grid, prognostic_field_names)
+    boundary_conditions = regularize_field_boundary_conditions(boundary_conditions, grid, hydrostatic_field_names)
 
     # Finally, we ensure that closure-specific boundary conditions, such as
     # those required by CATKEVerticalDiffusivity, are enforced:

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl

@@ -49,7 +49,7 @@ implicitly during time-stepping.
              - ∂xᶠᶜᶜ(i, j, k, grid, hydrostatic_pressure_anomaly)
              - ∂ⱼ_τ₁ⱼ(i, j, k, grid, closure, diffusivities, clock, model_fields, buoyancy)
              - immersed_∂ⱼ_τ₁ⱼ(i, j, k, grid, velocities, u_immersed_bc, closure, diffusivities, clock, model_fields)
-             + forcing(i, j, k, grid, clock, hydrostatic_prognostic_fields(velocities, free_surface, tracers)))
+             + forcing(i, j, k, grid, clock, model_fields))
 end
 
 """

src/Models/HydrostaticFreeSurfaceModels/nothing_free_surface.jl

@@ -2,5 +2,8 @@
 ##### There's no free surface with a rigid lid
 #####
 
+@inline materialize_free_surface(::Nothing, velocities, grid) = nothing
+
 @inline explicit_barotropic_pressure_x_gradient(i, j, k, grid, ::Nothing) = zero(grid)
 @inline explicit_barotropic_pressure_y_gradient(i, j, k, grid, ::Nothing) = zero(grid)
+

src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl

@@ -100,6 +100,7 @@ extract_boundary_conditions(::PrescribedVelocityFields) = NamedTuple()
 free_surface_displacement_field(::PrescribedVelocityFields, ::Nothing, grid) = nothing
 HorizontalVelocityFields(::PrescribedVelocityFields, grid) = nothing, nothing
 
+materialize_free_surface(::Nothing,                      ::PrescribedVelocityFields, grid) = nothing
 materialize_free_surface(::ExplicitFreeSurface{Nothing}, ::PrescribedVelocityFields, grid) = nothing
 materialize_free_surface(::ImplicitFreeSurface{Nothing}, ::PrescribedVelocityFields, grid) = nothing
 materialize_free_surface(::SplitExplicitFreeSurface,     ::PrescribedVelocityFields, grid) = nothing

test/test_hydrostatic_free_surface_models.jl

@@ -103,6 +103,7 @@ topos_3d = ((Periodic, Periodic, Bounded),
             @test grid isa SingleColumnGrid
             @test isnothing(model.free_surface)
             @test !(:η ∈ keys(fields(model))) # doesn't include free surface
+            @test field_names(model) == keys(fields(model))
         end
     end
 
@@ -114,6 +115,7 @@ topos_3d = ((Periodic, Periodic, Bounded),
                 model = HydrostaticFreeSurfaceModel(; grid)
                 @test model isa HydrostaticFreeSurfaceModel
                 @test :η ∈ keys(fields(model)) # contrary to the SingleColumnGrid case
+                @test field_names(model) == keys(fields(model))
             end
         end
     end
@@ -129,6 +131,18 @@ topos_3d = ((Periodic, Periodic, Bounded),
         end
     end
 
+    for FreeSurface in (ExplicitFreeSurface, ImplicitFreeSurface, SpliExplicitFreeSurface, Nothing)
+        @testset "$topo model construction" begin
+            @info "  Testing $free_surface model construction..."
+            for arch in archs, FT in float_types
+                grid = RectilinearGrid(arch, FT, topology=topo, size=(1, 1, 1), extent=(1, 2, 3))
+                model = HydrostaticFreeSurfaceModel(; grid, free_surface=FreeSurface())
+                @test model isa HydrostaticFreeSurfaceModel
+                @test field_names(model) == keys(fields(model))
+            end
+        end
+    end
+
     @testset "Halo size check in model constructor" begin
         for topo in topos_3d
             grid = RectilinearGrid(topology=topo, size=(1, 1, 1), extent=(1, 2, 3), halo=(1, 1, 1))