From 388219e6a37264d086c205760319608ea5c99caa Mon Sep 17 00:00:00 2001
From: Zhaoyi Shen <11598433+szy21@users.noreply.github.com>
Date: Wed, 6 Sep 2023 17:25:44 -0700
Subject: [PATCH] add shear production

---
 src/cache/diagnostic_edmf_precomputed_quantities.jl | 10 +++++++---
 src/cache/temporary_quantities.jl                   |  1 +
 src/utils/abbreviations.jl                          |  1 +
 3 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/src/cache/diagnostic_edmf_precomputed_quantities.jl b/src/cache/diagnostic_edmf_precomputed_quantities.jl
index 205767949de..de6e795cba6 100644
--- a/src/cache/diagnostic_edmf_precomputed_quantities.jl
+++ b/src/cache/diagnostic_edmf_precomputed_quantities.jl
@@ -654,9 +654,13 @@ function set_diagnostic_edmf_precomputed_quantities!(Y, p, t)
     # TODO: This is not correct with topography or with grid stretching
     ᶠ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)
+    ∇ᵥuvw = p.ᶜtemp_UVWxUVW
+    ∇ᵥuvw = Geometry.project.(Ref(Geometry.UVWAxis()), ᶜgradᵥ.(UVW.(ᶠu⁰)))
+    @. ᶜshear² = norm_sqr((∇ᵥuvw + adjoint(∇ᵥuvw)) / 2)
+
+    # ct3_unit = p.ᶜtemp_CT3
+    # @. ct3_unit = CT3(Geometry.WVector(FT(1)), ᶜlg)
+    # @. ᶜshear² = norm_sqr(adjoint(CA.ᶜgradᵥ(ᶠu⁰)) * ct3_unit)
     ᶜprandtl_nvec = p.ᶜtemp_scalar
     @. ᶜprandtl_nvec = turbulent_prandtl_number(
         params,
diff --git a/src/cache/temporary_quantities.jl b/src/cache/temporary_quantities.jl
index d19992a4fde..5c9c6bf3eec 100644
--- a/src/cache/temporary_quantities.jl
+++ b/src/cache/temporary_quantities.jl
@@ -51,6 +51,7 @@ 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), # ∇ᵥuvw
         # TODO: Remove this hack
         sfc_temp_C3 = Fields.Field(C3{FT}, Spaces.level(face_space, half)), # ρ_flux_χ
     )
diff --git a/src/utils/abbreviations.jl b/src/utils/abbreviations.jl
index 17acc225a52..a374cbd3960 100644
--- 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()