CliMA · bors · Sep 8, 2023 · Sep 7, 2023
diff --git a/perf/flame.jl b/perf/flame.jl
@@ -56,12 +56,12 @@ allocs = @allocated OrdinaryDiffEq.step!(integrator)
 allocs_limit = Dict()
 allocs_limit["flame_perf_target"] = 4384
 allocs_limit["flame_perf_target_tracers"] = 204016
-allocs_limit["flame_perf_target_edmfx"] = 301760
-allocs_limit["flame_perf_target_diagnostic_edmfx"] = 678672
-allocs_limit["flame_perf_target_edmf"] = 8953312080
+allocs_limit["flame_perf_target_edmfx"] = 304064
+allocs_limit["flame_perf_target_diagnostic_edmfx"] = 685456
+allocs_limit["flame_perf_target_edmf"] = 9015243600
 allocs_limit["flame_perf_target_threaded"] = 6175664
 allocs_limit["flame_perf_target_callbacks"] = 49850536
-allocs_limit["flame_perf_gw"] = 4964581344
+allocs_limit["flame_perf_gw"] = 4985829472
 
 if allocs < allocs_limit[job_id] * buffer
     @info "TODO: lower `allocs_limit[$job_id]` to: $(allocs)"

diff --git a/src/cache/diagnostic_edmf_precomputed_quantities.jl b/src/cache/diagnostic_edmf_precomputed_quantities.jl
@@ -146,7 +146,7 @@ function set_diagnostic_edmf_precomputed_quantities!(Y, p, t)
         ᶜS_e_totʲs_helper,
     ) = p
     (; ᶠu³⁰, ᶜu⁰, ᶜK⁰, ᶜh_tot⁰, ᶜtke⁰, ᶜS_q_tot⁰) = p
-    (; ᶜlinear_buoygrad, ᶜshear², ᶜmixing_length) = p
+    (; ᶜlinear_buoygrad, ᶜstrain_rate_norm, ᶜmixing_length) = p
     (; ᶜK_h, ᶜK_u, ρatke_flux) = p
     thermo_params = CAP.thermodynamics_params(params)
     microphys_params = CAP.microphysics_params(params)
@@ -651,18 +651,19 @@ function set_diagnostic_edmf_precomputed_quantities!(Y, p, t)
         ),
     )
 
-    # TODO: This is not correct with topography or with grid stretching
+    # TODO: Currently the shear production only includes vertical gradients
     ᶠu⁰ = p.ᶠtemp_C123
     @. ᶠu⁰ = C123(ᶠinterp(Y.c.uₕ)) + C123(ᶠu³⁰)
-    ct3_unit = p.ᶜtemp_CT3
-    @. ct3_unit = CT3(Geometry.WVector(FT(1)), ᶜlg)
-    @. ᶜshear² = norm_sqr(adjoint(CA.ᶜgradᵥ(ᶠu⁰)) * ct3_unit)
+    ᶜstrain_rate = p.ᶜtemp_UVWxUVW
+    compute_strain_rate!(ᶜstrain_rate, ᶠu⁰)
+    @. ᶜstrain_rate_norm = norm_sqr(ᶜstrain_rate)
+
     ᶜprandtl_nvec = p.ᶜtemp_scalar
     @. ᶜprandtl_nvec = turbulent_prandtl_number(
         params,
         obukhov_length,
         ᶜlinear_buoygrad,
-        ᶜshear²,
+        ᶜstrain_rate_norm,
     )
     ᶜtke_exch = p.ᶜtemp_scalar_2
     @. ᶜtke_exch = 0
@@ -686,7 +687,7 @@ function set_diagnostic_edmf_precomputed_quantities!(Y, p, t)
         ᶜlinear_buoygrad,
         max(ᶜtke⁰, 0),
         obukhov_length,
-        ᶜshear²,
+        ᶜstrain_rate_norm,
         ᶜprandtl_nvec,
         ᶜtke_exch,
     )

diff --git a/src/cache/edmf_precomputed_quantities.jl b/src/cache/edmf_precomputed_quantities.jl
@@ -24,7 +24,7 @@ function set_edmf_precomputed_quantities!(Y, p, ᶠuₕ³, t)
 
     (; ᶜspecific, ᶜp, ᶜΦ, ᶜh_tot, ᶜρ_ref) = p
     (; ᶜspecific⁰, ᶜρa⁰, ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜρ⁰, ᶜh_tot⁰) = p
-    (; ᶜmixing_length, ᶜlinear_buoygrad, ᶜshear², ᶜK_u, ᶜK_h) = p
+    (; ᶜmixing_length, ᶜlinear_buoygrad, ᶜstrain_rate_norm, ᶜK_u, ᶜK_h) = p
     (; ᶜspecificʲs, ᶜuʲs, ᶠu³ʲs, ᶜKʲs, ᶜtsʲs, ᶜρʲs, ᶜh_totʲs, ᶜentr_detrʲs) = p
     (; ustar, obukhov_length, buoyancy_flux) = p.sfc_conditions
 
