From cfaf34b4fee251c51ec837c325bfea665e72c01a Mon Sep 17 00:00:00 2001
From: Jago Strong-Wright <jagoosw@protonmail.com>
Date: Mon, 11 Nov 2024 15:18:53 +0000
Subject: [PATCH 01/18] added mld calculaion

---
 src/FlowDiagnostics.jl | 51 ++++++++++++++++++++++++++++++++++++++++--
 src/Oceanostics.jl     |  1 +
 2 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/src/FlowDiagnostics.jl b/src/FlowDiagnostics.jl
index abeb4240..a3279745 100755
--- a/src/FlowDiagnostics.jl
+++ b/src/FlowDiagnostics.jl
@@ -4,14 +4,18 @@ using DocStringExtensions
 export RichardsonNumber, RossbyNumber
 export ErtelPotentialVorticity, ThermalWindPotentialVorticity, DirectionalErtelPotentialVorticity
 export StrainRateTensorModulus, VorticityTensorModulus, Q, QVelocityGradientTensorInvariant
+export DensityCriteriaMixedLayerDepth
 
 using Oceanostics: validate_location, validate_dissipative_closure, add_background_fields
 
 using Oceananigans: NonhydrostaticModel, FPlane, ConstantCartesianCoriolis, BuoyancyField, BuoyancyTracer
+using Oceananigans.BuoyancyModels: get_temperature_and_salinity, SeawaterBuoyancy
 using Oceananigans.Operators
 using Oceananigans.AbstractOperations
 using Oceananigans.AbstractOperations: KernelFunctionOperation
-using Oceananigans.Grids: Center, Face, NegativeZDirection, ZDirection
+using Oceananigans.Grids: Center, Face, NegativeZDirection, ZDirection, znode
+
+using SeawaterPolynomials: ρ′
 
 #+++ Richardson number
 @inline ψ²(i, j, k, grid, ψ) = @inbounds ψ[i, j, k]^2
