rm pressure, h_tot from cache and begin refactoring velocity cache

src/cache/diagnostic_edmf_precomputed_quantities.jl

@@ -93,12 +93,12 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_bottom_bc!(
     FT = eltype(Y)
     n = n_mass_flux_subdomains(turbconv_model)
     (; ᶜΦ) = p.core
-    (; ᶜp, ᶠu³, ᶜh_tot, ᶜK) = p.precomputed
+    (; ᶠu³, ᶜK) = p.precomputed
     (; q_tot) = p.precomputed.ᶜspecific
     (; ustar, obukhov_length, buoyancy_flux, ρ_flux_h_tot, ρ_flux_q_tot) =
         p.precomputed.sfc_conditions
     (; ᶜρaʲs, ᶠu³ʲs, ᶜKʲs, ᶜmseʲs, ᶜq_totʲs, ᶜtsʲs, ᶜρʲs) = p.precomputed
-    (; ᶠu³⁰, ᶜK⁰) = p.precomputed
+    (; ᶠu³⁰, ᶜK⁰, ᶜts) = p.precomputed
 
     (; params) = p
     thermo_params = CAP.thermodynamics_params(params)
@@ -107,11 +107,11 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_bottom_bc!(
     ρ_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(Y, thermo_params, ᶜts)), 1))
     K_int_level = Fields.field_values(Fields.level(ᶜK, 1))
     q_tot_int_level = Fields.field_values(Fields.level(q_tot, 1))
 
-    p_int_level = Fields.field_values(Fields.level(ᶜp, 1))
+    p_int_level = Fields.field_values(Fields.level(Base.materialize(ᶜp(thermo_params, ᶜts)), 1))
     Φ_int_level = Fields.field_values(Fields.level(ᶜΦ, 1))
 
     local_geometry_int_level =
@@ -305,7 +305,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
     (; dt) = p
     dt = float(dt)
     (; ᶜΦ, ᶜgradᵥ_ᶠΦ) = p.core
-    (; ᶜp, ᶠu³, ᶜts, ᶜh_tot, ᶜK) = p.precomputed
+    (; ᶠu³, ᶜts, ᶜK) = p.precomputed
     (; q_tot) = p.precomputed.ᶜspecific
     (;
         ᶜρaʲs,
@@ -352,9 +352,9 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
         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))
+        h_tot_level = Fields.field_values(Fields.level(Base.materialize(ᶜh_tot(Y, thermo_params, ᶜts)), i))
         q_tot_level = Fields.field_values(Fields.level(q_tot, i))
-        p_level = Fields.field_values(Fields.level(ᶜp, i))
+        p_level = Fields.field_values(Fields.level(Base.materialize(ᶜp(thermo_params,ᶜts)), i))
         Φ_level = Fields.field_values(Fields.level(ᶜΦ, i))
         local_geometry_level = Fields.field_values(
             Fields.level(Fields.local_geometry_field(Y.c), i),
@@ -377,10 +377,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))
+        h_tot_prev_level = Fields.field_values(Fields.level(Base.materialize(ᶜh_tot(Y, thermo_params, ᶜts)), i - 1))
         q_tot_prev_level = Fields.field_values(Fields.level(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))
+        p_prev_level = Fields.field_values(Fields.level(Base.materialize(ᶜp(thermo_params, ᶜts)), i - 1))
         z_prev_level = Fields.field_values(Fields.level(ᶜz, i - 1))
         dz_prev_level = Fields.field_values(Fields.level(ᶜdz, i - 1))
 
@@ -960,7 +960,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
     ᶜdz = Fields.Δz_field(axes(Y.c))
     (; params) = p
     (; dt) = p
