Generalize ErtelPotentialVorticity and RichardsonNumber to work w…

…hen `model.buoyancy` is not `BuoyancyTracer` (#190) * Generalize EPV for non-buoyancyTracer models * cleanup FlowDiagnostics * started improving tests * allow EPV to be calculated with any tracer * couple of fixes for tests * disptach properly * formatting * some other test fixes * generalize RichardsonNumber and fix tests * fixed example and bump version
tomchor · Jan 3, 2025 · 5c16f89 · 5c16f89 · tomchor · Jan 3, 2025
1 parent 56f7f8f
commit 5c16f89
Show file tree

Hide file tree

Showing 5 changed files with 142 additions and 132 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "Oceanostics"
 uuid = "d0ccf422-c8fb-49b5-a76d-74acdde946ac"
 authors = ["tomchor <[email protected]>"]
-version = "0.14.6"
+version = "0.15.0"
 
 [deps]
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"

diff --git a/docs/examples/tilted_bottom_boundary_layer.jl b/docs/examples/tilted_bottom_boundary_layer.jl
@@ -146,9 +146,10 @@ simulation.callbacks[:progress] = Callback(progress, IterationInterval(400))
 
 using Oceanostics
 
-Ri = RichardsonNumber(model, add_background=true)
+b = model.tracers.b + model.background_fields.tracers.b
+Ri = RichardsonNumber(model, model.velocities..., b)
 Ro = RossbyNumber(model)
-PV = ErtelPotentialVorticity(model, add_background=true)
+PV = ErtelPotentialVorticity(model, model.velocities..., b, model.coriolis)
 
 # Note that the calculation of these quantities depends on the alignment with the true (geophysical)
 # vertical and the rotation axis. Oceanostics already takes that into consideration by using
@@ -159,7 +160,7 @@ PV = ErtelPotentialVorticity(model, add_background=true)
 #
 # Now we write these quantities to a NetCDF file:
 
-output_fields = (; Ri, Ro, PV, b = model.tracers.b + model.background_fields.tracers.b)
+output_fields = (; Ri, Ro, PV, b)
 
 filename = "tilted_bottom_boundary_layer"
 simulation.output_writers[:nc] = NetCDFOutputWriter(model, output_fields,

diff --git a/src/FlowDiagnostics.jl b/src/FlowDiagnostics.jl
@@ -5,12 +5,16 @@ export RichardsonNumber, RossbyNumber
 export ErtelPotentialVorticity, ThermalWindPotentialVorticity, DirectionalErtelPotentialVorticity
 export StrainRateTensorModulus, VorticityTensorModulus, Q, QVelocityGradientTensorInvariant
 
-using Oceanostics: validate_location, validate_dissipative_closure, add_background_fields
+using Oceanostics: validate_location,
+                   validate_dissipative_closure,
+                   add_background_fields,
+                   get_coriolis_frequency_components
 
 using Oceananigans: NonhydrostaticModel, FPlane, ConstantCartesianCoriolis, BuoyancyField, BuoyancyTracer
 using Oceananigans.Operators
 using Oceananigans.AbstractOperations
 using Oceananigans.AbstractOperations: KernelFunctionOperation
+using Oceananigans.BuoyancyFormulations: buoyancy
 using Oceananigans.Grids: Center, Face, NegativeZDirection, ZDirection
 
 #+++ Richardson number
@@ -64,21 +68,11 @@ Calculate the Richardson Number as
 
 where `z` is the true vertical direction (ie anti-parallel to gravity).
 """
-function RichardsonNumber(model; location = (Center, Center, Face), add_background=true)
+function RichardsonNumber(model; location = (Center, Center, Face))
     validate_location(location, "RichardsonNumber", (Center, Center, Face))
-
-    if (model isa NonhydrostaticModel) & add_background
-        full_fields = add_background_fields(model)
-        u, v, w, b = full_fields.u, full_fields.v, full_fields.w, full_fields.b
-    else
-        u, v, w = model.velocities
-        b = model.tracers.b
-    end
-
-    return RichardsonNumber(model, u, v, w, b; location)
+    return RichardsonNumber(model, model.velocities..., buoyancy(model); location)
 end
 
-
 function RichardsonNumber(model, u, v, w, b; location = (Center, Center, Face))
     validate_location(location, "RichardsonNumber", (Center, Center, Face))
 
@@ -89,6 +83,11 @@ function RichardsonNumber(model, u, v, w, b; location = (Center, Center, Face))
     else
         true_vertical_direction = .-model.buoyancy.gravity_unit_vector
     end
+    return RichardsonNumber(model, u, v, w, b, true_vertical_direction; location = (Center, Center, Face))
+end
+
+function RichardsonNumber(model, u, v, w, b, true_vertical_direction; location = (Center, Center, Face))
+    validate_location(location, "RichardsonNumber", (Center, Center, Face))
     return KernelFunctionOperation{Center, Center, Face}(richardson_number_ccf, model.grid,
                                                          u, v, w, b, true_vertical_direction)
 end
@@ -145,16 +144,7 @@ function RossbyNumber(model, u, v, w, coriolis; location = (Face, Face, Face),
                       dUdy_bg=0, dVdx_bg=0)
     validate_location(location, "RossbyNumber", (Face, Face, Face))
 
-    if coriolis isa FPlane
-        fx = fy = 0
-        fz = coriolis.f
-    elseif coriolis isa ConstantCartesianCoriolis
-        fx = coriolis.fx
-        fy = coriolis.fy
-        fz = coriolis.fz
-    else
-        throw(ArgumentError("RossbyNumber only implemented for FPlane and ConstantCartesianCoriolis"))
-    end
+    fx, fy, fz = get_coriolis_frequency_components(coriolis)
 
     parameters = (; fx, fy, fz, dWdy_bg, dVdz_bg, dUdz_bg, dWdx_bg, dUdy_bg, dVdx_bg)
     return KernelFunctionOperation{Face, Face, Face}(rossby_number_fff, model.grid,
@@ -192,15 +182,17 @@ is defined as
 where `f` is the Coriolis frequency, `ωᶻ` is the relative vorticity in the `z` direction, `b` is the buoyancy, and
 `∂U/∂z` and `∂V/∂z` comprise the thermal wind shear.
 """
-function ThermalWindPotentialVorticity(model; f=nothing, location = (Face, Face, Face))
+function ThermalWindPotentialVorticity(model; tracer = :b, location = (Face, Face, Face))
     validate_location(location, "ThermalWindPotentialVorticity", (Face, Face, Face))
     u, v, w = model.velocities
-    b = BuoyancyField(model)
-    if isnothing(f)
-        f = model.coriolis.f
-    end
+    return ThermalWindPotentialVorticity(model, u, v, model.tracers[tracer], model.coriolis; location)
+end
+
+function ThermalWindPotentialVorticity(model, u, v, tracer, coriolis; location = (Face, Face, Face))
+    validate_location(location, "ThermalWindPotentialVorticity", (Face, Face, Face))
+    fx, fy, fz = get_coriolis_frequency_components(coriolis)
     return KernelFunctionOperation{Face, Face, Face}(potential_vorticity_in_thermal_wind_fff, model.grid,
-                                                     u, v, b, f)
+                                                     u, v, tracer, fz)
 end
 
 @inline function ertel_potential_vorticity_fff(i, j, k, grid, u, v, w, b, fx, fy, fz)
@@ -279,45 +271,16 @@ Note that EPV values are correctly calculated both in the interior and the bound
 interior and top boundary, EPV = f×N² = 10⁻¹⁰, while EPV = 0 at the bottom boundary since ∂b/∂z
 is zero there.
 """
-function ErtelPotentialVorticity(model; location = (Face, Face, Face), add_background = true)
+function ErtelPotentialVorticity(model; tracer = :b, location = (Face, Face, Face))
     validate_location(location, "ErtelPotentialVorticity", (Face, Face, Face))
-
-    if model.buoyancy == nothing || !(model.buoyancy.formulation isa BuoyancyTracer)
-        throw(ArgumentError("`ErtelPotentialVorticity` is only implemented for `BuoyancyTracer` at the moment."))
-    end
-
-    u, v, w = model.velocities
-    b = model.tracers.b
-
-    if (model isa NonhydrostaticModel) & add_background
-        full_fields = add_background_fields(model)
-        u, v, w, b = full_fields.u, full_fields.v, full_fields.w, full_fields.b
-    else
-        u, v, w = model.velocities
-        b = model.tracers.b
-    end
-
-    return ErtelPotentialVorticity(model, u, v, w, b, model.coriolis; location)
+    return ErtelPotentialVorticity(model, model.velocities..., model.tracers[tracer], model.coriolis; location)
 end
 
-function ErtelPotentialVorticity(model, u, v, w, b, coriolis; location = (Face, Face, Face))
+function ErtelPotentialVorticity(model, u, v, w, tracer, coriolis; location = (Face, Face, Face))
     validate_location(location, "ErtelPotentialVorticity", (Face, Face, Face))
-
-    if coriolis isa FPlane
-        fx = fy = 0
-        fz = coriolis.f
-    elseif coriolis isa ConstantCartesianCoriolis
-        fx = coriolis.fx
-        fy = coriolis.fy
-        fz = coriolis.fz
-    elseif coriolis == nothing
-        fx = fy = fz = 0
-    else
-        throw(ArgumentError("ErtelPotentialVorticity is only implemented for FPlane and ConstantCartesianCoriolis"))
-    end
-
+    fx, fy, fz = get_coriolis_frequency_components(coriolis)
     return KernelFunctionOperation{Face, Face, Face}(ertel_potential_vorticity_fff, model.grid,
-                                                     u, v, w, b, fx, fy, fz)
+                                                     u, v, w, tracer, fx, fy, fz)
 end
 
 @inline function directional_ertel_potential_vorticity_fff(i, j, k, grid, u, v, w, b, params)
@@ -356,45 +319,21 @@ basde on a `model` and a `direction`. The Ertel Potential Vorticity is defined a
 where ωₜₒₜ is the total (relative + planetary) vorticity vector, `b` is the buoyancy and ∇ is the gradient
 operator.
 """
-function DirectionalErtelPotentialVorticity(model, direction; location = (Face, Face, Face))
+function DirectionalErtelPotentialVorticity(model, direction; tracer = :b, location = (Face, Face, Face))
     validate_location(location, "DirectionalErtelPotentialVorticity", (Face, Face, Face))
-
-    if model.buoyancy == nothing || !(model.buoyancy.formulation isa BuoyancyTracer)
-        throw(ArgumentError("`DirectionalErtelPotentialVorticity` is only implemented for `BuoyancyTracer` at the moment."))
-    end
-
-    if model isa NonhydrostaticModel
-        full_fields = add_background_fields(model)
-        u, v, w, b = full_fields.u, full_fields.v, full_fields.w, full_fields.b
-    else
-        u, v, w = model.velocities
-        b = model.tracers.b
-    end
-
-    return DirectionalErtelPotentialVorticity(model, direction, u, v, w, b, model.coriolis; location)
+    return DirectionalErtelPotentialVorticity(model, direction, model.velocities..., model.tracers[tracer], model.coriolis; location)
 end
 
 
-function DirectionalErtelPotentialVorticity(model, direction, u, v, w, b, coriolis; location = (Face, Face, Face))
+function DirectionalErtelPotentialVorticity(model, direction, u, v, w, tracer, coriolis; location = (Face, Face, Face))
     validate_location(location, "DirectionalErtelPotentialVorticity", (Face, Face, Face))
 
-    if coriolis isa FPlane
-        fx = fy = 0
-        fz = coriolis.f
-    elseif coriolis isa ConstantCartesianCoriolis
-        fx = coriolis.fx
-        fy = coriolis.fy
-        fz = coriolis.fz
-    elseif coriolis == nothing
-        fx = fy = fz = 0
-    else
-        throw(ArgumentError("ErtelPotentialVorticity is only implemented for FPlane and ConstantCartesianCoriolis"))
-    end
+    fx, fy, fz = get_coriolis_frequency_components(coriolis)
     f_dir = sum([fx, fy, fz] .* direction)
 
     dir_x, dir_y, dir_z = direction
     return KernelFunctionOperation{Face, Face, Face}(directional_ertel_potential_vorticity_fff, model.grid,
-                                                     u, v, w, b, (; f_dir, dir_x, dir_y, dir_z))
+                                                     u, v, w, tracer, (; f_dir, dir_x, dir_y, dir_z))
 end
 #---
 