@@ -502,6 +506,49 @@ function QVelocityGradientTensorInvariant(model; location = (Center, Center, Cen
 end
 
 const Q = QVelocityGradientTensorInvariant
-#---
 
+"""
+    $(SIGNATURES)
+
+Calculate the mixed layer depth defined as the depth at which the density is 
+some threshold (defaults to 0.125kg/m³) higher than the surface density.
+
+`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 DensityCriteriaMixedLayerDepth(grid, buoyancy, C; pertubation = convert(eltype(grid), 1/8))
+    validate_buoyancy_model(buoyancy)
+    return KernelFunctionOperation{Center, Center, Nothing}(_density_criteria_mixed_layer_depth!, grid, buoyancy, C, pertubation)
+end
+
+validate_buoyancy_model(buoyancy) = @error "Mixed layer depth is only implemented for `SeawaterBuoyancy`"
+validate_buoyancy_model(::SeawaterBuoyancy) = nothing
+
+@inline function density(i, j, k, grid, b::SeawaterBuoyancy, C)
+    T, S = get_temperature_and_salinity(b, C)
+
+    return ρ′(i, j, k, grid, b.equation_of_state, T, S)
+end
+
+function _density_criteria_mixed_layer_depth!(i, j, k, grid, buoyancy, C, δρ)
+    ρᵣ = density(i, j, grid.Nz, grid, buoyancy, C)
+
+    kₘₗ = -1
+
+    for k in grid.Nz-1:-1:1
+        ρₖ = density(i, j, k, grid, buoyancy, C)
+
+        kₘₗ = ifelse((ρₖ > ρᵣ + δρ) & (kₘₗ < 0), k, kₘₗ)
+    end
+
+    ρₖ = density(i, j, kₘₗ, grid, buoyancy, C)
+    ρ₊ = density(i, j, kₘₗ+1, grid, buoyancy, C)
+
+    zₖ = znode(i, j, kₘₗ, grid, Center(), Center(), Center())
+    z₊ = znode(i, j, kₘₗ+1, grid, Center(), Center(), Center())
+
+    zₘₗ = zₖ + (z₊ - zₖ) * (ρᵣ + δρ - ρₖ) / (ρ₊ - ρₖ)
+    
+    return ifelse(kₘₗ == -1, -Inf, zₘₗ)
+end
 end # module
diff --git a/src/Oceanostics.jl b/src/Oceanostics.jl
index a01157ac..35f04093 100755
--- a/src/Oceanostics.jl
+++ b/src/Oceanostics.jl
@@ -20,6 +20,7 @@ export RichardsonNumber, RossbyNumber
 export ErtelPotentialVorticity, ThermalWindPotentialVorticity
 export DirectionalErtelPotentialVorticity
 export StrainRateTensorModulus, VorticityTensorModulus, Q, QVelocityGradientTensorInvariant
+export DensityCriteriaMixedLayerDepth
 #---
 
 #+++ PotentialEnergyEquationTerms exports

From 71653f6148d59aca4c9c351c975bdefb893ce037 Mon Sep 17 00:00:00 2001
From: Jago Strong-Wright <jagoosw@protonmail.com>
Date: Mon, 11 Nov 2024 17:28:04 +0000
Subject: [PATCH 02/18] added tests and fixed

---
 src/FlowDiagnostics.jl | 12 ++++++------
 test/runtests.jl       | 29 +++++++++++++++++++++++++++--
 2 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/src/FlowDiagnostics.jl b/src/FlowDiagnostics.jl
index a3279745..eda4a5d2 100755
--- a/src/FlowDiagnostics.jl
+++ b/src/FlowDiagnostics.jl
@@ -15,7 +15,8 @@ using Oceananigans.AbstractOperations
 using Oceananigans.AbstractOperations: KernelFunctionOperation
 using Oceananigans.Grids: Center, Face, NegativeZDirection, ZDirection, znode
 
-using SeawaterPolynomials: ρ′
+using SeawaterPolynomials: ρ′, BoussinesqEquationOfState
+using SeawaterPolynomials.SecondOrderSeawaterPolynomials: SecondOrderSeawaterPolynomial
 
 #+++ Richardson number
 @inline ψ²(i, j, k, grid, ψ) = @inbounds ψ[i, j, k]^2
@@ -521,20 +522,19 @@ function DensityCriteriaMixedLayerDepth(grid, buoyancy, C; pertubation = convert
     return KernelFunctionOperation{Center, Center, Nothing}(_density_criteria_mixed_layer_depth!, grid, buoyancy, C, pertubation)
 end
 
-validate_buoyancy_model(buoyancy) = @error "Mixed layer depth is only implemented for `SeawaterBuoyancy`"
-validate_buoyancy_model(::SeawaterBuoyancy) = nothing
+validate_buoyancy_model(buoyancy) = @error "Mixed layer depth is not implemented for your buoyancy model as the density anomaly is not defined"
+validate_buoyancy_model(::SeawaterBuoyancy{<:Any, <:BoussinesqEquationOfState}) = nothing
 
-@inline function density(i, j, k, grid, b::SeawaterBuoyancy, C)
+@inline function density(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
 
 function _density_criteria_mixed_layer_depth!(i, j, k, grid, buoyancy, C, δρ)
-    ρᵣ = density(i, j, grid.Nz, grid, buoyancy, C)
+    ρᵣ = (density(i, j, grid.Nz, grid, buoyancy, C) + density(i, j, grid.Nz+1, grid, buoyancy, C)) * convert(eltype(grid), 0.5)
 
     kₘₗ = -1
-
     for k in grid.Nz-1:-1:1
         ρₖ = density(i, j, k, grid, buoyancy, C)
 
diff --git a/test/runtests.jl b/test/runtests.jl
index 14d55e7b..cc5c22b2 100755
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -2,12 +2,14 @@ using Test
 using CUDA
 
 using Oceananigans
+using Oceananigans: fill_halo_regions!
 using Oceananigans.AbstractOperations: AbstractOperation
 using Oceananigans.BuoyancyModels: buoyancy_perturbationᶜᶜᶜ
 using Oceananigans.Fields: @compute
 using Oceananigans.TurbulenceClosures: ThreeDimensionalFormulation
 using Oceananigans.Models: seawater_density, model_geopotential_height
 using SeawaterPolynomials: RoquetEquationOfState, TEOS10EquationOfState
+using SeawaterPolynomials.SecondOrderSeawaterPolynomials: LinearRoquetSeawaterPolynomial
 
 using Oceanostics
 using Oceanostics: TKEBudgetTerms, TracerVarianceBudgetTerms, FlowDiagnostics, PressureRedistributionTerm, BuoyancyProductionTerm, AdvectionTerm
@@ -535,6 +537,29 @@ function test_auxiliary_functions(model)
     @test all(Array(interior(fields_without_means.v)) .== 0)
     return
 end
+
+function test_mixed_layer_depth(grid; 
+                                buoyancy = SeawaterBuoyancy(equation_of_state=BoussinesqEquationOfState(LinearRoquetSeawaterPolynomial(), 1000),
+                                                            constant_salinity=0),
+                                zₘₓₗ = 0.5,
+                                δρ = 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 = DensityCriteriaMixedLayerDepth(grid, buoyancy, (; T); pertubation = δρ)
+
+    @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
+
 #---
 
 model_kwargs = (buoyancy = Buoyancy(model=BuoyancyTracer()),
@@ -618,9 +643,9 @@ model_types = (NonhydrostaticModel, HydrostaticFreeSurfaceModel)
                 end
 
             end
-            test_PEbuoyancytracer_equals_PElineareos(grid)
-
         end
+        test_PEbuoyancytracer_equals_PElineareos(grid)
+        test_mixed_layer_depth(grid)
 
         @info "Testing input validation for dissipation rates"
         invalid_closures = [HorizontalScalarDiffusivity(ν=1e-6, κ=1e-7),

From b88191cee6872024e6983756d0b9c305d4862a2b Mon Sep 17 00:00:00 2001
From: Jago Strong-Wright <jagoosw@protonmail.com>
Date: Mon, 11 Nov 2024 20:51:42 +0000
Subject: [PATCH 03/18] generalised mld computation

---
 src/FlowDiagnostics.jl | 93 +++++++++++++++++++++++++++++++-----------
 src/Oceanostics.jl     |  2 +-
 test/runtests.jl       |  4 +-
 3 files changed, 73 insertions(+), 26 deletions(-)

diff --git a/src/FlowDiagnostics.jl b/src/FlowDiagnostics.jl
index eda4a5d2..1b2efbd4 100755
--- a/src/FlowDiagnostics.jl
+++ b/src/FlowDiagnostics.jl
@@ -4,7 +4,7 @@ using DocStringExtensions
 export RichardsonNumber, RossbyNumber
 export ErtelPotentialVorticity, ThermalWindPotentialVorticity, DirectionalErtelPotentialVorticity
 export StrainRateTensorModulus, VorticityTensorModulus, Q, QVelocityGradientTensorInvariant
-export DensityCriteriaMixedLayerDepth
+export MixedLayerDepth, DensityAnomalyCriterion
 
 using Oceanostics: validate_location, validate_dissipative_closure, add_background_fields
 
@@ -507,23 +507,67 @@ function QVelocityGradientTensorInvariant(model; location = (Center, Center, Cen
 end
 
 const Q = QVelocityGradientTensorInvariant
+struct MixedLayerDepth{C}
+    criterion::C
+end
 
 """
     $(SIGNATURES)
 
-Calculate the mixed layer depth defined as the depth at which the density is 
-some threshold (defaults to 0.125kg/m³) higher than the surface density.
+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.
 
-`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.
+When `DensityAnomalyCriterion` is used, the arguments `buoyancy` and `C` should be 
+supploed 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 DensityCriteriaMixedLayerDepth(grid, buoyancy, C; pertubation = convert(eltype(grid), 1/8))
-    validate_buoyancy_model(buoyancy)
-    return KernelFunctionOperation{Center, Center, Nothing}(_density_criteria_mixed_layer_depth!, grid, buoyancy, C, pertubation)
+function MixedLayerDepth(grid, args...; criterion = DensityAnomalyCriterion(convert(eltype(grid), 1/8)))
+    validate_criterion_model(criterion, args...)
+
+    MLD = MixedLayerDepth(criterion)
+
+    return KernelFunctionOperation{Center, Center, Nothing}(MLD, grid, args...)
 end
 
-validate_buoyancy_model(buoyancy) = @error "Mixed layer depth is not implemented for your buoyancy model as the density anomaly is not defined"
-validate_buoyancy_model(::SeawaterBuoyancy{<:Any, <:BoussinesqEquationOfState}) = nothing
+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)
+
+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 
+supploed 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}
+    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) = 
+    @error "DensityAnomalyCriterion is not implemented for your buoyancy, model as the density anomaly is not defined"
+    
+validate_criterion_model(::DensityAnomalyCriterion, ::SeawaterBuoyancy{<:Any, <:BoussinesqEquationOfState}, C::NamedTuple) = nothing
 
 @inline function density(i, j, k, grid, b::SeawaterBuoyancy{<:Any, <:BoussinesqEquationOfState}, C)
     T, S = get_temperature_and_salinity(b, C)
@@ -531,24 +575,27 @@ validate_buoyancy_model(::SeawaterBuoyancy{<:Any, <:BoussinesqEquationOfState})
     return ρ′(i, j, k, grid, b.equation_of_state, T, S)
 end
 
-function _density_criteria_mixed_layer_depth!(i, j, k, grid, buoyancy, C, δρ)
+@inline function (density_criterion::DensityAnomalyCriterion)(i, j, k, grid, buoyancy, C)
+    δρ = density_criterion.pertubation
+
     ρᵣ = (density(i, j, grid.Nz, grid, buoyancy, C) + density(i, j, grid.Nz+1, grid, buoyancy, C)) * convert(eltype(grid), 0.5)
 
-    kₘₗ = -1
-    for k in grid.Nz-1:-1:1
-        ρₖ = density(i, j, k, grid, buoyancy, C)
+    ρₖ = density(i, j, k, grid, buoyancy, C)
 
-        kₘₗ = ifelse((ρₖ > ρᵣ + δρ) & (kₘₗ < 0), k, kₘₗ)
-    end
+    return ρₖ > ρᵣ + δρ
+end
+
+@inline function interpolate_from_nearest_cell(density_criterion::DensityAnomalyCriterion, i, j, k, grid, buoyancy, C)
+    δρ = density_criterion.pertubation
 
-    ρₖ = density(i, j, kₘₗ, grid, buoyancy, C)
-    ρ₊ = density(i, j, kₘₗ+1, grid, buoyancy, C)
+    ρᵣ = (density(i, j, grid.Nz, grid, buoyancy, C) + density(i, j, grid.Nz+1, grid, buoyancy, C)) * convert(eltype(grid), 0.5)
 
-    zₖ = znode(i, j, kₘₗ, grid, Center(), Center(), Center())
-    z₊ = znode(i, j, kₘₗ+1, grid, Center(), Center(), Center())
+    ρₖ = density(i, j, k, grid, buoyancy, C)
+    ρ₊ = density(i, j, k+1, grid, buoyancy, C)
 
-    zₘₗ = zₖ + (z₊ - zₖ) * (ρᵣ + δρ - ρₖ) / (ρ₊ - ρₖ)
-    
-    return ifelse(kₘₗ == -1, -Inf, zₘₗ)
+    zₖ = znode(i, j, k, grid, Center(), Center(), Center())
+    z₊ = znode(i, j, k+1, grid, Center(), Center(), Center())
+
+    return zₖ + (z₊ - zₖ) * (ρᵣ + δρ - ρₖ) / (ρ₊ - ρₖ)
 end
 end # module
diff --git a/src/Oceanostics.jl b/src/Oceanostics.jl
index 35f04093..5d42aab7 100755
--- a/src/Oceanostics.jl
+++ b/src/Oceanostics.jl
@@ -20,7 +20,7 @@ export RichardsonNumber, RossbyNumber
 export ErtelPotentialVorticity, ThermalWindPotentialVorticity
 export DirectionalErtelPotentialVorticity
 export StrainRateTensorModulus, VorticityTensorModulus, Q, QVelocityGradientTensorInvariant
-export DensityCriteriaMixedLayerDepth
+export MixedLayerDepth, DensityAnomalyCriterion
 #---
 
 #+++ PotentialEnergyEquationTerms exports
diff --git a/test/runtests.jl b/test/runtests.jl
index cc5c22b2..a50791b5 100755
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -8,7 +8,7 @@ using Oceananigans.BuoyancyModels: buoyancy_perturbationᶜᶜᶜ
 using Oceananigans.Fields: @compute
 using Oceananigans.TurbulenceClosures: ThreeDimensionalFormulation
 using Oceananigans.Models: seawater_density, model_geopotential_height
-using SeawaterPolynomials: RoquetEquationOfState, TEOS10EquationOfState
+using SeawaterPolynomials: RoquetEquationOfState, TEOS10EquationOfState, BoussinesqEquationOfState
 using SeawaterPolynomials.SecondOrderSeawaterPolynomials: LinearRoquetSeawaterPolynomial
 
 using Oceanostics
@@ -548,7 +548,7 @@ function test_mixed_layer_depth(grid;
 
     T = CenterField(grid; boundary_conditions)
 
-    mld = DensityCriteriaMixedLayerDepth(grid, buoyancy, (; T); pertubation = δρ)
+    mld = MixedLayerDepth(grid, buoyancy, (; T))
 
     @test isinf(mld[1, 1])
     

From 623995018155ac5d291ca33e50230f18a6267744 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1s=20Chor?= <tomaschor@gmail.com>
Date: Mon, 31 Mar 2025 12:34:27 -0700
Subject: [PATCH 04/18] Update test/runtests.jl

---
 test/runtests.jl | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index 6a090841..9b2dd9f9 100755
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -45,10 +45,7 @@ end
         include("test_canonical_flows.jl")
     end
 
-    test_PEbuoyancytracer_equals_PElineareos(grid)
     test_mixed_layer_depth(grid)
-
-
     if group == :progress_messengers || group == :all
         include("test_progress_messengers.jl")
     end

From 6d3e07b43dc7937a7eb733b77cd46b9783139cc6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1s=20Chor?= <tomaschor@gmail.com>
Date: Mon, 31 Mar 2025 17:02:25 -0700
Subject: [PATCH 05/18] Update src/FlowDiagnostics.jl

---
 src/FlowDiagnostics.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/FlowDiagnostics.jl b/src/FlowDiagnostics.jl
index b41af2ab..98137fd4 100755
--- a/src/FlowDiagnostics.jl
+++ b/src/FlowDiagnostics.jl
@@ -12,7 +12,7 @@ using Oceanostics: validate_location,
                    get_coriolis_frequency_components
 
 using Oceananigans: NonhydrostaticModel, FPlane, ConstantCartesianCoriolis, BuoyancyField, BuoyancyTracer
-using Oceananigans.BuoyancyModels: get_temperature_and_salinity, SeawaterBuoyancy
+using Oceananigans.BuoyancyFormulations: get_temperature_and_salinity, SeawaterBuoyancy
 using Oceananigans.Operators
 using Oceananigans.AbstractOperations
 using Oceananigans.AbstractOperations: KernelFunctionOperation

From 15012c1c3e1c7efd14d39c8ee5c95c5dcce00420 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1s=20Chor?= <tomaschor@gmail.com>
Date: Tue, 1 Apr 2025 18:40:19 -0700
Subject: [PATCH 06/18] Update test/runtests.jl

---
 test/runtests.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index 9b2dd9f9..7efc2615 100755
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -2,7 +2,7 @@ using Test
 
 group     = get(ENV, "TEST_GROUP", :all) |> Symbol
 
-function test_mixed_layer_depth(grid; 
+function test_mixed_layer_depth(grid;
                                 buoyancy = SeawaterBuoyancy(equation_of_state=BoussinesqEquationOfState(LinearRoquetSeawaterPolynomial(), 1000),
                                                             constant_salinity=0),
                                 zₘₓₗ = 0.5,

From 8018aad5eee8201e8d703ef708486bd3a191c6a4 Mon Sep 17 00:00:00 2001
From: tomaschor <tomaschor@gmail.com>
Date: Tue, 1 Apr 2025 19:50:16 -0700
Subject: [PATCH 07/18] move test to proper file

---
 test/runtests.jl                | 23 -----------------------
 test/test_tracer_diagnostics.jl | 28 ++++++++++++++++++++++++++--
 2 files changed, 26 insertions(+), 25 deletions(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index 7efc2615..1c305ccc 100755
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -2,28 +2,6 @@ using Test
 
 group     = get(ENV, "TEST_GROUP", :all) |> Symbol
 
-function test_mixed_layer_depth(grid;
-                                buoyancy = SeawaterBuoyancy(equation_of_state=BoussinesqEquationOfState(LinearRoquetSeawaterPolynomial(), 1000),
-                                                            constant_salinity=0),
-                                zₘₓₗ = 0.5,
-                                δρ = 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))
-
-    @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 "Oceanostics" begin
     if group == :vel_diagnostics || group == :all
         include("test_velocity_diagnostics.jl")
@@ -45,7 +23,6 @@ end
         include("test_canonical_flows.jl")
     end
 
-    test_mixed_layer_depth(grid)
     if group == :progress_messengers || group == :all
         include("test_progress_messengers.jl")
     end
diff --git a/test/test_tracer_diagnostics.jl b/test/test_tracer_diagnostics.jl
index 420b679d..f043e879 100644
--- a/test/test_tracer_diagnostics.jl
+++ b/test/test_tracer_diagnostics.jl
@@ -2,16 +2,18 @@ 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.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 +174,26 @@ function test_tracer_diagnostics(model)
     return nothing
 end
 
+function test_mixed_layer_depth(grid;
+                                buoyancy = SeawaterBuoyancy(equation_of_state=BoussinesqEquationOfState(LinearRoquetSeawaterPolynomial(), 1000),
+                                                            constant_salinity=0),
+                                zₘₓₗ = 0.5,
+                                δρ = 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))
+
+    @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
@@ -201,6 +223,8 @@ end
                     @info "            Testing buoyancy diagnostics"
                     test_buoyancy_diagnostics(model)
                 end
+                @info "            Testing mixed layer depth diagnostic"
+                test_mixed_layer_depth(grid)
             end
         end
     end

From 766458db1bb6201901613d9d1b9b6be435a37b41 Mon Sep 17 00:00:00 2001
From: Jago Strong-Wright <jagoosw@protonmail.com>
Date: Wed, 2 Apr 2025 15:21:31 -0400
Subject: [PATCH 08/18] Added buoyancy criterion and fixed tests

---
 src/FlowDiagnostics.jl          | 85 +++++++++++++++++++++++----------
 src/Oceanostics.jl              |  2 +-
 test/test_tracer_diagnostics.jl | 44 ++++++++++++-----
 3 files changed, 93 insertions(+), 38 deletions(-)

diff --git a/src/FlowDiagnostics.jl b/src/FlowDiagnostics.jl
index 98137fd4..506dcb96 100755
--- a/src/FlowDiagnostics.jl
+++ b/src/FlowDiagnostics.jl
@@ -4,7 +4,7 @@ using DocStringExtensions
 export RichardsonNumber, RossbyNumber
 export ErtelPotentialVorticity, ThermalWindPotentialVorticity, DirectionalErtelPotentialVorticity
 export StrainRateTensorModulus, VorticityTensorModulus, Q, QVelocityGradientTensorInvariant
-export MixedLayerDepth, DensityAnomalyCriterion
+export MixedLayerDepth, BuoyancyAnomalyCriterion, DensityAnomalyCriterion
 
 using Oceanostics: validate_location,
                    validate_dissipative_closure,
@@ -12,14 +12,13 @@ using Oceanostics: validate_location,
                    get_coriolis_frequency_components
 
 using Oceananigans: NonhydrostaticModel, FPlane, ConstantCartesianCoriolis, BuoyancyField, BuoyancyTracer
-using Oceananigans.BuoyancyFormulations: get_temperature_and_salinity, SeawaterBuoyancy
+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, znode
 
-
 using SeawaterPolynomials: ρ′, BoussinesqEquationOfState
 using SeawaterPolynomials.SecondOrderSeawaterPolynomials: SecondOrderSeawaterPolynomial
 
@@ -447,6 +446,7 @@ function QVelocityGradientTensorInvariant(model; location = (Center, Center, Cen
 end
 
 const Q = QVelocityGradientTensorInvariant
+
 struct MixedLayerDepth{C}
     criterion::C
 end
@@ -460,11 +460,11 @@ Defaults to `DensityAnomalyCriterion` where the depth is that at which the densi
 is some threshold (defaults to 0.125kg/m³) higher than the surface density.
 
 When `DensityAnomalyCriterion` is used, the arguments `buoyancy` and `C` should be 
-supploed where `buoyancy` should be the buoyancy model, and `C` should be a named 
+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 = DensityAnomalyCriterion(convert(eltype(grid), 1/8)))
+function MixedLayerDepth(grid, args...; criterion = BuoyancyAnomalyCriterion(convert(eltype(grid), -1/8 * 1020 / g_Earth)))
     validate_criterion_model(criterion, args...)
 
     MLD = MixedLayerDepth(criterion)
@@ -486,6 +486,41 @@ function (MLD::MixedLayerDepth)(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)
 
@@ -493,49 +528,47 @@ Defines the mixed layer to be the depth at which the density is more than `pertu
 greater than the surface density.
 
 When this model is used, the arguments `buoyancy` and `C` should be 
-supploed where `buoyancy` should be the buoyancy model, and `C` should be a named 
+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}
+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) = 
-    @error "DensityAnomalyCriterion is not implemented for your buoyancy, model as the density anomaly is not defined"
+validate_criterion_model(::DensityAnomalyCriterion, buoyancy, C) = nothing
     
 validate_criterion_model(::DensityAnomalyCriterion, ::SeawaterBuoyancy{<:Any, <:BoussinesqEquationOfState}, C::NamedTuple) = nothing
 
