restructure env, specific helper functions

Author: costachris

src/cache/cloud_fraction.jl

@@ -84,7 +84,6 @@ NVTX.@annotate function set_cloud_fraction!(
     (; turbconv_model) = p.atmos
 
     if isnothing(turbconv_model)
-        ᶜρa⁰ = @.lazy(ρa⁰(Y.c))
         if p.atmos.call_cloud_diagnostics_per_stage isa
            CallCloudDiagnosticsPerStage
             (; ᶜgradᵥ_θ_virt, ᶜgradᵥ_q_tot, ᶜgradᵥ_θ_liq_ice) = p.precomputed
@@ -191,7 +190,7 @@ NVTX.@annotate function set_cloud_fraction!(
     (; ᶜts⁰, ᶜmixing_length, cloud_diagnostics_tuple) = p.precomputed
     (; ᶜρʲs, ᶜtsʲs) = p.precomputed
     (; turbconv_model) = p.atmos
-    ᶜρa⁰ = @.lazy(ρa⁰(Y.c))
+    ᶜρa⁰_vals = ᶜρa⁰(Y.c, p)
 
     # TODO - we should make this default when using diagnostic edmf
     # environment
@@ -210,9 +209,9 @@ NVTX.@annotate function set_cloud_fraction!(
     # weight cloud diagnostics by environmental area
     @. cloud_diagnostics_tuple *= NamedTuple{(:cf, :q_liq, :q_ice)}(
         tuple(
-            draft_area(ᶜρa⁰, TD.air_density(thermo_params, ᶜts⁰)),
-            draft_area(ᶜρa⁰, TD.air_density(thermo_params, ᶜts⁰)),
-            draft_area(ᶜρa⁰, TD.air_density(thermo_params, ᶜts⁰)),
+            draft_area(ᶜρa⁰_vals, TD.air_density(thermo_params, ᶜts⁰)),
+            draft_area(ᶜρa⁰_vals, TD.air_density(thermo_params, ᶜts⁰)),
+            draft_area(ᶜρa⁰_vals, TD.air_density(thermo_params, ᶜts⁰)),
         ),
     )
     # updrafts

src/cache/diagnostic_edmf_precomputed_quantities.jl

@@ -105,21 +105,24 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_bottom_bc!(
     turbconv_params = CAP.turbconv_params(params)
     ᶜts = p.precomputed.ᶜts   #TODO replace
 
-    q_tot = specific(Y.c.ρq_tot, Y.c.ρ)
+    q_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ))
+    ᶜe_tot = @. lazy(specific(Y.c.ρe_tot, Y.c.ρ))
     ᶜh_tot = @. lazy(
         TD.total_specific_enthalpy(
             thermo_params,
             ᶜts,
-            specific(Y.c.ρe_tot, Y.c.ρ),
+            ᶜe_tot,
         ),
     )
 
     ρ_int_level = Fields.field_values(Fields.level(Y.c.ρ, 1))
     uₕ_int_level = Fields.field_values(Fields.level(Y.c.uₕ, 1))
     u³_int_halflevel = Fields.field_values(Fields.level(ᶠu³, half))
-    h_tot_int_level = Fields.field_values(Fields.level(ᶜh_tot, 1))
+    h_tot_int_level =
+        Fields.field_values(Fields.level(Base.materialize(ᶜh_tot), 1))
     K_int_level = Fields.field_values(Fields.level(ᶜK, 1))
-    q_tot_int_level = Fields.field_values(Fields.level(q_tot, 1))
+    q_tot_int_level =
+        Fields.field_values(Fields.level(Base.materialize(q_tot), 1))
 
     p_int_level = Fields.field_values(Fields.level(ᶜp, 1))
     Φ_int_level = Fields.field_values(Fields.level(ᶜΦ, 1))
@@ -330,7 +333,9 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
         ᶠnh_pressure³_buoyʲs,
         ᶠnh_pressure³_dragʲs,
     ) = p.precomputed
-    (; ᶠu³⁰, ᶜK⁰, ᶜtke⁰) = p.precomputed
+    (; ᶠu³⁰, ᶜK⁰) = p.precomputed
+
+
 
     if microphysics_model isa Microphysics1Moment
         q_rai = @. lazy(specific(Y.c.ρq_rai, Y.c.ρ))
@@ -354,14 +359,25 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
         Fields.field_values(Fields.level(Fields.coordinate_field(Y.f).z, half))
 
     # integral
-    q_tot = specific(Y.c.ρq_tot, Y.c.ρ)
+    q_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ))
+    ᶜh_tot = @. lazy(
+        TD.total_specific_enthalpy(
+            thermo_params,
+            ᶜts,
+            specific(Y.c.ρe_tot, Y.c.ρ),
+        ),
+    )
+    ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p)
+
     for i in 2:Spaces.nlevels(axes(Y.c))
         ρ_level = Fields.field_values(Fields.level(Y.c.ρ, i))
         uₕ_level = Fields.field_values(Fields.level(Y.c.uₕ, i))
         u³_halflevel = Fields.field_values(Fields.level(ᶠu³, i - half))
         K_level = Fields.field_values(Fields.level(ᶜK, i))
