From 347842003f69cee41c579aff5e79968b67cbc1c2 Mon Sep 17 00:00:00 2001
From: Zhaoyi Shen <11598433+szy21@users.noreply.github.com>
Date: Mon, 15 Apr 2024 11:10:26 -0700
Subject: [PATCH] move buoyancy gradient and strain rate norm to grid-mean
 cache

---
 src/cache/precomputed_quantities.jl         |  6 +--
 src/diagnostics/core_diagnostics.jl         | 32 ++++++++++++
 src/diagnostics/edmfx_diagnostics.jl        | 58 ---------------------
 src/prognostic_equations/gm_sgs_closures.jl |  6 +--
 4 files changed, 37 insertions(+), 65 deletions(-)

diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl
index 3376146682..c83bd1a2bc 100644
--- a/src/cache/precomputed_quantities.jl
+++ b/src/cache/precomputed_quantities.jl
@@ -53,6 +53,8 @@ function precomputed_quantities(Y, atmos)
         ᶜp = similar(Y.c, FT),
         ᶜh_tot = similar(Y.c, FT),
         ᶜmixing_length = similar(Y.c, FT),
+        ᶜlinear_buoygrad = similar(Y.c, FT),
+        ᶜstrain_rate_norm = similar(Y.c, FT),
         sfc_conditions = Fields.Field(SCT, Spaces.level(axes(Y.f), half)),
     )
     cloud_diagnostics = (; ᶜcloud_fraction = similar(Y.c, FT),)
@@ -83,8 +85,6 @@ function precomputed_quantities(Y, atmos)
             ᶜq_tot⁰ = similar(Y.c, FT),
             ᶜts⁰ = similar(Y.c, TST),
             ᶜρ⁰ = similar(Y.c, FT),
-            ᶜlinear_buoygrad = similar(Y.c, FT),
-            ᶜstrain_rate_norm = similar(Y.c, FT),
             ᶜK_u = similar(Y.c, FT),
             ᶜK_h = similar(Y.c, FT),
             ρatke_flux = similar(Fields.level(Y.f, half), C3{FT}),
@@ -123,8 +123,6 @@ function precomputed_quantities(Y, atmos)
             ᶜu⁰ = similar(Y.c, C123{FT}),
             ᶜK⁰ = similar(Y.c, FT),
             ᶜtke⁰ = similar(Y.c, FT),
-            ᶜlinear_buoygrad = similar(Y.c, FT),
-            ᶜstrain_rate_norm = similar(Y.c, FT),
             ᶜK_u = similar(Y.c, FT),
             ᶜK_h = similar(Y.c, FT),
             ρatke_flux = similar(Fields.level(Y.f, half), C3{FT}),
diff --git a/src/diagnostics/core_diagnostics.jl b/src/diagnostics/core_diagnostics.jl
index c710c7c057..de2aa0eafd 100644
--- a/src/diagnostics/core_diagnostics.jl
+++ b/src/diagnostics/core_diagnostics.jl
@@ -246,6 +246,38 @@ add_diagnostic_variable!(
     end,
 )
 
+###
+# Buoyancy gradient (3d)
+###
+add_diagnostic_variable!(
+    short_name = "bgrad",
+    long_name = "Linearized Buoyancy Gradient",
+    units = "s^-2",
+    compute! = (out, state, cache, time) -> begin
+        if isnothing(out)
+            return copy(cache.precomputed.ᶜlinear_buoygrad)
+        else
+            out .= cache.precomputed.ᶜlinear_buoygrad
+        end
+    end,
+)
+
+###
+# Strain rate magnitude (3d)
+###
+add_diagnostic_variable!(
+    short_name = "strain",
+    long_name = "String Rate Magnitude",
+    units = "s^-2",
+    compute! = (out, state, cache, time) -> begin
+        if isnothing(out)
+            return copy(cache.precomputed.ᶜstrain_rate_norm)
+        else
+            out .= cache.precomputed.ᶜstrain_rate_norm
+        end
+    end,
+)
+
 ###
 # Relative humidity (3d)
 ###