-@inline function density(i, j, k, grid, b::SeawaterBuoyancy{<:Any, <:BoussinesqEquationOfState}, C)
+@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
 
-@inline function (density_criterion::DensityAnomalyCriterion)(i, j, k, grid, buoyancy, C)
-    δρ = density_criterion.pertubation
-
-    ρᵣ = (density(i, j, grid.Nz, grid, buoyancy, C) + density(i, j, grid.Nz+1, grid, buoyancy, C)) * convert(eltype(grid), 0.5)
+"""
+    $(SIGNATURES)
 
-    ρₖ = density(i, j, k, grid, buoyancy, C)
+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).
 
-    return ρₖ > ρᵣ + δρ
+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
 
-@inline function interpolate_from_nearest_cell(density_criterion::DensityAnomalyCriterion, i, j, k, grid, buoyancy, C)
-    δρ = density_criterion.pertubation
-
-    ρᵣ = (density(i, j, grid.Nz, grid, buoyancy, C) + density(i, j, grid.Nz+1, grid, buoyancy, C)) * convert(eltype(grid), 0.5)
+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)"
 
-    ρₖ = density(i, j, k, grid, buoyancy, C)
-    ρ₊ = density(i, j, k+1, grid, buoyancy, C)
+validate_criterion_model(::BuoyancyAnomalyCriterion, buoyancy, C) = nothing
 
-    zₖ = znode(i, j, k, grid, Center(), Center(), Center())
-    z₊ = znode(i, j, k+1, grid, Center(), Center(), Center())
+@inline anomaly(::BuoyancyAnomalyCriterion, i, j, k, grid, buoyancy, C) = buoyancy_perturbationᶜᶜᶜ(i, j, k, grid, buoyancy, C)
 
-    return zₖ + (z₊ - zₖ) * (ρᵣ + δρ - ρₖ) / (ρ₊ - ρₖ)
-end
 end # module
diff --git a/src/Oceanostics.jl b/src/Oceanostics.jl
index c95e05fa..6a5b2837 100755
--- a/src/Oceanostics.jl
+++ b/src/Oceanostics.jl
@@ -20,7 +20,7 @@ export RichardsonNumber, RossbyNumber
 export ErtelPotentialVorticity, ThermalWindPotentialVorticity
 export DirectionalErtelPotentialVorticity
 export StrainRateTensorModulus, VorticityTensorModulus, Q, QVelocityGradientTensorInvariant
-export MixedLayerDepth, DensityAnomalyCriterion
+export MixedLayerDepth, BuoyancyAnomalyCriterion, DensityAnomalyCriterion
 #---
 
 #+++ PotentialEnergyEquationTerms exports
diff --git a/test/test_tracer_diagnostics.jl b/test/test_tracer_diagnostics.jl
index f043e879..bd957c60 100644
--- a/test/test_tracer_diagnostics.jl
+++ b/test/test_tracer_diagnostics.jl
@@ -1,5 +1,5 @@
 using Test
-using CUDA: has_cuda_gpu, @allowscalar
+#using CUDA: has_cuda_gpu, @allowscalar
 
 using Oceananigans
 using Oceananigans: fill_halo_regions!
@@ -16,7 +16,7 @@ using Oceanostics: get_coriolis_frequency_components
 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))
@@ -174,20 +174,41 @@ function test_tracer_diagnostics(model)
     return nothing
 end
 
-function test_mixed_layer_depth(grid;
-                                buoyancy = SeawaterBuoyancy(equation_of_state=BoussinesqEquationOfState(LinearRoquetSeawaterPolynomial(), 1000),
-                                                            constant_salinity=0),
-                                zₘₓₗ = 0.5,
-                                δρ = 0.125,
-                                ∂z_T = - δρ / zₘₓₗ / buoyancy.equation_of_state.seawater_polynomial.R₀₁₀)
+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))
+    mld = MixedLayerDepth(grid, buoyancy, (; T); criterion)
 
     @test isinf(mld[1, 1])
 
-    set!(T, (x, y, z) -> z * ∂z_T+10)
+    set!(T, (x, y, z) -> z * ∂z_T + 10)
     fill_halo_regions!(T)
 
     @test mld[1, 1] ≈ -zₘₓₗ
@@ -195,7 +216,7 @@ function test_mixed_layer_depth(grid;
     return nothing
 end
 #---
-
+#=
 @testset "Tracer diagnostics tests" begin
     @info "  Testing tracer diagnostics"
     for (grid_class, grid) in zip(keys(grids), values(grids))
@@ -229,3 +250,4 @@ end
         end
     end
 end
+=#
\ No newline at end of file

From 21952d62bd3eb92ceaca8d807a60e7dcd9098530 Mon Sep 17 00:00:00 2001
From: Jago Strong-Wright <jagoosw@protonmail.com>
Date: Wed, 2 Apr 2025 15:54:55 -0400
Subject: [PATCH 09/18] correct buoyancy default

---
 src/FlowDiagnostics.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/FlowDiagnostics.jl b/src/FlowDiagnostics.jl
index 506dcb96..3917cc10 100755
--- a/src/FlowDiagnostics.jl
+++ b/src/FlowDiagnostics.jl
@@ -464,7 +464,7 @@ supplied where `buoyancy` should be the buoyancy model, and `C` should be a name
 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)))
+function MixedLayerDepth(grid, args...; criterion = BuoyancyAnomalyCriterion(convert(eltype(grid), - 0.125 / 1020 * g_Earth)))
     validate_criterion_model(criterion, args...)
 
     MLD = MixedLayerDepth(criterion)
@@ -561,7 +561,7 @@ 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
+    pertubation :: FT = - 0.125 / 1020 * g_Earth
 end
 
 validate_criterion_model(::BuoyancyAnomalyCriterion, args...) = 

From dc443744eb52599ec839d46dbcc86f483694eedf Mon Sep 17 00:00:00 2001
From: Jago Strong-Wright <jagoosw@protonmail.com>
Date: Wed, 2 Apr 2025 16:09:46 -0400
Subject: [PATCH 10/18] updated

---
 src/FlowDiagnostics.jl          | 60 +++++++++++++++++++--------------
 test/test_tracer_diagnostics.jl |  8 ++---
 2 files changed, 38 insertions(+), 30 deletions(-)

diff --git a/src/FlowDiagnostics.jl b/src/FlowDiagnostics.jl
index 3917cc10..5a64e609 100755
--- a/src/FlowDiagnostics.jl
+++ b/src/FlowDiagnostics.jl
@@ -464,7 +464,7 @@ supplied where `buoyancy` should be the buoyancy model, and `C` should be a name
 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), - 0.125 / 1020 * g_Earth)))
+function MixedLayerDepth(grid, args...; criterion = BuoyancyAnomalyCriterion(convert(eltype(grid), -1e-4 * g_Earth)))
     validate_criterion_model(criterion, args...)
 
     MLD = MixedLayerDepth(criterion)
@@ -490,15 +490,15 @@ 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`.
+`anomaly` + `threshold` 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`.
+`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.pertubation
+    δ = 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)
 
@@ -508,7 +508,7 @@ abstract type AbstractAnomalyCriterion end
 end
 
 @inline function interpolate_from_nearest_cell(criterion::AbstractAnomalyCriterion, i, j, k, grid, args...)
-    δ = criterion.pertubation
+    δ = 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)
 
@@ -524,51 +524,61 @@ end
 """
     $(SIGNATURES)
 
