Adds a mixed layer depth diagnostic

src/FlowDiagnostics.jl

@@ -4,18 +4,23 @@ using DocStringExtensions
 export RichardsonNumber, RossbyNumber
 export ErtelPotentialVorticity, ThermalWindPotentialVorticity, DirectionalErtelPotentialVorticity
 export StrainRateTensorModulus, VorticityTensorModulus, Q, QVelocityGradientTensorInvariant
+export MixedLayerDepth, BuoyancyAnomalyCriterion, DensityAnomalyCriterion
 
 using Oceanostics: validate_location,
                    validate_dissipative_closure,
                    add_background_fields,
                    get_coriolis_frequency_components
 
 using Oceananigans: NonhydrostaticModel, FPlane, ConstantCartesianCoriolis, BuoyancyField, BuoyancyTracer
+using Oceananigans.BuoyancyFormulations: get_temperature_and_salinity, SeawaterBuoyancy, g_Earth, buoyancy_perturbationᶜᶜᶜ
 using Oceananigans.Operators
 using Oceananigans.AbstractOperations
 using Oceananigans.AbstractOperations: KernelFunctionOperation
 using Oceananigans.BuoyancyFormulations: buoyancy
-using Oceananigans.Grids: Center, Face, NegativeZDirection, ZDirection
+using Oceananigans.Grids: AbstractGrid, Center, Face, NegativeZDirection, ZDirection, znode
+
+using SeawaterPolynomials: ρ′, BoussinesqEquationOfState
+using SeawaterPolynomials.SecondOrderSeawaterPolynomials: SecondOrderSeawaterPolynomial
 
 #+++ Richardson number
 @inline ψ²(i, j, k, grid, ψ) = @inbounds ψ[i, j, k]^2