-    (; ᶜp, ᶜu, ᶜts) = p.precomputed
+    (; ᶜu, ᶜts) = p.precomputed
     (; ustar, obukhov_length) = p.precomputed.sfc_conditions
     (; ᶜtke⁰) = p.precomputed
     (;

src/cache/precomputed_quantities.jl

@@ -19,7 +19,6 @@ The following grid-scale quantities are treated implicitly:
     - `ᶜK`: kinetic energy on cell centers
     - `ᶜts`: thermodynamic state on cell centers
     - `ᶜp`: air pressure on cell centers
-    - `ᶜh_tot`: total enthalpy on cell centers
 If the `turbconv_model` is `PrognosticEDMFX`, there also two SGS versions of
 every quantity except for `ᶜp` (which is shared across all subdomains):
     - `_⁰`: value for the environment
@@ -48,8 +47,6 @@ function implicit_precomputed_quantities(Y, atmos)
         ᶠu = similar(Y.f, CT123{FT}),
         ᶜK = similar(Y.c, FT),
         ᶜts = similar(Y.c, TST),
-        ᶜp = similar(Y.c, FT),
-        ᶜh_tot = similar(Y.c, FT),
     )
     sgs_quantities =
         turbconv_model isa AbstractEDMF ? (; ᶜtke⁰ = similar(Y.c, FT)) : (;)
@@ -455,7 +452,7 @@ quantities are updated.
 NVTX.@annotate function set_implicit_precomputed_quantities!(Y, p, t)
     (; turbconv_model, moisture_model, precip_model) = p.atmos
     (; ᶜΦ) = p.core
-    (; ᶜspecific, ᶜu, ᶠu³, ᶠu, ᶜK, ᶜts, ᶜp, ᶜh_tot) = p.precomputed
+    (; ᶜspecific, ᶜu, ᶠu³, ᶠu, ᶜK, ᶜts) = p.precomputed
     ᶠuₕ³ = p.scratch.ᶠtemp_CT3
     n = n_mass_flux_subdomains(turbconv_model)
     thermo_params = CAP.thermodynamics_params(p.params)
@@ -488,12 +485,6 @@ NVTX.@annotate function set_implicit_precomputed_quantities!(Y, p, t)
         # TODO: We should think more about these increments before we use them.
     end
     @. ᶜts = ts_gs(thermo_args..., Y.c, ᶜK, ᶜΦ, Y.c.ρ)
-    @. ᶜp = TD.air_pressure(thermo_params, ᶜts)
-    @. ᶜh_tot = TD.total_specific_enthalpy(
-        thermo_params,
-        ᶜts,
-        specific(Y.c.ρe_tot, Y.c.ρ),
-    )
 
     if turbconv_model isa PrognosticEDMFX
         set_prognostic_edmf_precomputed_quantities_draft!(Y, p, ᶠuₕ³, t)
@@ -517,7 +508,7 @@ NVTX.@annotate function set_explicit_precomputed_quantities!(Y, p, t)
     (; turbconv_model, moisture_model, precip_model, cloud_model) = p.atmos
     (; vert_diff, call_cloud_diagnostics_per_stage) = p.atmos
     (; ᶜΦ) = p.core
-    (; ᶜu, ᶜts, ᶜp) = p.precomputed
+    (; ᶜu, ᶜts) = p.precomputed
     ᶠuₕ³ = p.scratch.ᶠtemp_CT3 # updated in set_implicit_precomputed_quantities!
     thermo_params = CAP.thermodynamics_params(p.params)
 
@@ -581,7 +572,8 @@ NVTX.@annotate function set_explicit_precomputed_quantities!(Y, p, t)
         @. ᶜK_h = $compute_eddy_diffusivity_coefficient(Y.c.ρ, vert_diff)
     elseif vert_diff isa VerticalDiffusion
         (; ᶜK_h) = p.precomputed
-        @. ᶜK_h = $compute_eddy_diffusivity_coefficient(Y.c.uₕ, ᶜp, vert_diff)
+        ᶜp_lazy = ᶜp(thermo_params, ᶜts)
+        @. ᶜK_h = $compute_eddy_diffusivity_coefficient(Y.c.uₕ, ᶜp_lazy, vert_diff)
     end
 
     # TODO

src/cache/prognostic_edmf_precomputed_quantities.jl