-Defines the mixed layer to be the depth at which the density is more than `pertubation`
-greater than the surface density.
+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` 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
+@kwdef struct BuoyancyAnomalyCriterion{FT} <: AbstractAnomalyCriterion
+    threshold :: FT = -1e-4 * g_Earth
 end
 
-validate_criterion_model(::DensityAnomalyCriterion, args...) = 
+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(::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)
+validate_criterion_model(::BuoyancyAnomalyCriterion, buoyancy, C) = nothing
 
-    return ρ′(i, j, k, grid, b.equation_of_state, T, S)
-end
+@inline anomaly(::BuoyancyAnomalyCriterion, i, j, k, grid, buoyancy, C) = buoyancy_perturbationᶜᶜᶜ(i, j, k, grid, buoyancy, C)
 
 """
     $(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).
+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` 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 = - 0.125 / 1020 * g_Earth
+@kwdef struct DensityAnomalyCriterion{FT} <: AbstractAnomalyCriterion
+             reference_density :: FT = 1020.0
+    gravitational_acceleration :: FT = g_Earth
+                     threshold :: FT = 0.125
 end
 
-validate_criterion_model(::BuoyancyAnomalyCriterion, args...) = 
+function DensityAnomalyCriterion(buoyancy::SeawaterBuoyancy{<:Any, <:BoussinesqEquationOfState}; threshold = 0.125)
+    ρᵣ = buoyancy.equation_of_state.reference_density
+    g  = buoyancy.gravitational_acceleration
+
+    return DensityAnomalyCriterion(ρᵣ, g, threshold)
+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(::BuoyancyAnomalyCriterion, buoyancy, C) = nothing
+validate_criterion_model(::DensityAnomalyCriterion, buoyancy, C) = nothing
+    
+@inline function anomaly(criterion::DensityAnomalyCriterion, i, j, k, grid, buoyancy, C)
+    ρᵣ = criterion.reference_density
+    g  = criterion.gravitational_acceleration
 