@@ -337,7 +342,7 @@ function DirectionalErtelPotentialVorticity(model, direction, u, v, w, tracer, c
 end
 #---
 
-#+++ Velocity gradient tensor
+#+++ Velocity gradient and vorticity tensors
 @inline fψ_plus_gφ²(i, j, k, grid, f, ψ, g, φ) = (f(i, j, k, grid, ψ) + g(i, j, k, grid, φ))^2
 
 function strain_rate_tensor_modulus_ccc(i, j, k, grid, u, v, w)
@@ -372,7 +377,6 @@ function StrainRateTensorModulus(model; location = (Center, Center, Center))
     return KernelFunctionOperation{Center, Center, Center}(strain_rate_tensor_modulus_ccc, model.grid, model.velocities...)
 end
 
-
 @inline fψ_minus_gφ²(i, j, k, grid, f, ψ, g, φ) = (f(i, j, k, grid, ψ) - g(i, j, k, grid, φ))^2
 
 function vorticity_tensor_modulus_ccc(i, j, k, grid, u, v, w)
@@ -414,7 +418,9 @@ end
     Ω² = vorticity_tensor_modulus_ccc(i, j, k, grid, u, v, w)^2
     return (Ω² - S²) / 2
 end
+#---
 
+#+++ Mixed layer depth
 """
     $(SIGNATURES)
 
@@ -441,6 +447,135 @@ function QVelocityGradientTensorInvariant(model; location = (Center, Center, Cen
 end
 
 const Q = QVelocityGradientTensorInvariant
+
+"""
+    $(TYPEDEF)
+"""
+struct MixedLayerDepthKernel{C}
+    criterion::C
+end
+
+"""
+    $(SIGNATURES)
+
+Returns the mixed layer depth defined as the depth at which `criterion` is true.
+
+Defaults to `DensityAnomalyCriterion` where the depth is that at which the density
+is some threshold (defaults to 0.125kg/m³) higher than the surface density.
+
+When `DensityAnomalyCriterion` is used, the arguments `buoyancy_formulation` and `C` should be
+supplied where `buoyancy_formulation` should be the buoyancy model, and `C` should be a named
+tuple of `(; T, S)`, `(; T)` or `(; S)` (the latter two if the buoyancy model
+specifies a constant salinity or temperature).
+"""
+function MixedLayerDepth(grid::AbstractGrid, args...; criterion = BuoyancyAnomalyCriterion(convert(eltype(grid), -1e-4 * g_Earth)))
+    validate_criterion_model(criterion, args...)
+    MLD = MixedLayerDepthKernel(criterion)
+    return KernelFunctionOperation{Center, Center, Nothing}(MLD, grid, args...)
+end
+
+function (MLD::MixedLayerDepthKernel)(i, j, k, grid, args...)
+    kₘₗ = -1
+
+    for k in grid.Nz-1:-1:1
+        below_mixed_layer = MLD.criterion(i, j, k, grid, args...)
+        kₘₗ = ifelse(below_mixed_layer & (kₘₗ < 0), k, kₘₗ)
+    end
+
+    zₘₗ = interpolate_from_nearest_cell(MLD.criterion, i, j, kₘₗ, grid, args...)
+    return ifelse(kₘₗ == -1, -Inf, zₘₗ)
+end
+
+"""
+    $(TYPEDEF)
+
+An abstract mixed layer depth criterion where the mixed layer is defined to be
+`anomaly` + `threshold` greater than the surface value of `anomaly`.
+
+`AbstractAnomalyCriterion` types should provide a method for the function `anomaly` in the form
+`anomaly(criterion, i, j, k, grid, args...)`, and should have a property `threshold`.
+"""
+abstract type AbstractAnomalyCriterion end
+
+@inline function (criterion::AbstractAnomalyCriterion)(i, j, k, grid, args...)
+    δ = criterion.threshold
+
+    ref = (anomaly(criterion, i, j, grid.Nz, grid, args...) + anomaly(criterion, i, j, grid.Nz+1, grid, args...)) * convert(eltype(grid), 0.5)
+    val = anomaly(criterion, i, j, k, grid,args...)
+
+    return val < ref + δ
+end
+
+@inline function interpolate_from_nearest_cell(criterion::AbstractAnomalyCriterion, i, j, k, grid, args...)
+    δ = criterion.threshold
+
+    ref = (anomaly(criterion, i, j, grid.Nz, grid, args...) + anomaly(criterion, i, j, grid.Nz + 1, grid, args...)) * convert(eltype(grid), 0.5)
+
+    k_val  = anomaly(criterion, i, j, k, grid, args...)
+    k⁺_val = anomaly(criterion, i, j, k + 1, grid, args...)
+
+    zₖ = znode(i, j, k, grid, Center(), Center(), Center())
+    z₊ = znode(i, j, k+1, grid, Center(), Center(), Center())
+
+    return zₖ + (z₊ - zₖ) * (ref + δ - k_val) / (k⁺_val - k_val)
+end
+
+"""
+    $(TYPEDEF)
+
+Defines the mixed layer to be the depth at which the buoyancy is more than `threshold` greater than
+the surface buoyancy (but the pertubaton is usually negative).
+
+When this model is used, the arguments `buoyancy_formulation` and `C` should be supplied where `C`
+should be the named tuple `(; b)`, with `b` the buoyancy tracer.
+"""
+@kwdef struct BuoyancyAnomalyCriterion{FT} <: AbstractAnomalyCriterion
+    threshold :: FT = -1e-4 * g_Earth
+end
+
+validate_criterion_model(::BuoyancyAnomalyCriterion, args...) =
+    @error "For BuoyancyAnomalyCriterion you must supply the arguments `buoyancy_formulation` and `C`, where `C` is the named tuple `(; b)`, with `b` the buoyancy tracer."
+
+validate_criterion_model(::BuoyancyAnomalyCriterion, buoyancy_formulation, C) = nothing
+
+@inline anomaly(::BuoyancyAnomalyCriterion, i, j, k, grid, buoyancy_formulation, C) = buoyancy_perturbationᶜᶜᶜ(i, j, k, grid, buoyancy_formulation, C)
+
+"""
+    $(TYPEDEF)
+
+Defines the mixed layer to be the depth at which the density is more than `threshold`
+greater than the surface density.
+
+When this model is used, the arguments `buoyancy_formulation` and `C` should be supplied where
+`buoyancy_formulation` should be the buoyancy model, and `C` should be a named tuple of `(; T, S)`,
+`(; T)` or `(; S)` (the latter two if the buoyancy model specifies a constant salinity or
+temperature).
+"""
+@kwdef struct DensityAnomalyCriterion{FT} <: AbstractAnomalyCriterion
+             reference_density :: FT = 1020.0
+    gravitational_acceleration :: FT = g_Earth
+                     threshold :: FT = 0.125
+end
+
+function DensityAnomalyCriterion(buoyancy_formulation::SeawaterBuoyancy{<:Any, <:BoussinesqEquationOfState}; threshold = 0.125)
+    ρᵣ = buoyancy_formulation.equation_of_state.reference_density
+    g  = buoyancy_formulation.gravitational_acceleration
+
+    return DensityAnomalyCriterion(ρᵣ, g, threshold)
+end
+
+validate_criterion_model(::DensityAnomalyCriterion, args...) =
+    @error "For DensityAnomalyCriterion you must supply the arguments buoyancy_formulation and C, where C is a named tuple of (; T, S), (; T) or (; S)"
+
+validate_criterion_model(::DensityAnomalyCriterion, buoyancy_formulation, C) = nothing
+
+@inline function anomaly(criterion::DensityAnomalyCriterion, i, j, k, grid, buoyancy_formulation, C)
+    b = buoyancy_perturbationᶜᶜᶜ(i, j, k, grid, buoyancy_formulation, C)
+
+    ρᵣ = criterion.reference_density
+    g  = criterion.gravitational_acceleration
+    return - ρᵣ * b / g
+end
 #---
 
 end # module

src/Oceanostics.jl

@@ -20,6 +20,7 @@ export RichardsonNumber, RossbyNumber
 export ErtelPotentialVorticity, ThermalWindPotentialVorticity
 export DirectionalErtelPotentialVorticity
 export StrainRateTensorModulus, VorticityTensorModulus, Q, QVelocityGradientTensorInvariant
+export MixedLayerDepth, BuoyancyAnomalyCriterion, DensityAnomalyCriterion
 #---
 
 #+++ PotentialEnergyEquationTerms exports

test/test_tracer_diagnostics.jl

@@ -2,16 +2,19 @@ using Test
 using CUDA: has_cuda_gpu, @allowscalar
 
 using Oceananigans
+using Oceananigans: fill_halo_regions!
 using Oceananigans.AbstractOperations: AbstractOperation
 using Oceananigans.Fields: @compute
+using Oceananigans.Grids: znode
 using Oceananigans.TurbulenceClosures.Smagorinskys: LagrangianAveraging
 using Oceananigans.BuoyancyFormulations: buoyancy
-using SeawaterPolynomials: RoquetEquationOfState, TEOS10EquationOfState
+using SeawaterPolynomials: RoquetEquationOfState, TEOS10EquationOfState, BoussinesqEquationOfState
+using SeawaterPolynomials.SecondOrderSeawaterPolynomials: LinearRoquetSeawaterPolynomial
 
 using Oceanostics
 using Oceanostics: get_coriolis_frequency_components
 
-const LinearBuoyancyForce = Union{BuoyancyTracer, SeawaterBuoyancy{<:Any, <:LinearEquationOfState}}
+LinearBuoyancyForce = Union{BuoyancyTracer, SeawaterBuoyancy{<:Any, <:LinearEquationOfState}}
 
 #+++ Default grids
 arch = has_cuda_gpu() ? GPU() : CPU()
@@ -172,6 +175,47 @@ function test_tracer_diagnostics(model)
     return nothing
 end
 
+function test_mixed_layer_depth(grid, buoyancy; zₘₓₗ = 0.5, δb = -1e-4 * 9.81, naive_thermal_expansion=0.000167)
+    density_is_defined = (!(buoyancy isa BuoyancyTracer)) && (buoyancy.equation_of_state isa BoussinesqEquationOfState)
+    ∂z_b = - δb / zₘₓₗ
+
+    if buoyancy isa BuoyancyTracer
+        boundary_conditions = FieldBoundaryConditions(grid, (Center, Center, Center); top = GradientBoundaryCondition(∂z_b))
+        C = (; b = CenterField(grid; boundary_conditions))
+
+    else
+        g = buoyancy.gravitational_acceleration
+        ∂z_T = ∂z_b / (g * naive_thermal_expansion)
+
+        boundary_conditions = FieldBoundaryConditions(grid, (Center, Center, Center); top = GradientBoundaryCondition(∂z_T))
+        C = (; T = CenterField(grid; boundary_conditions), S = CenterField(grid))
+    end
+
+    mld_b = MixedLayerDepth(grid, buoyancy, C; criterion = BuoyancyAnomalyCriterion(δb))
+
+    if density_is_defined
+        ρᵣ = buoyancy.equation_of_state.reference_density
+        δρ = - δb * ρᵣ / g
+
+        criterion = DensityAnomalyCriterion(buoyancy; threshold = convert(eltype(grid), δρ))
+        mld_ρ = MixedLayerDepth(grid, buoyancy, C; criterion)
+    end 
+
+    @test isinf(mld_b[1, 1])
+    density_is_defined && (@test isinf(mld_ρ[1, 1]) | (mld_ρ[1, 1] < znode(1, 1, 1, grid, Center(), Center(), Face()))) # for TEOS10 we don't get -Inf just a really deep depth
+
+    if buoyancy isa BuoyancyTracer
+        set!(C.b, (x, y, z) -> z * ∂z_b)
+    else
+        set!(C.T, (x, y, z) -> z * ∂z_T + 10)
+        set!(C.S, 35) # TEOS10SeawaterPolynomial doesn't seem to like it when this is zero
+    end
+
+    fill_halo_regions!(C)
+
+    @test isapprox(mld_b[1, 1], -zₘₓₗ + znode(1, 1, grid.Nz+1, grid, Center(), Center(), Face()), atol=0.02) # high tollerance from the approximation in ∂z_T
+    density_is_defined && (@test isapprox(mld_ρ[1, 1], -zₘₓₗ + znode(1, 1, grid.Nz+1, grid, Center(), Center(), Face()), atol=0.02)) # high tollerance from the approximation in ∂z_T
+end
 #---
 
 @testset "Tracer diagnostics tests" begin
@@ -201,6 +245,13 @@ end
                     @info "            Testing buoyancy diagnostics"
                     test_buoyancy_diagnostics(model)
                 end
+
+                @info "            Testing mixed layer depth diagnostic"
+                if !isnothing(buoyancy)
+                    test_mixed_layer_depth(grid, buoyancy)
+                else
+                    @test_throws ErrorException test_mixed_layer_depth(grid, buoyancy)
+                end
             end
         end
     end