Merge pull request #2951 from CliMA/ck/buoyancy_grad_perf

Flatten buoyancy_gradients call

Author: charleskawczynski

src/cache/diagnostic_edmf_precomputed_quantities.jl

@@ -782,16 +782,12 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
         BuoyGradMean(),
         thermo_params,
         moisture_model,
-        EnvBuoyGradVars(
-            ᶜts,
-            projected_vector_buoy_grad_vars(
-                C3,
-                p.precomputed.ᶜgradᵥ_θ_virt,    # ∂θv∂z_unsat
-                p.precomputed.ᶜgradᵥ_q_tot,     # ∂qt∂z_sat
-                p.precomputed.ᶜgradᵥ_θ_liq_ice, # ∂θl∂z_sat
-                ᶜlg,
-            ),
-        ),
+        ᶜts,
+        C3,
+        p.precomputed.ᶜgradᵥ_θ_virt,    # ∂θv∂z_unsat
+        p.precomputed.ᶜgradᵥ_q_tot,     # ∂qt∂z_sat
+        p.precomputed.ᶜgradᵥ_θ_liq_ice, # ∂θl∂z_sat
+        ᶜlg,
     )
 
     # TODO: Currently the shear production only includes vertical gradients

src/cache/precomputed_quantities.jl

@@ -101,6 +101,9 @@ function precomputed_quantities(Y, atmos)
             ᶜentrʲs = similar(Y.c, NTuple{n, FT}),
             ᶜdetrʲs = similar(Y.c, NTuple{n, FT}),
             ᶠnh_pressure₃ʲs = similar(Y.f, NTuple{n, C3{FT}}),
+            ᶜgradᵥ_θ_virt⁰ = Fields.Field(C3{FT}, cspace),
+            ᶜgradᵥ_q_tot⁰ = Fields.Field(C3{FT}, cspace),
+            ᶜgradᵥ_θ_liq_ice⁰ = Fields.Field(C3{FT}, cspace),
             precipitation_sgs_quantities...,
         ) : (;)
     sgs_quantities = (;

src/cache/prognostic_edmf_precomputed_quantities.jl

@@ -267,21 +267,22 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_closures!(
         )
     end
 
+    (; ᶜgradᵥ_θ_virt⁰, ᶜgradᵥ_q_tot⁰, ᶜgradᵥ_θ_liq_ice⁰) = p.precomputed
     # First order approximation: Use environmental mean fields.
+    @. ᶜgradᵥ_θ_virt⁰ = ᶜgradᵥ(ᶠinterp(TD.virtual_pottemp(thermo_params, ᶜts⁰)))       # ∂θv∂z_unsat
+    @. ᶜgradᵥ_q_tot⁰ = ᶜgradᵥ(ᶠinterp(ᶜq_tot⁰))                                        # ∂qt∂z_sat
+    @. ᶜgradᵥ_θ_liq_ice⁰ =
+        ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts⁰)))                    # ∂θl∂z_sat
     @. ᶜlinear_buoygrad = buoyancy_gradients(
         BuoyGradMean(),
         thermo_params,
         moisture_model,
-        EnvBuoyGradVars(
-            ᶜts⁰,
-            projected_vector_buoy_grad_vars(
-                C3,
-                ᶜgradᵥ(ᶠinterp(TD.virtual_pottemp(thermo_params, ᶜts⁰))),    # ∂θv∂z_unsat
-                ᶜgradᵥ(ᶠinterp(ᶜq_tot⁰)),                                    # ∂qt∂z_sat
-                ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts⁰))), # ∂θl∂z_sat
-                ᶜlg,
-            ),
-        ),
+        ᶜts⁰,
+        C3,
+        ᶜgradᵥ_θ_virt⁰,
+        ᶜgradᵥ_q_tot⁰,
+        ᶜgradᵥ_θ_liq_ice⁰,
+        ᶜlg,
     )
 
     # TODO: Currently the shear production only includes vertical gradients

src/prognostic_equations/buoyancy_gradients.jl

@@ -34,6 +34,36 @@ The dispatch on EnvBuoyGradVars type is performed at the EnvBuoyGradVars constru
 used here are consistent for both mean fields and conditional fields obtained from assumed distributions
 over the conserved thermodynamic variables.
 """
+function buoyancy_gradients end
+
+function buoyancy_gradients(
+    ebgc::AbstractEnvBuoyGradClosure,
+    thermo_params,
+    moisture_model,
+    ts,
+    ::Type{C3},
+    ∂θv∂z_unsat,
+    ∂qt∂z_sat,
+    ∂θl∂z_sat,
+    ᶜlg,
+) where {C3}
+    return buoyancy_gradients(
+        ebgc,
+        thermo_params,
+        moisture_model,
+        EnvBuoyGradVars(
+            ts,
+            projected_vector_buoy_grad_vars(
+                C3,
+                ∂θv∂z_unsat,
+                ∂qt∂z_sat,
+                ∂θl∂z_sat,
+                ᶜlg,
+            ),
+        ),
+    )
+end
+
 function buoyancy_gradients(
     ebgc::AbstractEnvBuoyGradClosure,
     thermo_params,

src/prognostic_equations/gm_sgs_closures.jl

@@ -36,16 +36,12 @@ NVTX.@annotate function compute_gm_mixing_length!(ᶜmixing_length, Y, p)
         BuoyGradMean(),
         thermo_params,
         p.atmos.moisture_model,
-        EnvBuoyGradVars(
-            ᶜts,
-            projected_vector_buoy_grad_vars(
-                C3,
-                p.precomputed.ᶜgradᵥ_θ_virt,    # ∂θv∂z_unsat
-                p.precomputed.ᶜgradᵥ_q_tot,     # ∂qt∂z_sat
-                p.precomputed.ᶜgradᵥ_θ_liq_ice, # ∂θl∂z_sat
-                ᶜlg,
-            ),
-        ),
+        ᶜts,
+        C3,
+        p.precomputed.ᶜgradᵥ_θ_virt,    # ∂θv∂z_unsat
+        p.precomputed.ᶜgradᵥ_q_tot,     # ∂qt∂z_sat
+        p.precomputed.ᶜgradᵥ_θ_liq_ice, # ∂θl∂z_sat
+        ᶜlg,
     )
 
     ᶠu = p.scratch.ᶠtemp_C123