-@inline anomaly(::BuoyancyAnomalyCriterion, i, j, k, grid, buoyancy, C) = buoyancy_perturbationᶜᶜᶜ(i, j, k, grid, buoyancy, C)
+    b = buoyancy_perturbationᶜᶜᶜ(i, j, k, grid, buoyancy, C)
+
+    return - ρᵣ * b / g
+end
 
 end # module
diff --git a/test/test_tracer_diagnostics.jl b/test/test_tracer_diagnostics.jl
index bd957c60..d0b5283f 100644
--- a/test/test_tracer_diagnostics.jl
+++ b/test/test_tracer_diagnostics.jl
@@ -193,10 +193,9 @@ function test_mixed_layer_depth(grid; zₘₓₗ = 0.5)
     @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)
 
-    buoyancy = SeawaterBuoyancy(equation_of_state=BoussinesqEquationOfState(LinearRoquetSeawaterPolynomial(), 1000),
-                                constant_salinity=0)
+    criterion = DensityAnomalyCriterion(buoyancy; threshold = convert(eltype(grid), 1/8))
 
     δρ = 0.125
     ∂z_T = - δρ / zₘₓₗ / buoyancy.equation_of_state.seawater_polynomial.R₀₁₀
@@ -216,7 +215,7 @@ function test_mixed_layer_depth(grid; zₘₓₗ = 0.5)
     return nothing
 end
 #---
