Add PotentialEnergy methods for BT and linearEOS

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,31 +5,22 @@ 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 Oceananigans.BuoyancyModels: Buoyancy, BuoyancyTracer, SeawaterBuoyancy, LinearEquationOfState
-using Oceananigans.Models: NonhydrostaticModel, HydrostaticFreeSurfaceModel
-using Oceanostics: validate_location, validate_buoyancy
+using Oceananigans.BuoyancyModels: buoyancy_perturbationᶜᶜᶜ
+using Oceananigans.Models: ShallowWaterModel
+using Oceanostics: validate_location
 using SeawaterPolynomials: BoussinesqEquationOfState
 
-const ModelsWithBoussinesqEOS = Union{NonhydrostaticModel{TS, E, A, G, T,
-                                    <:Buoyancy{<:SeawaterBuoyancy{FT, <:BoussinesqEquationOfState, Temp, Sal}},
-                                    R, SD, U, C, Φ, F, V, S, K, BG, P, BGC, I, AF},
-                                    HydrostaticFreeSurfaceModel{TS, E, A, S, G, T, V,
-                                    <:Buoyancy{<:SeawaterBuoyancy{FT, <:BoussinesqEquationOfState, Temp, Sal}},
-                                    R, F, P, BGC, U, C, Φ, K, AF}} where
-                                    {TS, E, A, G, T, R, SD, U, C, Φ, F, V, S, K, BG, P,
-                                    BGC, I, AF, FT, Temp, Sal}
-
-const ModelsWithBuoyancyTracer = Union{NonhydrostaticModel{TS, E, A, G, T,
-                                        <:Buoyancy{<:BuoyancyTracer, g},
-                                        R, SD, U, C, Φ, F, V, S, K, BG, P, BGC, I, AF},
-                                        HydrostaticFreeSurfaceModel{TS, E, A, S, G, T, V,
-                                        <:Buoyancy{<:BuoyancyTracer, g},
-                                        R, F, P, BGC, U, C, Φ, K, AF}} where
-                                        {TS, E, A, G, T, R, SD, U, C, Φ, F, V, S, K, BG, P,
-                                        BGC, I, AF, g}
+const NoBuoyancyModel = Union{Nothing, ShallowWaterModel}
+const BuoyancyTracerModel = Buoyancy{<:BuoyancyTracer, g} where g
+const BuoyancyLinearEOSModel = Buoyancy{<:SeawaterBuoyancy{FT, <:LinearEquationOfState, T, S}} where {FT, T, S}
+const BuoyancyBoussinesqEOSModel = Buoyancy{<:SeawaterBuoyancy{FT, <:BoussinesqEquationOfState, T, S}} where {FT, T, S}
+
+linear_eos_buoyancy(model) = linear_eos_buoyancy(model.grid, model.buoyancy.model, model.tracers)
+linear_eos_buoyancy(grid, buoyancy, tracers) = KernelFunctionOperation{Center, Center, Center}(buoyancy_perturbationᶜᶜᶜ, grid, buoyancy, tracers)
 
 """
     $(SIGNATURES)
@@ -38,16 +29,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 1 method)
+└── arguments: ("1×1×100 Field{Center, Center, Center} on RectilinearGrid on CPU", "KernelFunctionOperation at (Center, Center, Center)")
+```
+
+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
 
@@ -114,35 +137,48 @@ KernelFunctionOperation at (Center, Center, Center)
 └── arguments: ("KernelFunctionOperation at (Center, Center, Center)", "KernelFunctionOperation at (Center, Center, Center)", "(g=9.80665, ρ₀=1020.0)")
 ```
 """
-@inline function PotentialEnergy(model::ModelsWithBoussinesqEOS;
-                                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")
 
+    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
-    ρ = seawater_density(model; geopotential_height)
+    b = model.tracers.b
     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}(bz_ccc, grid, b, Z)
 end
 
-@inline g′z_ccc(i, j, k, grid, ρ, Z, p) = (p.g / p.ρ₀) * ρ[i, j, k] * Z[i, j, k]
-
-@inline function PotentialEnergy(model::ModelsWithBuoyancyTracer;
-                                location = (Center, Center, Center))
-
-    validate_location(location, "PotentialEnergy")
+@inline function PotentialEnergy(model, buoyancy_model::BuoyancyLinearEOSModel, geopotential_height)
 
     grid = model.grid
-    b = model.tracers.b
+    b = linear_eos_buoyancy(model)
     Z = model_geopotential_height(model)
 
     return KernelFunctionOperation{Center, Center, Center}(bz_ccc, grid, b, Z)
 end
 
 @inline bz_ccc(i, j, k, grid, b, Z) = b[i, j, k] * Z[i, j, k]
 
+@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)
+end
+
+@inline g′z_ccc(i, j, k, grid, ρ, Z, p) = (p.g / p.ρ₀) * ρ[i, j, k] * Z[i, j, k]
+
 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"