Make it easier to build simple HydrostaticFreeSurfaceModels (#3620)

* Make it easier to build simple HydrostaticFreeSurfaceModels

* Add a default buoyancy to basic HydrostaticFreeSurfaceModel test

* Update test_hydrostatic_free_surface_models.jl

* Add tracers=T in docstring

* Fix the doctest for buoyancy plus also show and usmmary

* Add summary for NegativeZDirection

* Fix doctests

* Update show for NegativeZDirection

docs/src/model_setup/buoyancy_and_equation_of_state.md

@@ -51,12 +51,10 @@ NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0)
 └── coriolis: Nothing
 ```
 
-To create a `HydrostaticFreeSurfaceModel` without a buoyancy term we explicitly
-specify `buoyancy = nothing`. The default tracers `T` and `S` for `HydrostaticFreeSurfaceModel`
-may be eliminated when `buoyancy = nothing` by specifying `tracers = ()`:
+The same is true for `HydrostaticFreeSurfaceModel`,
 
 ```jldoctest buoyancy
-julia> model = HydrostaticFreeSurfaceModel(; grid, buoyancy=nothing, tracers=())
+julia> model = HydrostaticFreeSurfaceModel(; grid)
 HydrostaticFreeSurfaceModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0)
 ├── grid: 8×8×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo
 ├── timestepper: QuasiAdamsBashforth2TimeStepper
@@ -99,7 +97,7 @@ HydrostaticFreeSurfaceModel{CPU, RectilinearGrid}(time = 0 seconds, iteration =
 ├── buoyancy: BuoyancyTracer with ĝ = NegativeZDirection()
 ├── free surface: ImplicitFreeSurface with gravitational acceleration 9.80665 m s⁻²
 │   └── solver: FFTImplicitFreeSurfaceSolver
-├── advection scheme: 
+├── advection scheme:
 │   ├── momentum: Centered reconstruction order 2
 │   └── b: Centered reconstruction order 2
 └── coriolis: Nothing
@@ -127,30 +125,11 @@ NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0)
 ├── buoyancy: SeawaterBuoyancy with g=9.80665 and LinearEquationOfState(thermal_expansion=0.000167, haline_contraction=0.00078) with ĝ = NegativeZDirection()
 └── coriolis: Nothing
 ```
-With `HydrostaticFreeSurfaceModel`, these are the default choices for `buoyancy` and `tracers` so,
-either including them or not we get:
-
-```jldoctest buoyancy
-julia> model = HydrostaticFreeSurfaceModel(; grid, buoyancy=SeawaterBuoyancy(), tracers=(:T, :S))
-HydrostaticFreeSurfaceModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0)
-├── grid: 8×8×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo
-├── timestepper: QuasiAdamsBashforth2TimeStepper
-├── tracers: (T, S)
-├── closure: Nothing
-├── buoyancy: SeawaterBuoyancy with g=9.80665 and LinearEquationOfState(thermal_expansion=0.000167, haline_contraction=0.00078) with ĝ = NegativeZDirection()
-├── free surface: ImplicitFreeSurface with gravitational acceleration 9.80665 m s⁻²
-│   └── solver: FFTImplicitFreeSurfaceSolver
-├── advection scheme: 
-│   ├── momentum: Centered reconstruction order 2
-│   ├── T: Centered reconstruction order 2
-│   └── S: Centered reconstruction order 2
-└── coriolis: Nothing
-```
 
-is identical to the default,
+and the same is true for `HydrostaticFreeSurfaceModel`,
 
 ```jldoctest buoyancy
-julia> model = HydrostaticFreeSurfaceModel(; grid)
+julia> model = HydrostaticFreeSurfaceModel(; grid, buoyancy=SeawaterBuoyancy(), tracers=(:T, :S))
 HydrostaticFreeSurfaceModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0)
 ├── grid: 8×8×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo
 ├── timestepper: QuasiAdamsBashforth2TimeStepper
@@ -159,7 +138,7 @@ HydrostaticFreeSurfaceModel{CPU, RectilinearGrid}(time = 0 seconds, iteration =
 ├── buoyancy: SeawaterBuoyancy with g=9.80665 and LinearEquationOfState(thermal_expansion=0.000167, haline_contraction=0.00078) with ĝ = NegativeZDirection()
 ├── free surface: ImplicitFreeSurface with gravitational acceleration 9.80665 m s⁻²
 │   └── solver: FFTImplicitFreeSurfaceSolver
-├── advection scheme: 
+├── advection scheme:
 │   ├── momentum: Centered reconstruction order 2
 │   ├── T: Centered reconstruction order 2
 │   └── S: Centered reconstruction order 2
@@ -245,19 +224,25 @@ we wrap the buoyancy model in
 `Buoyancy()` function call, which takes the keyword arguments `model` and `gravity_unit_vector`,
 
 ```jldoctest buoyancy
+julia> grid = RectilinearGrid(size=(8, 8, 8), extent=(1, 1, 1));
+
 julia> θ = 45; # degrees
 
 julia> g̃ = (0, sind(θ), cosd(θ));
 
-julia> model = NonhydrostaticModel(; grid,
-                                   buoyancy=Buoyancy(model=BuoyancyTracer(), gravity_unit_vector=g̃), 
-                                   tracers=:b)
+julia> buoyancy = Buoyancy(model=BuoyancyTracer(), gravity_unit_vector=g̃)
+Buoyancy:
+├── model: BuoyancyTracer
+└── gravity_unit_vector: (0.0, 0.707107, 0.707107)
+
+julia> model = NonhydrostaticModel(; grid, buoyancy, tracers=:b)
 NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0)
 ├── grid: 8×8×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo
 ├── timestepper: QuasiAdamsBashforth2TimeStepper
 ├── advection scheme: Centered reconstruction order 2
 ├── tracers: b
 ├── closure: Nothing
-├── buoyancy: BuoyancyTracer with ĝ = Tuple{Float64, Float64, Float64}
+├── buoyancy: BuoyancyTracer with ĝ = (0.0, 0.707107, 0.707107)
 └── coriolis: Nothing
 ```
+

src/BuoyancyModels/buoyancy.jl

@@ -35,7 +35,7 @@ NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0)
 ├── advection scheme: Centered reconstruction order 2
 ├── tracers: b
 ├── closure: Nothing
