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: Center, Face, NegativeZDirection, ZDirection, znode
+
+using SeawaterPolynomials: ρ′, BoussinesqEquationOfState
+using SeawaterPolynomials.SecondOrderSeawaterPolynomials: SecondOrderSeawaterPolynomial
 
 #+++ Richardson number
 @inline ψ²(i, j, k, grid, ψ) = @inbounds ψ[i, j, k]^2
@@ -441,6 +446,129 @@ function QVelocityGradientTensorInvariant(model; location = (Center, Center, Cen
 end
 
 const Q = QVelocityGradientTensorInvariant
-#---
+
+struct MixedLayerDepth{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` and `C` should be 
+supplied where `buoyancy` 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, args...; criterion = BuoyancyAnomalyCriterion(convert(eltype(grid), -1/8 * 1020 / g_Earth)))
+    validate_criterion_model(criterion, args...)
+
+    MLD = MixedLayerDepth(criterion)
+
+    return KernelFunctionOperation{Center, Center, Nothing}(MLD, grid, args...)
+end
+
+function (MLD::MixedLayerDepth)(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
+
+"""
+    $(SIGNATURES)
+
+An abstract mixed layer depth criterion where the mixed layer is defined to be
+`anomaly` + `pertubation` greater than the surface value of `anomaly`.
+
+$(SIGNATURES) types should provide a method for the function `anomaly` in the form
+`anomaly(criterion, i, j, k, grid, args...)`, and should have a property `pertubation`.
+"""
+abstract type AbstractAnomalyCriterion end
+
+@inline function (criterion::AbstractAnomalyCriterion)(i, j, k, grid, args...)
+    δ = criterion.pertubation
+
+    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.pertubation
+
+    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
+
+"""
+    $(SIGNATURES)
+
+Defines the mixed layer to be the depth at which the density is more than `pertubation`
+greater than the surface density.
+
+When this model is used, the arguments `buoyancy` and `C` should be 
+supplied where `buoyancy` 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).
+"""
+struct DensityAnomalyCriterion{FT} <: AbstractAnomalyCriterion
+    pertubation :: FT
+end
+
+validate_criterion_model(::DensityAnomalyCriterion, args...) = 
+    @error "For DensityAnomalyCriterion you must supply the arguments buoyancy and C, where C is a named tuple of (; T, S), (; T) or (; S)"
+
+validate_criterion_model(::DensityAnomalyCriterion, buoyancy, C) = nothing
+
+validate_criterion_model(::DensityAnomalyCriterion, ::SeawaterBuoyancy{<:Any, <:BoussinesqEquationOfState}, C::NamedTuple) = nothing
+
+@inline function anomaly(::DensityAnomalyCriterion, i, j, k, grid, b::SeawaterBuoyancy{<:Any, <:BoussinesqEquationOfState}, C)
+    T, S = get_temperature_and_salinity(b, C)
+
+    return ρ′(i, j, k, grid, b.equation_of_state, T, S)
+end
+
+"""
+    $(SIGNATURES)
+
+Defines the mixed layer to be the depth at which the buoyancy is more than `pertubation`
+greater than the surface buoyancy (but the pertubaton is usually negative).
+
+When this model is used, the arguments `buoyancy` and `C` should be 
+supplied where `buoyancy` 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 BuoyancyAnomalyCriterion{FT} <: AbstractAnomalyCriterion
+    pertubation :: FT = - 1/8 * 1020 / g_Earth
+end
+
+validate_criterion_model(::BuoyancyAnomalyCriterion, args...) = 
+    @error "For DensityAnomalyCriterion you must supply the arguments buoyancy and C, where C is a named tuple of (; T, S), (; T) or (; S)"
+
+validate_criterion_model(::BuoyancyAnomalyCriterion, buoyancy, C) = nothing
+
+@inline anomaly(::BuoyancyAnomalyCriterion, i, j, k, grid, buoyancy, C) = buoyancy_perturbationᶜᶜᶜ(i, j, k, grid, buoyancy, C)
 
 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

@@ -1,20 +1,22 @@
 using Test
-using CUDA: has_cuda_gpu, @allowscalar
+#using CUDA: has_cuda_gpu, @allowscalar
 
 using Oceananigans
+using Oceananigans: fill_halo_regions!
 using Oceananigans.AbstractOperations: AbstractOperation
 using Oceananigans.Fields: @compute
 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()
+arch = CPU()#has_cuda_gpu() ? GPU() : CPU()
 
 N = 6
 regular_grid = RectilinearGrid(arch, size=(N, N, N), extent=(1, 1, 1))
@@ -172,8 +174,49 @@ function test_tracer_diagnostics(model)
     return nothing
 end
 
-#---
+function test_mixed_layer_depth(grid; zₘₓₗ = 0.5)
+    # buoyancy criterion (default)
+    buoyancy = BuoyancyTracer()
+    δb = -1
+    ∂z_b = - δb / zₘₓₗ
+
+    boundary_conditions = FieldBoundaryConditions(grid, (Center, Center, Center); top = GradientBoundaryCondition(∂z_b))
+
+    b = CenterField(grid; boundary_conditions)
+    mld = MixedLayerDepth(grid, buoyancy, (; b); criterion = BuoyancyAnomalyCriterion(δb))
+
+    @test isinf(mld[1, 1])
+
+    set!(b, (x, y, z) -> z * ∂z_b)
+    fill_halo_regions!(b)
+
+    @test mld[1, 1] ≈ -zₘₓₗ
+
+    # density criterion
+    criterion = DensityAnomalyCriterion(convert(eltype(grid), 1/8))
+
+    buoyancy = SeawaterBuoyancy(equation_of_state=BoussinesqEquationOfState(LinearRoquetSeawaterPolynomial(), 1000),
+                                constant_salinity=0)
+
+    δρ = 0.125
+    ∂z_T = - δρ / zₘₓₗ / buoyancy.equation_of_state.seawater_polynomial.R₀₁₀
 
+    boundary_conditions = FieldBoundaryConditions(grid, (Center, Center, Center); top = GradientBoundaryCondition(∂z_T))
+
+    T = CenterField(grid; boundary_conditions)
+    mld = MixedLayerDepth(grid, buoyancy, (; T); criterion)
+
+    @test isinf(mld[1, 1])
+
+    set!(T, (x, y, z) -> z * ∂z_T + 10)
+    fill_halo_regions!(T)
+
+    @test mld[1, 1] ≈ -zₘₓₗ
+
+    return nothing
+end
+#---
+#=
 @testset "Tracer diagnostics tests" begin
     @info "  Testing tracer diagnostics"
     for (grid_class, grid) in zip(keys(grids), values(grids))
@@ -201,7 +244,10 @@ end
                     @info "            Testing buoyancy diagnostics"
                     test_buoyancy_diagnostics(model)
                 end
+                @info "            Testing mixed layer depth diagnostic"
+                test_mixed_layer_depth(grid)
             end
         end
     end
 end
+=#