-#=
+
 @testset "Tracer diagnostics tests" begin
     @info "  Testing tracer diagnostics"
     for (grid_class, grid) in zip(keys(grids), values(grids))
@@ -250,4 +249,3 @@ end
         end
     end
 end
-=#
\ No newline at end of file

From 3d5303b779fee4ff1cc10183ce6484c86c1b0ec3 Mon Sep 17 00:00:00 2001
From: Jago Strong-Wright <jagoosw@protonmail.com>
Date: Thu, 3 Apr 2025 12:08:47 -0400
Subject: [PATCH 11/18] fixed tests

---
 test/test_tracer_diagnostics.jl | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/test/test_tracer_diagnostics.jl b/test/test_tracer_diagnostics.jl
index d0b5283f..5a0475e4 100644
--- a/test/test_tracer_diagnostics.jl
+++ b/test/test_tracer_diagnostics.jl
@@ -1,10 +1,11 @@
 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.Grids: znode
 using Oceananigans.TurbulenceClosures.Smagorinskys: LagrangianAveraging
 using Oceananigans.BuoyancyFormulations: buoyancy
 using SeawaterPolynomials: RoquetEquationOfState, TEOS10EquationOfState, BoussinesqEquationOfState
@@ -16,7 +17,7 @@ using Oceanostics: get_coriolis_frequency_components
 LinearBuoyancyForce = Union{BuoyancyTracer, SeawaterBuoyancy{<:Any, <:LinearEquationOfState}}
 
 #+++ Default grids
-arch = CPU()#has_cuda_gpu() ? GPU() : CPU()
+arch = has_cuda_gpu() ? GPU() : CPU()
 
 N = 6
 regular_grid = RectilinearGrid(arch, size=(N, N, N), extent=(1, 1, 1))
@@ -190,7 +191,7 @@ function test_mixed_layer_depth(grid; zₘₓₗ = 0.5)
     set!(b, (x, y, z) -> z * ∂z_b)
     fill_halo_regions!(b)
 
-    @test mld[1, 1] ≈ -zₘₓₗ
+    @test mld[1, 1] ≈ -zₘₓₗ + znode(1, 1, grid.Nz+1, grid, Center(), Center(), Face())
 
     # density criterion
     buoyancy = SeawaterBuoyancy(equation_of_state=BoussinesqEquationOfState(LinearRoquetSeawaterPolynomial(), 1000), constant_salinity=0)
@@ -210,7 +211,7 @@ function test_mixed_layer_depth(grid; zₘₓₗ = 0.5)
     set!(T, (x, y, z) -> z * ∂z_T + 10)
     fill_halo_regions!(T)
 
-    @test mld[1, 1] ≈ -zₘₓₗ
+    @test isapprox(mld[1, 1], -zₘₓₗ + znode(1, 1, grid.Nz+1, grid, Center(), Center(), Face()), atol=1e-10)
 
     return nothing
 end

From 6d9390be73c4045b948c237060b5c6703a357fbf Mon Sep 17 00:00:00 2001
From: Jago Strong-Wright <jagoosw@protonmail.com>
Date: Thu, 3 Apr 2025 12:52:50 -0400
Subject: [PATCH 12/18] test on all the buoyancy models

---
 test/test_tracer_diagnostics.jl | 67 +++++++++++++++++++--------------
 1 file changed, 38 insertions(+), 29 deletions(-)

diff --git a/test/test_tracer_diagnostics.jl b/test/test_tracer_diagnostics.jl
index 5a0475e4..6400073b 100644
--- a/test/test_tracer_diagnostics.jl
+++ b/test/test_tracer_diagnostics.jl
@@ -4,7 +4,7 @@ using CUDA: has_cuda_gpu, @allowscalar
 using Oceananigans
 using Oceananigans: fill_halo_regions!
 using Oceananigans.AbstractOperations: AbstractOperation
-using Oceananigans.Fields: @compute
+using Oceananigans.Fields: @compute, ConstantField
 using Oceananigans.Grids: znode
 using Oceananigans.TurbulenceClosures.Smagorinskys: LagrangianAveraging
 using Oceananigans.BuoyancyFormulations: buoyancy
@@ -175,45 +175,51 @@ 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))
+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)
 
-    b = CenterField(grid; boundary_conditions)
-    mld = MixedLayerDepth(grid, buoyancy, (; b); criterion = BuoyancyAnomalyCriterion(δb))
+    ∂z_b = - δb / zₘₓₗ
 
-    @test isinf(mld[1, 1])
+    if buoyancy isa BuoyancyTracer
+        boundary_conditions = FieldBoundaryConditions(grid, (Center, Center, Center); top = GradientBoundaryCondition(∂z_b))
 