@@ -21,20 +21,23 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!(
     thermo_params = CAP.thermodynamics_params(p.params)
     (; turbconv_model) = p.atmos
     (; ᶜΦ,) = p.core
-    (; ᶜp, ᶜh_tot, ᶜK) = p.precomputed
-    (; ᶜtke⁰, ᶜρa⁰, ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜρ⁰, ᶜmse⁰, ᶜq_tot⁰) =
+    (; ᶜK) = p.precomputed
+    (; ᶜtke⁰, ᶜρa⁰, ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜρ⁰, ᶜmse⁰, ᶜq_tot⁰, ᶜts) =
         p.precomputed
     if p.atmos.moisture_model isa NonEquilMoistModel &&
        p.atmos.precip_model isa Microphysics1Moment
         (; ᶜq_liq⁰, ᶜq_ice⁰, ᶜq_rai⁰, ᶜq_sno⁰) = p.precomputed
     end
 
+    ᶜp_lazy = ᶜp(thermo_params, ᶜts)
+    ᶜh_tot_lazy = ᶜh_tot(Y,thermo_params, ᶜts)
+
     @. ᶜρa⁰ = ρa⁰(Y.c)
     @. ᶜtke⁰ = divide_by_ρa(Y.c.sgs⁰.ρatke, ᶜρa⁰, 0, Y.c.ρ, turbconv_model)
     @. ᶜmse⁰ = divide_by_ρa(
-        Y.c.ρ * (ᶜh_tot - ᶜK) - ρamse⁺(Y.c.sgsʲs),
+        Y.c.ρ * (ᶜh_tot_lazy - ᶜK) - ρamse⁺(Y.c.sgsʲs),
         ᶜρa⁰,
-        Y.c.ρ * (ᶜh_tot - ᶜK),
+        Y.c.ρ * (ᶜh_tot_lazy - ᶜK),
         Y.c.ρ,
         turbconv_model,
     )
@@ -83,12 +86,15 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!(
        p.atmos.precip_model isa Microphysics1Moment
         @. ᶜts⁰ = TD.PhaseNonEquil_phq(
             thermo_params,
-            ᶜp,
+            ᶜp_lazy,
             ᶜmse⁰ - ᶜΦ,
             TD.PhasePartition(ᶜq_tot⁰, ᶜq_liq⁰ + ᶜq_rai⁰, ᶜq_ice⁰ + ᶜq_sno⁰),
         )
     else
-        @. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmse⁰ - ᶜΦ, ᶜq_tot⁰)
+        @. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, 
+                                    ᶜp_lazy,
+                                    ᶜmse⁰ - ᶜΦ, 
+                                    ᶜq_tot⁰)
     end
     @. ᶜρ⁰ = TD.air_density(thermo_params, ᶜts⁰)
     return nothing
