Merge pull request #167 from tomchor/jib-dispatch-on-buoyancymodel

Add `PotentialEnergy` methods for `BuoyancyTracer` and `SeawaterBuoyancy{<:LinearEquationOfState}`

Author: jbisits

src/Oceanostics.jl

@@ -33,8 +33,6 @@ export ProgressMessengers
 #+++ Utils for validation
 # Right now, all kernels must be located at ccc
 using Oceananigans.TurbulenceClosures: AbstractScalarDiffusivity, ThreeDimensionalFormulation
-using Oceananigans.BuoyancyModels: Buoyancy, BuoyancyTracer, SeawaterBuoyancy, LinearEquationOfState
-using SeawaterPolynomials: BoussinesqEquationOfState
 using Oceananigans.Grids: Center, Face
 
 validate_location(location, type, valid_location=(Center, Center, Center)) =
@@ -45,13 +43,6 @@ validate_dissipative_closure(closure) = error("Cannot calculate dissipation rate
 validate_dissipative_closure(::AbstractScalarDiffusivity{<:Any, ThreeDimensionalFormulation}) = nothing
 validate_dissipative_closure(closure_tuple::Tuple) = Tuple(validate_dissipative_closure(c) for c in closure_tuple)
 
-validate_buoyancy(buoyancy::Nothing) = throw(ArgumentError("Cannot calculate gravitational potential energy without a buoyancy model."))
-validate_buoyancy(buoyancy::Buoyancy{<:BuoyancyTracer, g}) where g =
-    throw(ArgumentError("Cannot calculate gravitational potential energy for the buoyancy model $(buoyancy.model)."))
-validate_buoyancy(buoyancy::Buoyancy{<:SeawaterBuoyancy{FT, <:LinearEquationOfState, T, S}, g}) where {FT, T, S, g} =
-    throw(ArgumentError("To calculate gravitational potential energy with a linear equation of state use a linear `BoussinesqEquationOfState`."))
-validate_buoyancy(buoyancy::Buoyancy{<:SeawaterBuoyancy{FT, <:BoussinesqEquationOfState, T, S}, g}) where {FT, T, S, g} = nothing
-
 #---
 
 #+++ Utils for background fields

src/PotentialEnergyEquationTerms.jl

@@ -5,11 +5,24 @@ using DocStringExtensions
 export PotentialEnergy
 
 using Oceananigans.AbstractOperations: KernelFunctionOperation
-using Oceananigans: Models.seawater_density
-using Oceananigans: Models.model_geopotential_height
+using Oceananigans.Models: seawater_density
+using Oceananigans.Models: model_geopotential_height
 using Oceananigans.Grids: Center, Face
-using Oceanostics: validate_location, validate_buoyancy
-
+using Oceananigans.Grids: NegativeZDirection
+using Oceananigans.BuoyancyModels: Buoyancy, BuoyancyTracer, SeawaterBuoyancy, LinearEquationOfState
+using Oceananigans.BuoyancyModels: buoyancy_perturbationᶜᶜᶜ, Zᶜᶜᶜ
+using Oceananigans.Models: ShallowWaterModel
+using Oceanostics: validate_location
+using SeawaterPolynomials: BoussinesqEquationOfState
+
+const NoBuoyancyModel = Union{Nothing, ShallowWaterModel}
+const BuoyancyTracerModel = Buoyancy{<:BuoyancyTracer, g} where g
+const BuoyancyLinearEOSModel = Buoyancy{<:SeawaterBuoyancy{FT, <:LinearEquationOfState, T, S} where {FT, T, S}, g} where {g}
+const BuoyancyBoussinesqEOSModel = Buoyancy{<:SeawaterBuoyancy{FT, <:BoussinesqEquationOfState, T, S} where {FT, T, S}, g} where {g}
+
+validate_gravity_unit_vector(gravity_unit_vector::NegativeZDirection) = nothing
+validate_gravity_unit_vector(gravity_unit_vector) =
+    throw(ArgumentError("`PotentialEnergy` is curently only defined for models that have a `NegativeZDirection` gravity unit vector."))
 
 """
     $(SIGNATURES)
@@ -18,16 +31,48 @@ Return a `KernelFunctionOperation` to compute the `PotentialEnergy` per unit vol
 ```math
 Eₚ = \\frac{gρz}{ρ₀}
 ```
-at each grid `location` in `model`.
+at each grid `location` in `model`. `PotentialEnergy` is defined for both `BuoyancyTracer`
+and `SeawaterBuoyancy`. See the relevant Oceananigans.jl documentation on
+[buoyancy models](https://clima.github.io/OceananigansDocumentation/dev/model_setup/buoyancy_and_equation_of_state/)
+for more information about available options.
 
-**NOTE:** A `BoussinesqEquationOfState` must be used in the `model` to calculate
-`seawater_density`. See the [relevant documentation](https://clima.github.io/OceananigansDocumentation/dev/model_setup/buoyancy_and_equation_of_state/#Idealized-nonlinear-equations-of-state)
-for how to set `SeawaterBuoyancy` using a `BoussinesqEquationOfState`.
+The optional keyword argument `geopotential_height` is only used
+if ones wishes to calculate `Eₚ` with a potential density referenced to `geopotential_height`,
+rather than in-situ density, when using a `BoussinesqEquationOfState`.
 
 Example
 =======
 
-The default behaviour of `PotentialEnergy` uses the *in-situ density* in the calculation:
+Usage with a `BuoyancyTracer` buoyacny model
+```jldoctest
+julia> using Oceananigans
+
+julia> using Oceanostics.PotentialEnergyEquationTerms: PotentialEnergy
+
+julia> grid = RectilinearGrid(size=100, z=(-1000, 0), topology=(Flat, Flat, Bounded))
+1×1×100 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
+├── Flat x
+├── Flat y
+└── Bounded  z ∈ [-1000.0, 0.0]   regularly spaced with Δz=10.0
+
+julia> model = NonhydrostaticModel(; grid, buoyancy=BuoyancyTracer(), tracers=(:b,))
+NonhydrostaticModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0)
+├── grid: 1×1×100 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
+├── timestepper: QuasiAdamsBashforth2TimeStepper
+├── tracers: b
+├── closure: Nothing
+├── buoyancy: BuoyancyTracer with ĝ = NegativeZDirection()
+└── coriolis: Nothing
+
+julia> PotentialEnergy(model)
+KernelFunctionOperation at (Center, Center, Center)
+├── grid: 1×1×100 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
+├── kernel_function: bz_ccc (generic function with 2 methods)
+└── arguments: ("1×1×100 Field{Center, Center, Center} on RectilinearGrid on CPU",)
+```
+
+The default behaviour of `PotentialEnergy` uses the *in-situ density* in the calculation
+when the equation of state is a `BoussinesqEquationOfState`:
 ```jldoctest
 julia> using Oceananigans, SeawaterPolynomials.TEOS10
 
@@ -65,7 +110,7 @@ julia> PotentialEnergy(model)
 KernelFunctionOperation at (Center, Center, Center)
 ├── grid: 1×1×100 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
 ├── kernel_function: g′z_ccc (generic function with 1 method)
-└── arguments: ("KernelFunctionOperation at (Center, Center, Center)", "KernelFunctionOperation at (Center, Center, Center)", "(g=9.80665, ρ₀=1020.0)")
+└── arguments: ("KernelFunctionOperation at (Center, Center, Center)", "(g=9.80665, ρ₀=1020.0)")
 ```
 
 To use a reference density set a constant value for the keyword argument `geopotential_height`
@@ -87,27 +132,56 @@ julia> model = NonhydrostaticModel(; grid, buoyancy, tracers);
 
 julia> geopotential_height = 0; # density variable will be σ₀
 
-julia> PotentialEnergy(model; geopotential_height)
+julia> PotentialEnergy(model)
 KernelFunctionOperation at (Center, Center, Center)
 ├── grid: 1×1×100 RectilinearGrid{Float64, Flat, Flat, Bounded} on CPU with 0×0×3 halo
 ├── kernel_function: g′z_ccc (generic function with 1 method)
-└── arguments: ("KernelFunctionOperation at (Center, Center, Center)", "KernelFunctionOperation at (Center, Center, Center)", "(g=9.80665, ρ₀=1020.0)")
+└── arguments: ("KernelFunctionOperation at (Center, Center, Center)", "(g=9.80665, ρ₀=1020.0)")
 ```
 """
-@inline function PotentialEnergy(model; geopotential_height = model_geopotential_height(model), location = (Center, Center, Center))
+@inline function PotentialEnergy(model; location = (Center, Center, Center),
+                                 geopotential_height = model_geopotential_height(model))
 
     validate_location(location, "PotentialEnergy")
-    validate_buoyancy(model.buoyancy)
+    isnothing(model.buoyancy) ? nothing : validate_gravity_unit_vector(model.buoyancy.gravity_unit_vector)
+
+    return PotentialEnergy(model, model.buoyancy, geopotential_height)
+end
+
+@inline PotentialEnergy(model, buoyancy_model::NoBuoyancyModel, geopotential_height) =
+    throw(ArgumentError("Cannot calculate gravitational potential energy without a Buoyancy model."))
+
+@inline function PotentialEnergy(model, buoyancy_model::BuoyancyTracerModel, geopotential_height)
+
+    grid = model.grid
+    b = model.tracers.b
+
+    return KernelFunctionOperation{Center, Center, Center}(bz_ccc, grid, b)
+end
+
+@inline bz_ccc(i, j, k, grid, b) = b[i, j, k] * Zᶜᶜᶜ(i, j, k, grid)
+
+@inline function PotentialEnergy(model, buoyancy_model::BuoyancyLinearEOSModel, geopotential_height)
+
+    grid = model.grid
+    C = model.tracers
+    b = buoyancy_model.model
+
+    return KernelFunctionOperation{Center, Center, Center}(bz_ccc, grid, b, C)
+end
+
+@inline bz_ccc(i, j, k, grid, b, C) = buoyancy_perturbationᶜᶜᶜ(i, j, k, grid, b, C) * Zᶜᶜᶜ(i, j, k, grid)
+
+@inline function PotentialEnergy(model, buoyancy_model::BuoyancyBoussinesqEOSModel, geopotential_height)
 
     grid = model.grid
     ρ = seawater_density(model; geopotential_height)
-    Z = model_geopotential_height(model)
     parameters = (g = model.buoyancy.model.gravitational_acceleration,
                   ρ₀ = model.buoyancy.model.equation_of_state.reference_density)
 
-    return KernelFunctionOperation{Center, Center, Center}(g′z_ccc, grid, ρ, Z, parameters)
+    return KernelFunctionOperation{Center, Center, Center}(g′z_ccc, grid, ρ, parameters)
 end
 
-@inline g′z_ccc(i, j, k, grid, ρ, Z, p) = (p.g / p.ρ₀) * ρ[i, j, k] * Z[i, j, k]
+@inline g′z_ccc(i, j, k, grid, ρ, p) = (p.g / p.ρ₀) * ρ[i, j, k] * Zᶜᶜᶜ(i, j, k, grid)
 
 end # module

test/runtests.jl

@@ -11,7 +11,7 @@ using SeawaterPolynomials: RoquetEquationOfState, TEOS10EquationOfState
 using Oceanostics
 using Oceanostics: TKEBudgetTerms, TracerVarianceBudgetTerms, FlowDiagnostics, PressureRedistributionTerm, BuoyancyProductionTerm, AdvectionTerm
 using Oceanostics.TKEBudgetTerms: AdvectionTerm
-using Oceanostics: PotentialEnergy
+using Oceanostics: PotentialEnergy, PotentialEnergyEquationTerms.BuoyancyBoussinesqEOSModel
 using Oceanostics.ProgressMessengers
 
 include("test_budgets.jl")
@@ -324,15 +324,15 @@ function test_tracer_diagnostics(model)
     return nothing
 end
 
-function test_density_equation_terms_errors(model)
+function test_potential_energy_equation_terms_errors(model)
 
     @test_throws ArgumentError PotentialEnergy(model)
     @test_throws ArgumentError PotentialEnergy(model, geopotential_height = 0)
 
     return nothing
 end
 
-function test_density_equation_terms(model; geopotential_height = nothing)
+function test_potential_energy_equation_terms(model; geopotential_height = nothing)
 
     Eₚ = isnothing(geopotential_height) ? PotentialEnergy(model) :
                                           PotentialEnergy(model; geopotential_height)
@@ -342,18 +342,20 @@ function test_density_equation_terms(model; geopotential_height = nothing)
     @test Eₚ_field isa Field
     compute!(Eₚ_field)
 
-    ρ = isnothing(geopotential_height) ? Field(seawater_density(model)) :
-                                         Field(seawater_density(model; geopotential_height))
+    if model.buoyancy isa BuoyancyBoussinesqEOSModel
+        ρ = isnothing(geopotential_height) ? Field(seawater_density(model)) :
+                                             Field(seawater_density(model; geopotential_height))
 
-    compute!(ρ)
-    Z = Field(model_geopotential_height(model))
-    compute!(Z)
-    ρ₀ = model.buoyancy.model.equation_of_state.reference_density
-    g = model.buoyancy.model.gravitational_acceleration
+        compute!(ρ)
+        Z = Field(model_geopotential_height(model))
+        compute!(Z)
+        ρ₀ = model.buoyancy.model.equation_of_state.reference_density
+        g = model.buoyancy.model.gravitational_acceleration
 
-    true_value = (g / ρ₀) .* ρ.data .* Z.data
+        true_value = (g / ρ₀) .* ρ.data .* Z.data
 
-    @test isequal(Eₚ_field.data, true_value)
+        @test isequal(Eₚ_field.data, true_value)
+    end
 
     return nothing
 end
@@ -464,9 +466,9 @@ closures = (ScalarDiffusivity(ν=1e-6, κ=1e-7),
             SmagorinskyLilly(),
             (ScalarDiffusivity(ν=1e-6, κ=1e-7), AnisotropicMinimumDissipation()),)
 
-invalid_buoyancy_models = (nothing, BuoyancyTracer(), SeawaterBuoyancy())
-valid_buoyancy_models = (SeawaterBuoyancy(equation_of_state=TEOS10EquationOfState()),
-                         SeawaterBuoyancy(equation_of_state=RoquetEquationOfState(:Linear)))
+buoyancy_models = (nothing, BuoyancyTracer(), SeawaterBuoyancy(),
+                   SeawaterBuoyancy(equation_of_state=TEOS10EquationOfState()),
+                   SeawaterBuoyancy(equation_of_state=RoquetEquationOfState(:Linear)))
 
 grids = (regular_grid, stretched_grid)
 
@@ -519,23 +521,22 @@ model_types = (NonhydrostaticModel, HydrostaticFreeSurfaceModel)
                 test_tracer_diagnostics(model)
 
             end
-            @info "Testing error throws for invalid buoyancy models"
-            for buoyancy in invalid_buoyancy_models
+
+            @info "Testing `PotentialEnergy`"
+            for buoyancy in buoyancy_models
 
                 tracers = buoyancy isa BuoyancyTracer ? :b : (:S, :T)
                 model = model_type(; grid, buoyancy, tracers)
-                test_density_equation_terms_errors(model)
+                buoyancy isa BuoyancyTracer ? set!(model, b = 9.87) : set!(model, S = 34.7, T = 0.5)
+                if isnothing(buoyancy)
+                    test_potential_energy_equation_terms_errors(model)
+                else
+                    test_potential_energy_equation_terms(model)
+                    test_potential_energy_equation_terms(model, geopotential_height = 0)
+                end
 
             end
-            @info "Testing valid buoyancy models"
-            for buoyancy in valid_buoyancy_models
 
-                model = model_type(; grid, buoyancy, tracers = (:S, :T))
-                set!(model, S = 34.7, T = 0.5)
-                test_density_equation_terms(model)
-                test_density_equation_terms(model, geopotential_height = 0)
-
-            end
         end
 
         @info "Testing input validation for dissipation rates"