-    set!(b, (x, y, z) -> z * ∂z_b)
-    fill_halo_regions!(b)
+        C = (; b = CenterField(grid; boundary_conditions))
+    else
+        g = buoyancy.gravitational_acceleration
 
-    @test mld[1, 1] ≈ -zₘₓₗ + znode(1, 1, grid.Nz+1, grid, Center(), Center(), Face())
+        ∂z_T = ∂z_b / (g * naive_thermal_expansion)
 
-    # density criterion
-    buoyancy = SeawaterBuoyancy(equation_of_state=BoussinesqEquationOfState(LinearRoquetSeawaterPolynomial(), 1000), constant_salinity=0)
+        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))
 
-    criterion = DensityAnomalyCriterion(buoyancy; threshold = convert(eltype(grid), 1/8))
+    if density_is_defined
+        ρᵣ = buoyancy.equation_of_state.reference_density
 
-    δρ = 0.125
-    ∂z_T = - δρ / zₘₓₗ / buoyancy.equation_of_state.seawater_polynomial.R₀₁₀
+        δρ = - δb * ρᵣ / g
 
-    boundary_conditions = FieldBoundaryConditions(grid, (Center, Center, Center); top = GradientBoundaryCondition(∂z_T))
+        criterion = DensityAnomalyCriterion(buoyancy; threshold = convert(eltype(grid), δρ))
 
-    T = CenterField(grid; boundary_conditions)
-    mld = MixedLayerDepth(grid, buoyancy, (; T); criterion)
+        mld_ρ = MixedLayerDepth(grid, buoyancy, C; criterion)
+    end 
 
-    @test isinf(mld[1, 1])
+    @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
 
-    set!(T, (x, y, z) -> z * ∂z_T + 10)
-    fill_halo_regions!(T)
+    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
 
-    @test isapprox(mld[1, 1], -zₘₓₗ + znode(1, 1, grid.Nz+1, grid, Center(), Center(), Face()), atol=1e-10)
+    fill_halo_regions!(C)
 
-    return nothing
+    @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
 #---
 
@@ -243,9 +249,12 @@ end
 
                     @info "            Testing buoyancy diagnostics"
                     test_buoyancy_diagnostics(model)
+
+                    if !isnothing(buoyancy)
+                        @info "            Testing mixed layer depth diagnostic"
+                        test_mixed_layer_depth(grid, buoyancy)
+                    end
                 end
-                @info "            Testing mixed layer depth diagnostic"
-                test_mixed_layer_depth(grid)
             end
         end
     end

From f9ac0e10f522b3c577102b17b99352944832ab4a Mon Sep 17 00:00:00 2001
From: tomaschor <tomaschor@gmail.com>
Date: Sat, 5 Apr 2025 22:35:26 -0300
Subject: [PATCH 13/18] remove signature where it does nothing

---
 src/FlowDiagnostics.jl | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/FlowDiagnostics.jl b/src/FlowDiagnostics.jl
index 5a64e609..45c3f6c9 100755
--- a/src/FlowDiagnostics.jl
+++ b/src/FlowDiagnostics.jl
@@ -487,12 +487,10 @@ function (MLD::MixedLayerDepth)(i, j, k, grid, args...)
 end
 
 """
-    $(SIGNATURES)
-
 An abstract mixed layer depth criterion where the mixed layer is defined to be
 `anomaly` + `threshold` greater than the surface value of `anomaly`.
 
-$(SIGNATURES) types should provide a method for the function `anomaly` in the form
+`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

From 53049f28600f032dcf762049ee097d8937468621 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1s=20Chor?= <tomaschor@gmail.com>
Date: Wed, 9 Apr 2025 15:16:11 -0700
Subject: [PATCH 14/18] Update src/FlowDiagnostics.jl

---
 src/FlowDiagnostics.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/FlowDiagnostics.jl b/src/FlowDiagnostics.jl
index 45c3f6c9..79475879 100755
--- a/src/FlowDiagnostics.jl
+++ b/src/FlowDiagnostics.jl
@@ -464,7 +464,7 @@ supplied where `buoyancy` should be the buoyancy model, and `C` should be a name
 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), -1e-4 * g_Earth)))
+function MixedLayerDepth(grid::AbstractGrid, args...; criterion = BuoyancyAnomalyCriterion(convert(eltype(grid), -1e-4 * g_Earth)))
     validate_criterion_model(criterion, args...)
 
     MLD = MixedLayerDepth(criterion)

From 0492f0fac4494eb5d8e4dd21f96a742f2f489e8d Mon Sep 17 00:00:00 2001
From: tomaschor <tomaschor@gmail.com>
Date: Wed, 9 Apr 2025 19:48:13 -0300
Subject: [PATCH 15/18] move MLD test one loop out

---
 src/FlowDiagnostics.jl          | 64 +++++++++++++++------------------
 test/test_tracer_diagnostics.jl |  8 ++---
 2 files changed, 32 insertions(+), 40 deletions(-)

diff --git a/src/FlowDiagnostics.jl b/src/FlowDiagnostics.jl
index 79475879..3b55ee61 100755
--- a/src/FlowDiagnostics.jl
+++ b/src/FlowDiagnostics.jl
@@ -17,7 +17,7 @@ using Oceananigans.Operators
 using Oceananigans.AbstractOperations
 using Oceananigans.AbstractOperations: KernelFunctionOperation
 using Oceananigans.BuoyancyFormulations: buoyancy
-using Oceananigans.Grids: Center, Face, NegativeZDirection, ZDirection, znode
+using Oceananigans.Grids: AbstractGrid, Center, Face, NegativeZDirection, ZDirection, znode
 
 using SeawaterPolynomials: ρ′, BoussinesqEquationOfState
 using SeawaterPolynomials.SecondOrderSeawaterPolynomials: SecondOrderSeawaterPolynomial
@@ -447,7 +447,7 @@ end
 
 const Q = QVelocityGradientTensorInvariant
 
-struct MixedLayerDepth{C}
+struct MixedLayerDepthKernel{C}
     criterion::C
 end
 
@@ -459,30 +459,26 @@ 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 
+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 = MixedLayerDepth(criterion)
-
+    MLD = MixedLayerDepthKernel(criterion)
     return KernelFunctionOperation{Center, Center, Nothing}(MLD, grid, args...)
 end
 
-function (MLD::MixedLayerDepth)(i, j, k, grid, args...)
+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
 
@@ -499,7 +495,6 @@ abstract type AbstractAnomalyCriterion end
     δ = 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 + δ
@@ -522,24 +517,22 @@ end
 """
     $(SIGNATURES)
 