-├── buoyancy: BuoyancyTracer with ĝ = Tuple{Float64, Float64, Float64}
+├── buoyancy: BuoyancyTracer with ĝ = (0.0, -0.707107, -0.707107)
 └── coriolis: Nothing
 ```
 """
@@ -70,6 +70,18 @@ end
 regularize_buoyancy(b) = b
 regularize_buoyancy(b::AbstractBuoyancyModel) = Buoyancy(model=b)
 
-Base.summary(buoyancy::Buoyancy) = string(summary(buoyancy.model), " with ĝ = ", summary(buoyancy.gravity_unit_vector))
+Base.summary(buoyancy::Buoyancy) = string(summary(buoyancy.model),
+                                          " with ĝ = ",
+                                          summarize_vector(buoyancy.gravity_unit_vector))
 
-Base.show(io::IO, buoyancy::Buoyancy) = print(io, sprint(show, buoyancy.model), "\nwith `gravity_unit_vector` = ", summary(buoyancy.gravity_unit_vector))
+summarize_vector(n) = string("(", prettysummary(n[1]), ", ",
+                                  prettysummary(n[2]), ", ",
+                                  prettysummary(n[3]), ")")
+
+summarize_vector(::NegativeZDirection) = "NegativeZDirection()"
+
+function Base.show(io::IO, buoyancy::Buoyancy)
+    print(io, "Buoyancy:", '\n',
+              "├── model: ", prettysummary(buoyancy.model), '\n',
+              "└── gravity_unit_vector: ", summarize_vector(buoyancy.gravity_unit_vector))
+end

src/Fields/Fields.jl

@@ -47,6 +47,8 @@ end
 
 @inline field(loc, a::Function, grid) = FunctionField(loc, a, grid)
 @inline field(loc, a::Number, grid) = ConstantField(a)
+@inline field(loc, a::ZeroField, grid) = a
+@inline field(loc, a::ConstantField, grid) = a
 
 @inline function field(loc, f::Field, grid)
     loc === location(f) && grid === f.grid && return f

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl

@@ -106,10 +106,10 @@ function HydrostaticFreeSurfaceModel(; grid,
                                              clock = Clock{eltype(grid)}(time = 0),
                                 momentum_advection = CenteredSecondOrder(),
                                   tracer_advection = CenteredSecondOrder(),
-                                          buoyancy = SeawaterBuoyancy(eltype(grid)),
+                                          buoyancy = nothing,
                                           coriolis = nothing,
                                       free_surface = default_free_surface(grid, gravitational_acceleration=g_Earth),
-                                           tracers = (:T, :S),
+                                           tracers = nothing,
                                       forcing::NamedTuple = NamedTuple(),
                                            closure = nothing,
                    boundary_conditions::NamedTuple = NamedTuple(),

src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl

@@ -54,7 +54,7 @@ function PrescribedVelocityFields(; u = ZeroField(),
 end
 
 wrap_prescribed_field(X, Y, Z, f::Function, grid; kwargs...) = FunctionField{X, Y, Z}(f, grid; kwargs...)
-wrap_prescribed_field(X, Y, Z, f, grid; kwargs...) = f
+wrap_prescribed_field(X, Y, Z, f, grid; kwargs...) = field((X, Y, Z), f, grid)
 
 function HydrostaticFreeSurfaceVelocityFields(velocities::PrescribedVelocityFields, grid, clock, bcs)
 

src/Models/HydrostaticFreeSurfaceModels/set_hydrostatic_free_surface_model.jl

@@ -18,8 +18,8 @@ Example
 
 ```jldoctest
 using Oceananigans
-
-model = HydrostaticFreeSurfaceModel(grid = RectilinearGrid(size=(16, 16, 16), extent=(1, 1, 1)))
+grid = RectilinearGrid(size=(16, 16, 16), extent=(1, 1, 1))
+model = HydrostaticFreeSurfaceModel(; grid, tracers=:T)
 
 # Set u to a parabolic function of z, v to random numbers damped
 # at top and bottom, and T to some silly array of half zeros,

test/test_hydrostatic_free_surface_models.jl

@@ -14,16 +14,12 @@ function time_step_hydrostatic_model_works(grid;
                                            closure = nothing,
                                            velocities = nothing)
 
-    tracers = [:T, :S]
+    tracers = [:b]
+    buoyancy = BuoyancyTracer()
     closure isa CATKEVerticalDiffusivity && push!(tracers, :e)
 
-    model = HydrostaticFreeSurfaceModel(grid = grid,
-                                        momentum_advection = momentum_advection,
-                                        free_surface = free_surface,
-                                        coriolis = coriolis,
-                                        tracers = tracers,
-                                        velocities = velocities,
-                                        closure = closure)
+    model = HydrostaticFreeSurfaceModel(; grid, coriolis, tracers, velocities, buoyancy,
+                                        momentum_advection, free_surface, closure)
 
     simulation = Simulation(model, Δt=1.0, stop_iteration=1)