@@ -197,7 +197,7 @@ function set_edmf_precomputed_quantities!(Y, p, ᶠuₕ³, t)
         ),
     )
 
-    @. ᶜshear² = $(FT(1e-4))
+    @. ᶜstrain_rate_norm = $(FT(1e-4))
     ᶜprandtl_nvec = p.ᶜtemp_scalar
     @. ᶜprandtl_nvec = FT(1) / 3
     ᶜtke_exch = p.ᶜtemp_scalar_2
@@ -223,7 +223,7 @@ function set_edmf_precomputed_quantities!(Y, p, ᶠuₕ³, t)
         ᶜlinear_buoygrad,
         ᶜspecific⁰.tke,
         obukhov_length,
-        ᶜshear²,
+        ᶜstrain_rate_norm,
         ᶜprandtl_nvec,
         ᶜtke_exch,
     )

diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl
@@ -68,7 +68,7 @@ function precomputed_quantities(Y, atmos)
             ᶜts⁰ = similar(Y.c, TST),
             ᶜρ⁰ = similar(Y.c, FT),
             ᶜlinear_buoygrad = similar(Y.c, FT),
-            ᶜshear² = similar(Y.c, FT),
+            ᶜstrain_rate_norm = similar(Y.c, FT),
             ᶜmixing_length = similar(Y.c, FT),
             ᶜK_u = similar(Y.c, FT),
             ᶜK_h = similar(Y.c, FT),
@@ -115,7 +115,7 @@ function precomputed_quantities(Y, atmos)
             ᶜh_tot⁰ = similar(Y.c, FT),
             ᶜtke⁰ = similar(Y.c, FT),
             ᶜlinear_buoygrad = similar(Y.c, FT),
-            ᶜshear² = similar(Y.c, FT),
+            ᶜstrain_rate_norm = similar(Y.c, FT),
             ᶜmixing_length = similar(Y.c, FT),
             ᶜK_u = similar(Y.c, FT),
             ᶜK_h = similar(Y.c, FT),

diff --git a/src/cache/temporary_quantities.jl b/src/cache/temporary_quantities.jl
@@ -51,6 +51,10 @@ function temporary_quantities(atmos, center_space, face_space)
         ᶠtemp_CT12ʲs = Fields.Field(NTuple{n, CT12{FT}}, face_space), # ᶠω¹²ʲs
         ᶜtemp_C123 = Fields.Field(C123{FT}, center_space), # χ₁₂₃
         ᶠtemp_C123 = Fields.Field(C123{FT}, face_space), # χ₁₂₃
+        ᶜtemp_UVWxUVW = Fields.Field(
+            typeof(UVW(FT(0), FT(0), FT(0)) * UVW(FT(0), FT(0), FT(0))'),
+            center_space,
+        ), # ᶜstrain_rate
         # TODO: Remove this hack
         sfc_temp_C3 = Fields.Field(C3{FT}, Spaces.level(face_space, half)), # ρ_flux_χ
     )