-        h_tot_level = Fields.field_values(Fields.level(ᶜh_tot, i))
-        q_tot_level = Fields.field_values(Fields.level(q_tot, i))
+        h_tot_level =
+            Fields.field_values(Fields.level(Base.materialize(ᶜh_tot), i))
+        q_tot_level =
+            Fields.field_values(Fields.level(Base.materialize(q_tot), i))
         p_level = Fields.field_values(Fields.level(ᶜp, i))
         Φ_level = Fields.field_values(Fields.level(ᶜΦ, i))
         local_geometry_level = Fields.field_values(
@@ -385,8 +401,10 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
             Fields.field_values(Fields.level(ᶠu³⁰, i - 1 - half))
         u³⁰_data_prev_halflevel = u³⁰_prev_halflevel.components.data.:1
         K_prev_level = Fields.field_values(Fields.level(ᶜK, i - 1))
-        h_tot_prev_level = Fields.field_values(Fields.level(ᶜh_tot, i - 1))
-        q_tot_prev_level = Fields.field_values(Fields.level(q_tot, i - 1))
+        h_tot_prev_level =
+            Fields.field_values(Fields.level(Base.materialize(ᶜh_tot), i - 1))
+        q_tot_prev_level =
+            Fields.field_values(Fields.level(Base.materialize(q_tot), i - 1))
         ts_prev_level = Fields.field_values(Fields.level(ᶜts, i - 1))
         p_prev_level = Fields.field_values(Fields.level(ᶜp, i - 1))
         z_prev_level = Fields.field_values(Fields.level(ᶜz, i - 1))
@@ -486,7 +504,9 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
                     Fields.field_values(Fields.level(ᶜq_iceʲ, i - 1))
             end
 
-            tke_prev_level = Fields.field_values(Fields.level(ᶜtke⁰, i - 1))
+            tke_prev_level = Fields.field_values(
+                Fields.level(Base.materialize(ᶜtke⁰), i - 1),
+            )
 
             @. entrʲ_prev_level = entrainment(
                 thermo_params,
@@ -1015,6 +1035,8 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
     @. ᶜprandtl_nvec =
         turbulent_prandtl_number(params, ᶜlinear_buoygrad, ᶜstrain_rate_norm)
 
+    ᶜtke⁰ = ᶜspecific_tke(Y.c.sgs⁰, Y.c, p)
+
     ᶜtke_exch = p.scratch.ᶜtemp_scalar_2
     @. ᶜtke_exch = 0
     # using ᶜu⁰ would be more correct, but this is more consistent with the
@@ -1025,7 +1047,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
             (1 / 2 * norm_sqr(ᶜinterp(ᶠu³⁰) - ᶜinterp(ᶠu³ʲs.:($$j))) - ᶜtke⁰)
     end
 
-    sfc_tke = Fields.level(ᶜtke⁰, 1)
+    sfc_tke = Fields.level(Base.materialize(ᶜtke⁰), 1)
     z_sfc = Fields.level(Fields.coordinate_field(Y.f).z, half)
     @. ᶜmixing_length_tuple = mixing_length(
         params,
@@ -1090,12 +1112,12 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_precipita
     (; ᶜts, ᶜSqₜᵖ⁰) = p.precomputed
 
     # Environment precipitation sources (to be applied to grid mean)
-    q_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ))
+    ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ))
     @. ᶜSqₜᵖ⁰ = q_tot_0M_precipitation_sources(
         thermo_params,
         microphys_0m_params,
         dt,
-        specific(Y.c.ρq_tot, Y.c.ρ),
+        ᶜq_tot,
         ᶜts,
     )
     return nothing