@@ -112,7 +118,9 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_draft!(
     thermo_params = CAP.thermodynamics_params(p.params)
 
     (; ᶜΦ,) = p.core
-    (; ᶜp, ᶜuʲs, ᶠu³ʲs, ᶜKʲs, ᶠKᵥʲs, ᶜtsʲs, ᶜρʲs) = p.precomputed
+    (; ᶜuʲs, ᶠu³ʲs, ᶜKʲs, ᶠKᵥʲs, ᶜtsʲs, ᶜρʲs, ᶜts) = p.precomputed
+
+    ᶜp_lazy = ᶜp(thermo_params, ᶜts)
 
     for j in 1:n
         ᶜuʲ = ᶜuʲs.:($j)
@@ -138,7 +146,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_draft!(
            p.atmos.precip_model isa Microphysics1Moment
             @. ᶜtsʲ = TD.PhaseNonEquil_phq(
                 thermo_params,
-                ᶜp,
+                ᶜp_lazy,
                 ᶜmseʲ - ᶜΦ,
                 TD.PhasePartition(
                     ᶜq_totʲ,
@@ -147,7 +155,10 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_draft!(
                 ),
             )
         else
-            @. ᶜtsʲ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmseʲ - ᶜΦ, ᶜq_totʲ)
+            @. ᶜtsʲ = TD.PhaseEquil_phq(thermo_params, 
+                                        ᶜp_lazy,
+                                        ᶜmseʲ - ᶜΦ, 
+                                        ᶜq_totʲ)
         end
         @. ᶜρʲ = TD.air_density(thermo_params, ᶜtsʲ)
     end
@@ -173,9 +184,12 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_bottom_bc!(
     turbconv_params = CAP.turbconv_params(p.params)
 
     (; ᶜΦ,) = p.core
-    (; ᶜspecific, ᶜp, ᶜh_tot, ᶜK, ᶜtsʲs, ᶜρʲs) = p.precomputed
+    (; ᶜspecific, ᶜK, ᶜtsʲs, ᶜρʲs, ᶜts) = p.precomputed
     (; ustar, obukhov_length, buoyancy_flux) = p.precomputed.sfc_conditions
 
+    ᶜp_lazy = ᶜp(thermo_params, ᶜts)
+    ᶜh_tot_lazy = ᶜh_tot(Y, thermo_params, ᶜts)
+
     for j in 1:n
         ᶜtsʲ = ᶜtsʲs.:($j)
         ᶜmseʲ = Y.c.sgsʲs.:($j).mse
@@ -197,7 +211,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_bottom_bc!(
             Fields.level(Fields.coordinate_field(Y.f).z, Fields.half),
         )
         ᶜρ_int_val = Fields.field_values(Fields.level(Y.c.ρ, 1))
-        ᶜp_int_val = Fields.field_values(Fields.level(ᶜp, 1))
+        ᶜp_int_val = Fields.field_values(Fields.level(Base.materialize(ᶜp_lazy), 1))
         (; ρ_flux_h_tot, ρ_flux_q_tot, ustar, obukhov_length) =
             p.precomputed.sfc_conditions
 
@@ -217,7 +231,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_bottom_bc!(
         # TODO: replace this with the actual surface area fraction when
         # using prognostic surface area
         @. ᶜaʲ_int_val = FT(turbconv_params.surface_area)
-        ᶜh_tot_int_val = Fields.field_values(Fields.level(ᶜh_tot, 1))
+        ᶜh_tot_int_val = Fields.field_values(Fields.level(Base.materialize(ᶜh_tot_lazy), 1))
         ᶜK_int_val = Fields.field_values(Fields.level(ᶜK, 1))
         ᶜmseʲ_int_val = Fields.field_values(Fields.level(ᶜmseʲ, 1))
         @. ᶜmseʲ_int_val = sgs_scalar_first_interior_bc(
@@ -367,7 +381,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
     FT = eltype(params)
     n = n_mass_flux_subdomains(turbconv_model)
 
-    (; ᶜtke⁰, ᶜu, ᶜp, ᶜρa⁰, ᶠu³⁰, ᶜts⁰, ᶜq_tot⁰) = p.precomputed
+    (; ᶜtke⁰, ᶜu, ᶜρa⁰, ᶠu³⁰, ᶜts⁰, ᶜq_tot⁰) = p.precomputed
     (;
         ᶜmixing_length_tuple,
         ᶜmixing_length,
@@ -386,6 +400,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
         ᶜdetrʲs,
         ᶜturb_entrʲs,
         ᶠnh_pressure₃_buoyʲs,
+        ᶜts,
     ) = p.precomputed
     (; ustar, obukhov_length) = p.precomputed.sfc_conditions
 
@@ -398,14 +413,18 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
     ᶜvert_div = p.scratch.ᶜtemp_scalar
     ᶜmassflux_vert_div = p.scratch.ᶜtemp_scalar_2
     ᶜw_vert_div = p.scratch.ᶜtemp_scalar_3
+
+    ᶜp_lazy = ᶜp(thermo_params, ᶜts)
+    ᶜh_tot_lazy = ᶜp(thermo_params, ᶜts)
+
     for j in 1:n
         # entrainment/detrainment
         @. ᶜentrʲs.:($$j) = entrainment(
             thermo_params,
             turbconv_params,
             ᶜz,
             z_sfc,
-            ᶜp,
+            ᶜp_lazy,
             Y.c.ρ,
             draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)),
             get_physical_w(ᶜuʲs.:($$j), ᶜlg),
@@ -441,7 +460,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
             turbconv_params,
             ᶜz,
             z_sfc,
-            ᶜp,
+            ᶜp_lazy,
             Y.c.ρ,
             Y.c.sgsʲs.:($$j).ρa,
             draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)),

src/diagnostics/core_diagnostics.jl

@@ -166,9 +166,9 @@ add_diagnostic_variable!(
     compute! = (out, state, cache, time) -> begin
         thermo_params = CAP.thermodynamics_params(cache.params)
         if isnothing(out)
-            return TD.specific_enthalpy.(thermo_params, cache.precomputed.ᶜts)
+            return TD.air_pressure.(thermo_params, cache.precomputed.ᶜts)
         else
-            out .= TD.specific_enthalpy.(thermo_params, cache.precomputed.ᶜts)
+            out .= TD.air_pressure.(thermo_params, cache.precomputed.ᶜts)
         end
     end,
 )
@@ -181,10 +181,11 @@ add_diagnostic_variable!(
     long_name = "Pressure at Model Full-Levels",
     units = "Pa",
     compute! = (out, state, cache, time) -> begin
+    thermo_params = CAP.thermodynamics_params(cache.params)
         if isnothing(out)
-            return copy(cache.precomputed.ᶜp)
+            return copy(TD.air_pressure.(thermo_params, cache.precomputed.ᶜts))
         else
-            out .= cache.precomputed.ᶜp
+            out .= TD.air_pressure.(thermo_params, cache.precomputed.ᶜts)
         end
     end,
 )
@@ -1434,7 +1435,7 @@ function compute_cape!(out, state, cache, time)
         lazy.(
             TD.PhaseEquil_pθq.(
                 thermo_params,
-                cache.precomputed.ᶜp,
+                TD.air_pressure.(thermo_params, cache.precomputed.ᶜts),
                 surface_θ_liq_ice,
                 surface_q,
             ),

src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl

@@ -141,6 +141,7 @@ function non_orographic_gravity_wave_compute_tendency!(
     #unpack
     ᶜT = p.scratch.ᶜtemp_scalar
     (; ᶜts) = p.precomputed
+    thermo_params = CAP.thermodynamics_params(p.params)
     (; params) = p
     (;
         ᶜdTdz,
@@ -159,7 +160,6 @@ function non_orographic_gravity_wave_compute_tendency!(
     ᶜz = Fields.coordinate_field(Y.c).z
     FT = Spaces.undertype(axes(Y.c))
     # parameters
-    thermo_params = CAP.thermodynamics_params(params)
     grav = CAP.grav(params)
 
     # compute buoyancy frequency
@@ -205,12 +205,11 @@ function non_orographic_gravity_wave_compute_tendency!(
         fill!(damp_level, Spaces.nlevels(axes(ᶜz)))
 
     elseif issphere(axes(Y.c))
-        (; ᶜp) = p.precomputed
         (; gw_source_pressure, gw_damp_pressure, source_p_ρ_z_u_v_level) =
             p.non_orographic_gravity_wave
         # source level: the index of the highest level whose pressure is higher than source pressure
 
-        input = Base.Broadcast.broadcasted(tuple, ᶜp, ᶜρ, ᶜz, ᶜu, ᶜv, ᶜlevel)
+        input = Base.Broadcast.broadcasted(tuple, Base.materialize(ᶜp(thermo_params, ᶜts)), ᶜρ, ᶜz, ᶜu, ᶜv, ᶜlevel)
         Operators.column_reduce!(
             source_p_ρ_z_u_v_level,
             input,
@@ -231,7 +230,7 @@ function non_orographic_gravity_wave_compute_tendency!(
 
         # damp level: the index of the lowest level whose pressure is lower than the damp pressure
 
-        input = Base.Broadcast.broadcasted(tuple, ᶜlevel, ᶜp)
+        input = Base.Broadcast.broadcasted(tuple, ᶜlevel, Base.materialize(ᶜp(thermo_params, ᶜts)))
         Operators.column_reduce!(
             damp_level,
             input;