diff --git a/src/prognostic_equations/edmfx_closures.jl b/src/prognostic_equations/edmfx_closures.jl
@@ -250,7 +250,7 @@ function lamb_smooth_minimum(
 end
 
 """
-    mixing_length(params, ustar, ᶜz, sfc_tke, ᶜlinear_buoygrad, ᶜtke, obukhov_length, shear², ᶜPr, ᶜtke_exch)
+    mixing_length(params, ustar, ᶜz, sfc_tke, ᶜlinear_buoygrad, ᶜtke, obukhov_length, ᶜstrain_rate_norm, ᶜPr, ᶜtke_exch)
 
 where:
 - `params`: set with model parameters
@@ -260,7 +260,7 @@ where:
 - `ᶜlinear_buoygrad`: buoyancy gradient
 - `ᶜtke`: env turbulent kinetic energy
 - `obukhov_length`: surface Monin Obukhov length
-- `ᶜshear²`: shear term
+- `ᶜstrain_rate_norm`: Frobenius norm of strain rate tensor
 - `ᶜPr`: Prandtl number
 - `ᶜtke_exch`: subdomain exchange term
 
@@ -280,7 +280,7 @@ function mixing_length(
     ᶜlinear_buoygrad::FT,
     ᶜtke::FT,
     obukhov_length::FT,
-    ᶜshear²::FT,
+    ᶜstrain_rate_norm::FT,
     ᶜPr::FT,
     ᶜtke_exch::FT,
 ) where {FT}
@@ -306,7 +306,7 @@ function mixing_length(
     end
 
     # compute l_TKE - the production-dissipation balanced length scale
-    a_pd = c_m * (ᶜshear² - ᶜlinear_buoygrad / ᶜPr) * sqrt(ᶜtke)
+    a_pd = c_m * (ᶜstrain_rate_norm - ᶜlinear_buoygrad / ᶜPr) * sqrt(ᶜtke)
     # Dissipation term
     c_neg = c_d * ᶜtke * sqrt(ᶜtke)
     if abs(a_pd) > eps(FT) && 4 * a_pd * c_neg > -(ᶜtke_exch * ᶜtke_exch)
@@ -334,7 +334,8 @@ function mixing_length(
     c_smag = FT(0.2)
     N_eff = sqrt(max(ᶜlinear_buoygrad, 0))
     if N_eff > 0.0
-        l_smag = c_smag * ᶜdz * max(0, 1 - N_eff^2 / ᶜPr / ᶜshear²)^(1 / 4)
+        l_smag =
+            c_smag * ᶜdz * max(0, 1 - N_eff^2 / ᶜPr / ᶜstrain_rate_norm)^(1 / 4)
     else
         l_smag = c_smag * ᶜdz
     end
@@ -366,13 +367,13 @@ function turbulent_prandtl_number(
     params,
     obukhov_length::FT,
     ᶜlinear_buoygrad::FT,
-    ᶜshear²::FT,
+    ᶜstrain_rate_norm::FT,
 ) where {FT}
     turbconv_params = CAP.turbconv_params(params)
     Ri_c = TCP.Ri_crit(turbconv_params)
     ω_pr = TCP.Prandtl_number_scale(turbconv_params)
     Pr_n = TCP.Prandtl_number_0(turbconv_params)
-    ᶜRi_grad = min(ᶜlinear_buoygrad / max(ᶜshear², eps(FT)), Ri_c)
+    ᶜRi_grad = min(ᶜlinear_buoygrad / max(ᶜstrain_rate_norm, eps(FT)), Ri_c)
     if obukhov_length > 0 && ᶜRi_grad > 0 #stable
         # CSB (Dan Li, 2019, eq. 75), where ω_pr = ω_1 + 1 = 53.0 / 13.0
         prandtl_nvec =

diff --git a/src/prognostic_equations/edmfx_tke.jl b/src/prognostic_equations/edmfx_tke.jl
@@ -19,7 +19,7 @@ function edmfx_tke_tendency!(
     turbconv_params = CAP.turbconv_params(params)
     c_d = TCP.tke_diss_coeff(turbconv_params)
     (; ᶜentr_detrʲs, ᶠu³ʲs) = p
-    (; ᶠu³⁰, ᶜshear², ᶜlinear_buoygrad, ᶜtke⁰, ᶜmixing_length) = p
+    (; ᶠu³⁰, ᶜstrain_rate_norm, ᶜlinear_buoygrad, ᶜtke⁰, ᶜmixing_length) = p
     (; ᶜK_u, ᶜK_h, ρatke_flux) = p
     ᶠgradᵥ = Operators.GradientC2F()
 
@@ -36,7 +36,7 @@ function edmfx_tke_tendency!(
             ᶜdivᵥ_ρatke(-(ᶠρaK_u[colidx] * ᶠgradᵥ(ᶜtke⁰[colidx])))
         # shear production
         @. Yₜ.c.sgs⁰.ρatke[colidx] +=
-            2 * Y.c.ρ[colidx] * ᶜK_u[colidx] * ᶜshear²[colidx]
+            2 * Y.c.ρ[colidx] * ᶜK_u[colidx] * ᶜstrain_rate_norm[colidx]
         # buoyancy production
         @. Yₜ.c.sgs⁰.ρatke[colidx] -=
             Y.c.ρ[colidx] * ᶜK_h[colidx] * ᶜlinear_buoygrad[colidx]

diff --git a/src/utils/abbreviations.jl b/src/utils/abbreviations.jl
@@ -11,6 +11,7 @@ const CT2 = Geometry.Contravariant2Vector
 const CT12 = Geometry.Contravariant12Vector
 const CT3 = Geometry.Contravariant3Vector
 const CT123 = Geometry.Contravariant123Vector
+const UVW = Geometry.UVWVector
 
 const divₕ = Operators.Divergence()
 const wdivₕ = Operators.WeakDivergence()

diff --git a/src/utils/utilities.jl b/src/utils/utilities.jl
@@ -74,11 +74,27 @@ end
     compute_kinetic!(κ::Field, Y::FieldVector)
 
 Compute the specific kinetic energy at cell centers, storing in `κ`, where `Y`
-is the model state.s
+is the model state.
 """
 compute_kinetic!(κ::Fields.Field, Y::Fields.FieldVector) =
     compute_kinetic!(κ, Y.c.uₕ, Y.f.u₃)
 
+"""
+    compute_strain_rate!(ϵ::Field, u::Field)
+
+Compute the strain_rate at cell centers, storing in `ϵ` from
+velocity at cell faces.
+"""
+function compute_strain_rate!(ϵ::Fields.Field, u::Fields.Field)
+    @assert eltype(u) <: C123
+    axis_uvw = Geometry.UVWAxis()
+    @. ϵ =
+        (
+            Geometry.project((axis_uvw,), ᶜgradᵥ(UVW(u))) +
+            adjoint(Geometry.project((axis_uvw,), ᶜgradᵥ(UVW(u))))
+        ) / 2
+end
+
 """
     g³³_field(field)