use existing scratch, use .= to avoid allocation

Author: costachris

src/cache/cloud_fraction.jl

@@ -144,9 +144,8 @@ NVTX.@annotate function set_cloud_fraction!(
     # environment
     diagnostic_covariance_coeff = CAP.diagnostic_covariance_coeff(params)
 
-    ᶜmixing_length = p.scratch.ᶜtemp_scalar_6
-    ᶜmixing_length_lazy = mixing_length(Y, p)
-    ᶜmixing_length = Base.Broadcast.materialize(ᶜmixing_length_lazy)
+    ᶜmixing_length = p.scratch.ᶜtemp_scalar
+    ᶜmixing_length .= mixing_length(Y, p)
 
     @. cloud_diagnostics_tuple = quad_loop(
         SG_quad,
@@ -202,9 +201,8 @@ NVTX.@annotate function set_cloud_fraction!(
     # environment
     diagnostic_covariance_coeff = CAP.diagnostic_covariance_coeff(params)
 
-    ᶜmixing_length = p.scratch.ᶜtemp_scalar_6
-    ᶜmixing_length_lazy = mixing_length(Y, p)
-    ᶜmixing_length = Base.Broadcast.materialize(ᶜmixing_length_lazy)
+    ᶜmixing_length = p.scratch.ᶜtemp_scalar
+    ᶜmixing_length .= mixing_length(Y, p)
 
     @. cloud_diagnostics_tuple = quad_loop(
         SG_quad,

src/cache/temporary_quantities.jl

@@ -18,8 +18,6 @@ function temporary_quantities(Y, atmos)
         ᶜtemp_scalar_3 = Fields.Field(FT, center_space),
         ᶜtemp_scalar_4 = Fields.Field(FT, center_space),
         ᶜtemp_scalar_5 = Fields.Field(FT, center_space),
-        ᶜtemp_scalar_6 = Fields.Field(FT, center_space),
-        ᶜtemp_scalar_7 = Fields.Field(FT, center_space),
         ᶠtemp_field_level = Fields.level(Fields.Field(FT, face_space), half),
         temp_field_level = Fields.level(Fields.Field(FT, center_space), 1),
         temp_field_level_2 = Fields.level(Fields.Field(FT, center_space), 1),

src/diagnostics/core_diagnostics.jl

@@ -284,11 +284,12 @@ add_diagnostic_variable!(
             (; params) = cache
             (; ᶜlinear_buoygrad, ᶜstrain_rate_norm) = cache.precomputed
             ᶜdz = Fields.Δz_field(axes(state.c))
-            ᶜprandtl_nvec = cache.scratch.ᶜtemp_scalar_2
-            @. ᶜprandtl_nvec = turbulent_prandtl_number(
-                params,
-                ᶜlinear_buoygrad,
-                ᶜstrain_rate_norm,
+            ᶜprandtl_nvec = @. lazy(
+                turbulent_prandtl_number(
+                    params,
+                    ᶜlinear_buoygrad,
+                    ᶜstrain_rate_norm,
+                ),
             )
             ᶜmixing_length = @. lazy(
                 smagorinsky_lilly_length(

src/prognostic_equations/eddy_diffusion_closures.jl

@@ -576,9 +576,8 @@ function mixing_length(Y, p, mixing_length_property = :master)
     ᶜdz = Fields.Δz_field(axes(Y.c))
     sfc_tke = Fields.level(ᶜtke⁰, 1)
 
-    ᶜprandtl_nvec_lazy = turbulent_prandtl_number(p)
-    ᶜprandtl_nvec = p.scratch.ᶜtemp_scalar_7
-    ᶜprandtl_nvec = Base.Broadcast.materialize(ᶜprandtl_nvec_lazy)
+    ᶜprandtl_nvec = p.scratch.ᶜtemp_scalar_5
+    ᶜprandtl_nvec .= ᶜturbulent_prandtl_number(p)
 
     ᶜtke_exch = tke_exchange(Y, p)
 
@@ -700,7 +699,7 @@ function turbulent_prandtl_number(params, ᶜN²_eff, ᶜstrain_rate_norm)
 end
 
 
-function turbulent_prandtl_number(p)
+function ᶜturbulent_prandtl_number(p)
     (; params) = p
     (; ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed
     return @. lazy(
@@ -890,6 +889,6 @@ end
 
 
 function eddy_diffusivity(p, K_u)
-    ᶜprandtl_nvec = turbulent_prandtl_number(p)
+    ᶜprandtl_nvec = ᶜturbulent_prandtl_number(p)
     return @. lazy(K_u / ᶜprandtl_nvec) # prandtl_nvec is already bounded by eps_FT and Pr_max
 end

src/prognostic_equations/edmfx_sgs_flux.jl

@@ -353,9 +353,8 @@ function edmfx_sgs_diffusive_flux_tendency!(
 
     if p.atmos.edmfx_model.sgs_diffusive_flux isa Val{true}
 
-        ᶜmixing_length = p.scratch.ᶜtemp_scalar_6
-        ᶜmixing_length_lazy = mixing_length(Y, p)
-        ᶜmixing_length = Base.Broadcast.materialize(ᶜmixing_length_lazy)
+        ᶜmixing_length = p.scratch.ᶜtemp_scalar_2
+        ᶜmixing_length .= mixing_length(Y, p)
         ᶜK_u = eddy_viscosity(turbconv_params, ᶜtke⁰, ᶜmixing_length)
         ᶜK_h = eddy_diffusivity(p, ᶜK_u)
         ᶠρaK_h = p.scratch.ᶠtemp_scalar
@@ -458,9 +457,8 @@ function edmfx_sgs_diffusive_flux_tendency!(
 
     if p.atmos.edmfx_model.sgs_diffusive_flux isa Val{true}
 
-        ᶜmixing_length = p.scratch.ᶜtemp_scalar_6
-        ᶜmixing_length_lazy = mixing_length(Y, p)
-        ᶜmixing_length = Base.Broadcast.materialize(ᶜmixing_length_lazy)
+        ᶜmixing_length = p.scratch.ᶜtemp_scalar_2
+        ᶜmixing_length .= mixing_length(Y, p)
         ᶜK_u = eddy_viscosity(turbconv_params, ᶜtke⁰, ᶜmixing_length)
         ᶜK_h = eddy_diffusivity(p, ᶜK_u)
 

src/prognostic_equations/edmfx_tke.jl

@@ -29,9 +29,8 @@ function edmfx_tke_tendency!(Yₜ, Y, p, t, turbconv_model::EDOnlyEDMFX)
     (; ᶜstrain_rate_norm, ᶜlinear_buoygrad, ᶜtke⁰) = p.precomputed
     turbconv_params = CAP.turbconv_params(p.params)
 
-    ᶜmixing_length = p.scratch.ᶜtemp_scalar_6
-    ᶜmixing_length_lazy = mixing_length(Y, p)
-    ᶜmixing_length = Base.Broadcast.materialize(ᶜmixing_length_lazy)
+    ᶜmixing_length = p.scratch.ᶜtemp_scalar
+    ᶜmixing_length .= mixing_length(Y, p)
     ᶜK_u = eddy_viscosity(turbconv_params, ᶜtke⁰, ᶜmixing_length)
     ᶜK_h = eddy_diffusivity(p, ᶜK_u)
 
@@ -77,9 +76,8 @@ function edmfx_tke_tendency!(
 
     if use_prognostic_tke(turbconv_model)
 
-        ᶜmixing_length = p.scratch.ᶜtemp_scalar_6
-        ᶜmixing_length_lazy = mixing_length(Y, p)
-        ᶜmixing_length = Base.Broadcast.materialize(ᶜmixing_length_lazy)
+        ᶜmixing_length = p.scratch.ᶜtemp_scalar_2
+        ᶜmixing_length .= mixing_length(Y, p)
         ᶜK_u = eddy_viscosity(turbconv_params, ᶜtke⁰, ᶜmixing_length)
         ᶜK_h = eddy_diffusivity(p, ᶜK_u)
 

src/prognostic_equations/implicit/manual_sparse_jacobian.jl

@@ -555,9 +555,8 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
             ᶜK_u = ᶜK_h
         else
             (; ᶜtke⁰,) = p.precomputed
-            ᶜmixing_length = p.scratch.ᶜtemp_scalar_6
-            ᶜmixing_length_lazy = mixing_length(Y, p)
-            ᶜmixing_length = Base.Broadcast.materialize(ᶜmixing_length_lazy)
+            ᶜmixing_length = p.scratch.ᶜtemp_scalar_3
+            ᶜmixing_length .= mixing_length(Y, p)
             ᶜK_u = eddy_viscosity(turbconv_params, ᶜtke⁰, ᶜmixing_length)
             ᶜK_h = eddy_diffusivity(p, ᶜK_u)
         end