@@ -425,7 +364,7 @@ symmetric part of the velocity gradient tensor:
 Its modulus is then defined (using Einstein summation notation) as
 
 ```
-    || Sᵢⱼ || = √( Sᵢⱼ Sᵢⱼ)
+    || Sᵢⱼ || = √(Sᵢⱼ Sᵢⱼ)
 ```
 """
 function StrainRateTensorModulus(model; location = (Center, Center, Center))
@@ -460,7 +399,7 @@ antisymmetric part of the velocity gradient tensor:
 Its modulus is then defined (using Einstein summation notation) as
 
 ```
-    || Ωᵢⱼ || = √( Ωᵢⱼ Ωᵢⱼ)
+    || Ωᵢⱼ || = √(Ωᵢⱼ Ωᵢⱼ)
 ```
 """
 function VorticityTensorModulus(model; location = (Center, Center, Center))
@@ -491,7 +430,7 @@ gradient tensor `∂ⱼuᵢ`:
 from where `Q` is defined as
 
 ```
-    Q = ½ ( ΩᵢⱼΩᵢⱼ - SᵢⱼSᵢⱼ)
+    Q = ½ (ΩᵢⱼΩᵢⱼ - SᵢⱼSᵢⱼ)
 ```
 and where `Sᵢⱼ= ½(∂ⱼuᵢ + ∂ᵢuⱼ)` and `Ωᵢⱼ= ½(∂ⱼuᵢ - ∂ᵢuⱼ)`. More info about it can be found in
 doi:10.1063/1.5124245.