-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).
+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` 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).
+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 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, 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, C) = nothing
+validate_criterion_model(::BuoyancyAnomalyCriterion, buoyancy_formulation, C) = nothing
 
-@inline anomaly(::BuoyancyAnomalyCriterion, i, j, k, grid, buoyancy, C) = buoyancy_perturbationᶜᶜᶜ(i, j, k, grid, buoyancy, C)
+@inline anomaly(::BuoyancyAnomalyCriterion, i, j, k, grid, buoyancy_formulation, C) = buoyancy_perturbationᶜᶜᶜ(i, j, k, grid, buoyancy_formulation, C)
 
 """
     $(SIGNATURES)
@@ -547,10 +540,10 @@ validate_criterion_model(::BuoyancyAnomalyCriterion, buoyancy, C) = nothing
 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` 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).
+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
@@ -558,24 +551,23 @@ specifies a constant salinity or temperature).
                      threshold :: FT = 0.125
 end
 
-function DensityAnomalyCriterion(buoyancy::SeawaterBuoyancy{<:Any, <:BoussinesqEquationOfState}; threshold = 0.125)
-    ρᵣ = buoyancy.equation_of_state.reference_density
-    g  = buoyancy.gravitational_acceleration
+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 and C, where C is a named tuple of (; T, S), (; T) or (; S)"
+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, C) = nothing
+validate_criterion_model(::DensityAnomalyCriterion, buoyancy_formulation, C) = nothing
     
-@inline function anomaly(criterion::DensityAnomalyCriterion, i, j, k, grid, buoyancy, C)
+@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
-
-    b = buoyancy_perturbationᶜᶜᶜ(i, j, k, grid, buoyancy, C)
-
     return - ρᵣ * b / g
 end
 
diff --git a/test/test_tracer_diagnostics.jl b/test/test_tracer_diagnostics.jl
index 6400073b..1ccecf69 100644
--- a/test/test_tracer_diagnostics.jl
+++ b/test/test_tracer_diagnostics.jl
@@ -249,11 +249,11 @@ end
 
                     @info "            Testing buoyancy diagnostics"
                     test_buoyancy_diagnostics(model)
+                end
 
-                    if !isnothing(buoyancy)
-                        @info "            Testing mixed layer depth diagnostic"
-                        test_mixed_layer_depth(grid, buoyancy)
-                    end
+                if !isnothing(buoyancy)
+                    @info "            Testing mixed layer depth diagnostic"
+                    test_mixed_layer_depth(grid, buoyancy)
                 end
             end
         end

From 7225e51ceb2e258d00ce4563ea878b0bc67e8339 Mon Sep 17 00:00:00 2001
From: tomaschor <tomaschor@gmail.com>
Date: Wed, 9 Apr 2025 19:59:40 -0300
Subject: [PATCH 16/18] remove SIGNATURES that dont do anything

---
 src/FlowDiagnostics.jl | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/src/FlowDiagnostics.jl b/src/FlowDiagnostics.jl
index 3b55ee61..6788bdc1 100755
--- a/src/FlowDiagnostics.jl
+++ b/src/FlowDiagnostics.jl
@@ -342,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)
@@ -377,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)
@@ -419,7 +418,9 @@ end
     Ω² = vorticity_tensor_modulus_ccc(i, j, k, grid, u, v, w)^2
     return (Ω² - S²) / 2
 end
+#---
 
+#+++ Mixed layer depth
 """
     $(SIGNATURES)
 
@@ -515,8 +516,6 @@ end
 end
 
 """
-    $(SIGNATURES)
-
 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).
 
@@ -535,8 +534,6 @@ validate_criterion_model(::BuoyancyAnomalyCriterion, buoyancy_formulation, C) =
 @inline anomaly(::BuoyancyAnomalyCriterion, i, j, k, grid, buoyancy_formulation, C) = buoyancy_perturbationᶜᶜᶜ(i, j, k, grid, buoyancy_formulation, C)
 
 """
-    $(SIGNATURES)
-
 Defines the mixed layer to be the depth at which the density is more than `threshold`
 greater than the surface density.
 
@@ -570,5 +567,6 @@ validate_criterion_model(::DensityAnomalyCriterion, buoyancy_formulation, C) = n
     g  = criterion.gravitational_acceleration
     return - ρᵣ * b / g
 end
+#---
 
 end # module

From 0554ff2e49533776ed418cc2bd6327ada1305c85 Mon Sep 17 00:00:00 2001
From: tomaschor <tomaschor@gmail.com>
Date: Wed, 9 Apr 2025 20:06:22 -0300
Subject: [PATCH 17/18] added TYPEDEF to structs and types

---
 src/FlowDiagnostics.jl | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/src/FlowDiagnostics.jl b/src/FlowDiagnostics.jl
index 6788bdc1..8578abac 100755
--- a/src/FlowDiagnostics.jl
+++ b/src/FlowDiagnostics.jl
@@ -448,6 +448,9 @@ end
 
 const Q = QVelocityGradientTensorInvariant
 
+"""
+    $(TYPEDEF)
+"""
 struct MixedLayerDepthKernel{C}
     criterion::C
 end
@@ -484,6 +487,8 @@ function (MLD::MixedLayerDepthKernel)(i, j, k, grid, args...)
 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`.
 
@@ -516,6 +521,8 @@ end
 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).
 
@@ -534,6 +541,8 @@ validate_criterion_model(::BuoyancyAnomalyCriterion, buoyancy_formulation, C) =
 @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.
 

From 32a152cd46455e6d46212a63133a95d89c1456c2 Mon Sep 17 00:00:00 2001
From: tomaschor <tomaschor@gmail.com>
Date: Wed, 9 Apr 2025 20:51:07 -0300
Subject: [PATCH 18/18] test that MLD doesnt work for buoyancy==Nothing

---
 test/test_tracer_diagnostics.jl | 15 ++++++---------
 1 file changed, 6 insertions(+), 9 deletions(-)

diff --git a/test/test_tracer_diagnostics.jl b/test/test_tracer_diagnostics.jl
index 1ccecf69..3359719f 100644
--- a/test/test_tracer_diagnostics.jl
+++ b/test/test_tracer_diagnostics.jl
@@ -4,7 +4,7 @@ using CUDA: has_cuda_gpu, @allowscalar
 using Oceananigans
 using Oceananigans: fill_halo_regions!
 using Oceananigans.AbstractOperations: AbstractOperation
-using Oceananigans.Fields: @compute, ConstantField
+using Oceananigans.Fields: @compute
 using Oceananigans.Grids: znode
 using Oceananigans.TurbulenceClosures.Smagorinskys: LagrangianAveraging
 using Oceananigans.BuoyancyFormulations: buoyancy
@@ -177,32 +177,27 @@ 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 
 
@@ -251,9 +246,11 @@ end
                     test_buoyancy_diagnostics(model)
                 end
 
+                @info "            Testing mixed layer depth diagnostic"
                 if !isnothing(buoyancy)
-                    @info "            Testing mixed layer depth diagnostic"
                     test_mixed_layer_depth(grid, buoyancy)
+                else
+                    @test_throws ErrorException test_mixed_layer_depth(grid, buoyancy)
                 end
             end
         end