src/cache/precipitation_precomputed_quantities.jl

@@ -262,14 +262,14 @@ function set_precipitation_cache!(
     (; ᶜS_ρq_tot, ᶜS_ρe_tot) = p.precomputed
     (; ᶜts⁰, ᶜtsʲs) = p.precomputed
     thermo_params = CAP.thermodynamics_params(p.params)
-    ᶜρa⁰ = @.lazy(ρa⁰(Y.c))
+    ᶜρa⁰_vals = ᶜρa⁰(Y.c, p)
 
     n = n_mass_flux_subdomains(p.atmos.turbconv_model)
 
-    @. ᶜS_ρq_tot = ᶜSqₜᵖ⁰ * ᶜρa⁰
+    @. ᶜS_ρq_tot = ᶜSqₜᵖ⁰ * ᶜρa⁰_vals
     @. ᶜS_ρe_tot =
         ᶜSqₜᵖ⁰ *
-        ᶜρa⁰ *
+        ᶜρa⁰_vals *
         e_tot_0M_precipitation_sources_helper(thermo_params, ᶜts⁰, ᶜΦ)
     for j in 1:n
         @. ᶜS_ρq_tot += ᶜSqₜᵖʲs.:($$j) * Y.c.sgsʲs.:($$j).ρa
@@ -289,8 +289,11 @@ function set_precipitation_cache!(Y, p, ::Microphysics1Moment, _)
     (; ᶜts, ᶜwᵣ, ᶜwₛ, ᶜu) = p.precomputed
     (; ᶜSqₗᵖ, ᶜSqᵢᵖ, ᶜSqᵣᵖ, ᶜSqₛᵖ) = p.precomputed
 
-    q_rai = specific(Y.c.ρq_rai, Y.c.ρ)
-    q_sno = specific(Y.c.ρq_sno, Y.c.ρ)
+    q_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ))
+    q_rai = @. lazy(specific(Y.c.ρq_rai, Y.c.ρ))
+    q_sno = @. lazy(specific(Y.c.ρq_sno, Y.c.ρ))
+    q_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ))
+    q_ice = @. lazy(specific(Y.c.ρq_ice, Y.c.ρ))
 
     ᶜSᵖ = p.scratch.ᶜtemp_scalar
     ᶜSᵖ_snow = p.scratch.ᶜtemp_scalar_2
@@ -364,10 +367,10 @@ function set_precipitation_cache!(Y, p, ::Microphysics2Moment, _)
     (; ᶜSqₗᵖ, ᶜSqᵢᵖ, ᶜSqᵣᵖ, ᶜSqₛᵖ) = p.precomputed
     (; ᶜSnₗᵖ, ᶜSnᵣᵖ) = p.precomputed
 
-    q_liq = specific(Y.c.ρq_liq, Y.c.ρ)
-    q_rai = specific(Y.c.ρq_rai, Y.c.ρ)
-    n_liq = specific(Y.c.ρn_liq, Y.c.ρ)
-    n_rai = specific(Y.c.ρn_rai, Y.c.ρ)
+    q_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ))
+    q_rai = @. lazy(specific(Y.c.ρq_rai, Y.c.ρ))
+    n_liq = @. lazy(specific(Y.c.ρn_liq, Y.c.ρ))
+    n_rai = @. lazy(specific(Y.c.ρn_rai, Y.c.ρ))
 
     ᶜSᵖ = p.scratch.ᶜtemp_scalar
     ᶜS₂ᵖ = p.scratch.ᶜtemp_scalar_2