diff --git a/src/Oceanostics.jl b/src/Oceanostics.jl
@@ -57,15 +57,15 @@ function add_background_fields(model)
     velocities = model.velocities
     # Adds background velocities to their perturbations only if background velocity isn't ZeroField
     full_velocities = NamedTuple{keys(velocities)}((model.background_fields.velocities[key] isa ZeroField) ?
-                                                   val :
-                                                   Field(val + model.background_fields.velocities[key])
-                                                   for (key,val) in zip(keys(velocities), velocities))
+                                                    val :
+                                                    Field(val + model.background_fields.velocities[key])
+                                                    for (key, val) in zip(keys(velocities), velocities))
     tracers = model.tracers
     # Adds background tracer fields to their perturbations only if background tracer field isn't ZeroField
     full_tracers = NamedTuple{keys(tracers)}((model.background_fields.tracers[key] isa ZeroField) ?
                                               val :
                                               Field(val + model.background_fields.tracers[key])
-                                              for (key,val) in zip(keys(tracers), tracers))
+                                              for (key, val) in zip(keys(tracers), tracers))
 
     return merge(full_velocities, full_tracers)
 end
@@ -98,6 +98,27 @@ function perturbation_fields(model; kwargs...)
 end
 #---
 
+#+++ Utils for Coriolis frequency
+using Oceananigans: FPlane, ConstantCartesianCoriolis, AbstractModel
+function get_coriolis_frequency_components(coriolis)
+    if coriolis isa FPlane
+        fx = fy = 0
+        fz = coriolis.f
+    elseif coriolis isa ConstantCartesianCoriolis
+        fx = coriolis.fx
+        fy = coriolis.fy
+        fz = coriolis.fz
+    elseif coriolis == nothing
+        fx = fy = fz = 0
+    else
+        throw(ArgumentError("Extraction of rotation components is only implemented for `FPlane`, `ConstantCartesianCoriolis` and `nothing`."))
+    end
+    return fx, fy, fz
+end
+
+get_coriolis_frequency_components(model::AbstractModel) = get_coriolis_frequency_components(model.coriolis)
+#---
+
 #+++ A few utils for closure tuples:
 using Oceananigans.TurbulenceClosures: νᶜᶜᶜ