diff --git a/src/diagnostics/edmfx_diagnostics.jl b/src/diagnostics/edmfx_diagnostics.jl
index a7eb2f2d0d..cfd6982a86 100644
--- a/src/diagnostics/edmfx_diagnostics.jl
+++ b/src/diagnostics/edmfx_diagnostics.jl
@@ -882,64 +882,6 @@ add_diagnostic_variable!(
     compute! = compute_tke!,
 )
 
-###
-# Buoyancy gradient (3d)
-###
-compute_bgrad!(out, state, cache, time) =
-    compute_bgrad!(out, state, cache, time, cache.atmos.turbconv_model)
-compute_bgrad!(_, _, _, _, turbconv_model::T) where {T} =
-    error_diagnostic_variable("bgrad", turbconv_model)
-
-function compute_bgrad!(
-    out,
-    state,
-    cache,
-    time,
-    turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX},
-)
-    if isnothing(out)
-        return copy(cache.precomputed.ᶜlinear_buoygrad)
-    else
-        out .= cache.precomputed.ᶜlinear_buoygrad
-    end
-end
-
-add_diagnostic_variable!(
-    short_name = "bgrad",
-    long_name = "Linearized Buoyancy Gradient",
-    units = "s^-2",
-    compute! = compute_bgrad!,
-)
-
-###
-# Strain rate magnitude (3d)
-###
-compute_strain!(out, state, cache, time) =
-    compute_strain!(out, state, cache, time, cache.atmos.turbconv_model)
-compute_strain!(_, _, _, _, turbconv_model::T) where {T} =
-    error_diagnostic_variable("strain", turbconv_model)
-
-function compute_strain!(
-    out,
-    state,
-    cache,
-    time,
-    turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX},
-)
-    if isnothing(out)
-        return copy(cache.precomputed.ᶜstrain_rate_norm)
-    else
-        out .= cache.precomputed.ᶜstrain_rate_norm
-    end
-end
-
-add_diagnostic_variable!(
-    short_name = "strain",
-    long_name = "String Rate Magnitude",
-    units = "s^-2",
-    compute! = compute_strain!,
-)
-
 ###
 # Diffusivity of heat (3d)
 ###
diff --git a/src/prognostic_equations/gm_sgs_closures.jl b/src/prognostic_equations/gm_sgs_closures.jl
index 6635d2a4ce..84adff6544 100644
--- a/src/prognostic_equations/gm_sgs_closures.jl
+++ b/src/prognostic_equations/gm_sgs_closures.jl
@@ -29,10 +29,9 @@ NVTX.@annotate function compute_gm_mixing_length!(ᶜmixing_length, Y, p)
 
     ᶜdz = Fields.Δz_field(axes(Y.c))
     ᶜlg = Fields.local_geometry_field(Y.c)
-    (; ᶜts, ᶜp, ᶠu³) = p.precomputed
+    (; ᶜts, ᶠu³, ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed
     (; obukhov_length) = p.precomputed.sfc_conditions
 
-    ᶜlinear_buoygrad = p.scratch.ᶜtemp_scalar
     @. ᶜlinear_buoygrad = buoyancy_gradients(
         BuoyGradMean(),
         thermo_params,
@@ -53,13 +52,14 @@ NVTX.@annotate function compute_gm_mixing_length!(ᶜmixing_length, Y, p)
     @. ᶠu = C123(ᶠinterp(Y.c.uₕ)) + C123(ᶠu³)
     ᶜstrain_rate = p.scratch.ᶜtemp_UVWxUVW
     compute_strain_rate_center!(ᶜstrain_rate, ᶠu)
+    @. ᶜstrain_rate_norm = norm_sqr(ᶜstrain_rate)
 
     ᶜprandtl_nvec = p.scratch.ᶜtemp_scalar_2
     @. ᶜprandtl_nvec = turbulent_prandtl_number(
         params,
         obukhov_length,
         ᶜlinear_buoygrad,
-        norm_sqr(ᶜstrain_rate),
+        ᶜstrain_rate_norm,
     )
 
     @. ᶜmixing_length = smagorinsky_lilly_length(