@@ -481,20 +484,24 @@ function set_precipitation_surface_fluxes!(
     sfc_ρ = @. lazy(int_ρ * int_J / sfc_J)
 
     # Constant extrapolation to surface, consistent with simple downwinding
+    ᶜq_rai = @. lazy(specific(Y.c.ρq_rai, Y.c.ρ))
+    ᶜq_sno = @. lazy(specific(Y.c.ρq_sno, Y.c.ρ))
+    ᶜq_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ))
+    ᶜq_ice = @. lazy(specific(Y.c.ρq_ice, Y.c.ρ))
     sfc_qᵣ = Fields.Field(
-        Fields.field_values(Fields.level(specific(Y.c.ρq_rai, Y.c.ρ), 1)),
+        Fields.field_values(Fields.level(Base.materialize(ᶜq_rai), 1)),
         sfc_space,
     )
     sfc_qₛ = Fields.Field(
-        Fields.field_values(Fields.level(specific(Y.c.ρq_sno, Y.c.ρ), 1)),
+        Fields.field_values(Fields.level(Base.materialize(ᶜq_sno), 1)),
         sfc_space,
     )
     sfc_qₗ = Fields.Field(
-        Fields.field_values(Fields.level(specific(Y.c.ρq_liq, Y.c.ρ), 1)),
+        Fields.field_values(Fields.level(Base.materialize(ᶜq_liq), 1)),
         sfc_space,
     )
     sfc_qᵢ = Fields.Field(
-        Fields.field_values(Fields.level(specific(Y.c.ρq_ice, Y.c.ρ), 1)),
+        Fields.field_values(Fields.level(Base.materialize(ᶜq_ice), 1)),
         sfc_space,
     )
     sfc_wᵣ = Fields.Field(Fields.field_values(Fields.level(ᶜwᵣ, 1)), sfc_space)

src/cache/precomputed_quantities.jl

@@ -319,8 +319,9 @@ function set_sgs_ᶠu₃!(w_function, ᶠu₃, Y, turbconv_model)
     return nothing
 end
 
-function add_sgs_ᶜK!(ᶜK, Y, ᶜρa⁰, ᶠu₃⁰, turbconv_model)
-    @. ᶜK += ᶜρa⁰ * ᶜinterp(dot(ᶠu₃⁰ - Yf.u₃, CT3(ᶠu₃⁰ - Yf.u₃))) / 2 / Yc.ρ
+function add_sgs_ᶜK!(ᶜK, Y, ᶜρa⁰_vals, ᶠu₃⁰, turbconv_model)
+    @. ᶜK +=
+        ᶜρa⁰_vals * ᶜinterp(dot(ᶠu₃⁰ - Yf.u₃, CT3(ᶠu₃⁰ - Yf.u₃))) / 2 / Yc.ρ
     for j in 1:n_mass_flux_subdomains(turbconv_model)
         ᶜρaʲ = Y.c.sgsʲs.:($j).ρa
         ᶠu₃ʲ = Y.f.sgsʲs.:($j).u₃
@@ -473,8 +474,6 @@ NVTX.@annotate function set_implicit_precomputed_quantities!(Y, p, t)
         set_prognostic_edmf_precomputed_quantities_draft!(Y, p, ᶠuₕ³, t)
         set_prognostic_edmf_precomputed_quantities_environment!(Y, p, ᶠuₕ³, t)
         set_prognostic_edmf_precomputed_quantities_implicit_closures!(Y, p, t)
-    elseif turbconv_model isa DiagnosticEDMFX
-        set_diagnostic_edmf_precomputed_quantities!(Y, p, t)
     elseif !(isnothing(turbconv_model))
         # Do nothing for other turbconv models for now
     end