Merge pull request #3845 from CliMA/ck/slim_cache

charleskawczynski · web-flow · commit 4607ec2a9b49 · 2025-06-12T16:33:09.000-04:00
Reduce dependency on specific cache
diff --git a/src/cache/cloud_fraction.jl b/src/cache/cloud_fraction.jl
@@ -64,15 +64,16 @@ NVTX.@annotate function set_cloud_fraction!(
     else
         @. cloud_diagnostics_tuple = make_named_tuple(
             ifelse(
-                p.precomputed.ᶜspecific.q_liq + p.precomputed.ᶜspecific.q_ice > 0,
+                specific(Y.c.ρq_liq, Y.c.ρ) + specific(Y.c.ρq_ice, Y.c.ρ) > 0,
                 1,
                 0,
             ),
-            p.precomputed.ᶜspecific.q_liq,
-            p.precomputed.ᶜspecific.q_ice,
+            specific(Y.c.ρq_liq, Y.c.ρ),
+            specific(Y.c.ρq_ice, Y.c.ρ),
         )
     end
 end
+
 NVTX.@annotate function set_cloud_fraction!(
     Y,
     p,
diff --git a/src/cache/precipitation_precomputed_quantities.jl b/src/cache/precipitation_precomputed_quantities.jl
@@ -105,7 +105,6 @@ function set_precipitation_velocities!(
     precip_model::Microphysics2Moment,
 )
     (; ᶜwₗ, ᶜwᵢ, ᶜwᵣ, ᶜwₛ, ᶜwnₗ, ᶜwnᵣ, ᶜwₜqₜ, ᶜwₕhₜ, ᶜts, ᶜu) = p.precomputed
-    (; q_liq, q_ice, q_rai, q_sno) = p.precomputed.ᶜspecific
     (; ᶜΦ) = p.core
 
     cm1c = CAP.microphysics_cloud_params(p.params)
@@ -116,23 +115,53 @@ function set_precipitation_velocities!(
     # compute the precipitation terminal velocity [m/s]
     # TODO sedimentation of snow is based on the 1M scheme
     @. ᶜwnᵣ = getindex(
-        CM2.rain_terminal_velocity(cm2p.sb, cm2p.tv, q_rai, Y.c.ρ, Y.c.ρn_rai),
+        CM2.rain_terminal_velocity(
+            cm2p.sb,
+            cm2p.tv,
+            specific(Y.c.ρq_rai, Y.c.ρ),
+            Y.c.ρ,
+            Y.c.ρn_rai,
+        ),
         1,
     )
     @. ᶜwᵣ = getindex(
-        CM2.rain_terminal_velocity(cm2p.sb, cm2p.tv, q_rai, Y.c.ρ, Y.c.ρn_rai),
+        CM2.rain_terminal_velocity(
+            cm2p.sb,
+            cm2p.tv,
+            specific(Y.c.ρq_rai, Y.c.ρ),
+            Y.c.ρ,
+            Y.c.ρn_rai,
+        ),
         2,
     )
-    @. ᶜwₛ = CM1.terminal_velocity(cm1p.ps, cm1p.tv.snow, Y.c.ρ, q_sno)
+    @. ᶜwₛ = CM1.terminal_velocity(
+        cm1p.ps,
+        cm1p.tv.snow,
+        Y.c.ρ,
+        specific(Y.c.ρq_sno, Y.c.ρ),
+    )
     # compute sedimentation velocity for cloud condensate [m/s]
     # TODO sedimentation velocities of cloud condensates are based
     # on the 1M scheme. Sedimentation velocity of cloud number concentration
     # is equal to that of the mass.
-    @. ᶜwnₗ =
-        CMNe.terminal_velocity(cm1c.liquid, cm1c.Ch2022.rain, Y.c.ρ, q_liq)
-    @. ᶜwₗ = CMNe.terminal_velocity(cm1c.liquid, cm1c.Ch2022.rain, Y.c.ρ, q_liq)
-    @. ᶜwᵢ =
-        CMNe.terminal_velocity(cm1c.ice, cm1c.Ch2022.small_ice, Y.c.ρ, q_ice)
+    @. ᶜwnₗ = CMNe.terminal_velocity(
+        cm1c.liquid,
+        cm1c.Ch2022.rain,
+        Y.c.ρ,
+        specific(Y.c.ρq_liq, Y.c.ρ),
+    )
+    @. ᶜwₗ = CMNe.terminal_velocity(
+        cm1c.liquid,
+        cm1c.Ch2022.rain,
+        Y.c.ρ,
+        specific(Y.c.ρq_liq, Y.c.ρ),
+    )
+    @. ᶜwᵢ = CMNe.terminal_velocity(
+        cm1c.ice,
+        cm1c.Ch2022.small_ice,
+        Y.c.ρ,
+        specific(Y.c.ρq_ice, Y.c.ρ),
+    )
 
     # compute their contributions to energy and total water advection
     @. ᶜwₜqₜ =
diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl
@@ -398,36 +398,34 @@ function thermo_state(
     return get_ts(ρ, p, θ, e_int, q_tot, q_pt)
 end
 
-function thermo_vars(moisture_model, precip_model, specific, K, Φ)
-    energy_var = (; e_int = specific.e_tot - K - Φ)
+function thermo_vars(moisture_model, precip_model, ᶜY, K, Φ)
+    energy_var = (; e_int = specific(ᶜY.ρe_tot, ᶜY.ρ) - K - Φ)
     moisture_var = if moisture_model isa DryModel
         (;)
     elseif moisture_model isa EquilMoistModel
-        (; specific.q_tot)
+        (; q_tot = specific(ᶜY.ρq_tot, ᶜY.ρ))
     elseif moisture_model isa NonEquilMoistModel
-        q_pt_args = (
-            specific.q_tot,
-            specific.q_liq + specific.q_rai,
-            specific.q_ice + specific.q_sno,
+        q_pt_args = (;
+            q_tot = specific(ᶜY.ρq_tot, ᶜY.ρ),
+            q_liq = specific(ᶜY.ρq_liq, ᶜY.ρ) + specific(ᶜY.ρq_rai, ᶜY.ρ),
+            q_ice = specific(ᶜY.ρq_ice, ᶜY.ρ) + specific(ᶜY.ρq_sno, ᶜY.ρ),
         )
         (; q_pt = TD.PhasePartition(q_pt_args...))
     end
     return (; energy_var..., moisture_var...)
 end
 
-ts_gs(thermo_params, moisture_model, precip_model, specific, K, Φ, ρ) =
-    thermo_state(
-        thermo_params;
-        thermo_vars(moisture_model, precip_model, specific, K, Φ)...,
-        ρ,
-    )
+ts_gs(thermo_params, moisture_model, precip_model, ᶜY, K, Φ, ρ) = thermo_state(
+    thermo_params;
+    thermo_vars(moisture_model, precip_model, ᶜY, K, Φ)...,
+    ρ,
+)
 
-ts_sgs(thermo_params, moisture_model, precip_model, specific, K, Φ, p) =
-    thermo_state(
-        thermo_params;
-        thermo_vars(moisture_model, precip_model, specific, K, Φ)...,
-        p,
-    )
+ts_sgs(thermo_params, moisture_model, precip_model, ᶜY, K, Φ, p) = thermo_state(
+    thermo_params;
+    thermo_vars(moisture_model, precip_model, ᶜY, K, Φ)...,
+    p,
+)
 
 function eddy_diffusivity_coefficient_H(D₀, H, z_sfc, z)
     return D₀ * exp(-(z - z_sfc) / H)
@@ -489,9 +487,13 @@ NVTX.@annotate function set_implicit_precomputed_quantities!(Y, p, t)
         # @. ᶜK += Y.c.sgs⁰.ρatke / Y.c.ρ
         # TODO: We should think more about these increments before we use them.
     end
-    @. ᶜts = ts_gs(thermo_args..., ᶜspecific, ᶜK, ᶜΦ, Y.c.ρ)
+    @. ᶜ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.e_tot)
+    @. ᶜ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)
diff --git a/src/parameterized_tendencies/microphysics/cloud_condensate.jl b/src/parameterized_tendencies/microphysics/cloud_condensate.jl
@@ -29,13 +29,16 @@ function cloud_condensate_tendency!(
 )
     (; ᶜts) = p.precomputed
     (; params, dt) = p
-    (; q_rai, q_sno) = p.precomputed.ᶜspecific
     FT = eltype(params)
     thp = CAP.thermodynamics_params(params)
     cmc = CAP.microphysics_cloud_params(params)
 
-    @. Yₜ.c.ρq_liq += Y.c.ρ * cloud_sources(cmc.liquid, thp, ᶜts, q_rai, dt)
-    @. Yₜ.c.ρq_ice += Y.c.ρ * cloud_sources(cmc.ice, thp, ᶜts, q_sno, dt)
+    @. Yₜ.c.ρq_liq +=
+        Y.c.ρ *
+        cloud_sources(cmc.liquid, thp, ᶜts, specific(Y.c.ρq_rai, Y.c.ρ), dt)
+    @. Yₜ.c.ρq_ice +=
+        Y.c.ρ *
+        cloud_sources(cmc.ice, thp, ᶜts, specific(Y.c.ρq_sno, Y.c.ρ), dt)
 end
 
 function cloud_condensate_tendency!(
@@ -47,21 +50,24 @@ function cloud_condensate_tendency!(
 )
     (; ᶜts) = p.precomputed
     (; params, dt) = p
-    (; q_rai, q_sno, n_liq) = p.precomputed.ᶜspecific
     FT = eltype(params)
     thp = CAP.thermodynamics_params(params)
     cmc = CAP.microphysics_cloud_params(params)
 
-    @. Yₜ.c.ρq_liq += Y.c.ρ * cloud_sources(cmc.liquid, thp, ᶜts, q_rai, dt)
-    @. Yₜ.c.ρq_ice += Y.c.ρ * cloud_sources(cmc.ice, thp, ᶜts, q_sno, dt)
+    @. Yₜ.c.ρq_liq +=
+        Y.c.ρ *
+        cloud_sources(cmc.liquid, thp, ᶜts, specific(Y.c.ρq_rai, Y.c.ρ), dt)
+    @. Yₜ.c.ρq_ice +=
+        Y.c.ρ *
+        cloud_sources(cmc.ice, thp, ᶜts, specific(Y.c.ρq_sno, Y.c.ρ), dt)
 
     @. Yₜ.c.ρn_liq +=
         Y.c.ρ * aerosol_activation_sources(
             cmc.liquid,
             thp,
             ᶜts,
-            q_rai,
-            n_liq,
+            specific(Y.c.ρq_rai, Y.c.ρ),
+            specific(Y.c.ρn_liq, Y.c.ρ),
             cmc.N_cloud_liquid_droplets,
             dt,
         )
diff --git a/src/prognostic_equations/edmfx_sgs_flux.jl b/src/prognostic_equations/edmfx_sgs_flux.jl
@@ -90,7 +90,7 @@ function edmfx_sgs_mass_flux_tendency!(
             for j in 1:n
                 @. ᶠu³_diff = ᶠu³ʲs.:($$j) - ᶠu³
                 @. ᶜa_scalar =
-                    (Y.c.sgsʲs.:($$j).q_tot - ᶜspecific.q_tot) *
+                    (Y.c.sgsʲs.:($$j).q_tot - specific(Y.c.ρq_tot, Y.c.ρ)) *
                     draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j))
                 vtt = vertical_transport(
                     ᶜρʲs.:($j),
@@ -103,7 +103,8 @@ function edmfx_sgs_mass_flux_tendency!(
             end
             # Add the environment fluxes
             @. ᶠu³_diff = ᶠu³⁰ - ᶠu³
-            @. ᶜa_scalar = (ᶜq_tot⁰ - ᶜspecific.q_tot) * draft_area(ᶜρa⁰, ᶜρ⁰)
+            @. ᶜa_scalar =
+                (ᶜq_tot⁰ - specific(Y.c.ρq_tot, Y.c.ρ)) * draft_area(ᶜρa⁰, ᶜρ⁰)
             vtt = vertical_transport(
                 ᶜρ⁰,
                 ᶠu³_diff,
@@ -123,7 +124,7 @@ function edmfx_sgs_mass_flux_tendency!(
                 @. ᶠu³_diff = ᶠu³ʲs.:($$j) - ᶠu³
 
                 @. ᶜa_scalar =
-                    (Y.c.sgsʲs.:($$j).q_liq - ᶜspecific.q_liq) *
+                    (Y.c.sgsʲs.:($$j).q_liq - specific(Y.c.ρq_liq, Y.c.ρ)) *
                     draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j))
                 vtt = vertical_transport(
                     ᶜρʲs.:($j),
@@ -135,7 +136,7 @@ function edmfx_sgs_mass_flux_tendency!(
                 @. Yₜ.c.ρq_liq += vtt
 
                 @. ᶜa_scalar =
-                    (Y.c.sgsʲs.:($$j).q_ice - ᶜspecific.q_ice) *
+                    (Y.c.sgsʲs.:($$j).q_ice - specific(Y.c.ρq_ice, Y.c.ρ)) *
                     draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j))
                 vtt = vertical_transport(
                     ᶜρʲs.:($j),
@@ -147,7 +148,7 @@ function edmfx_sgs_mass_flux_tendency!(
                 @. Yₜ.c.ρq_ice += vtt
 
                 @. ᶜa_scalar =
-                    (Y.c.sgsʲs.:($$j).q_rai - ᶜspecific.q_rai) *
+                    (Y.c.sgsʲs.:($$j).q_rai - specific(Y.c.ρq_rai, Y.c.ρ)) *
                     draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j))
                 vtt = vertical_transport(
                     ᶜρʲs.:($j),
@@ -159,7 +160,7 @@ function edmfx_sgs_mass_flux_tendency!(
                 @. Yₜ.c.ρq_rai += vtt
 
                 @. ᶜa_scalar =
-                    (Y.c.sgsʲs.:($$j).q_sno - ᶜspecific.q_sno) *
+                    (Y.c.sgsʲs.:($$j).q_sno - specific(Y.c.ρq_sno, Y.c.ρ)) *
                     draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j))
                 vtt = vertical_transport(
                     ᶜρʲs.:($j),
@@ -172,7 +173,8 @@ function edmfx_sgs_mass_flux_tendency!(
             end
             @. ᶠu³_diff = ᶠu³⁰ - ᶠu³
 
-            @. ᶜa_scalar = (ᶜq_liq⁰ - ᶜspecific.q_liq) * draft_area(ᶜρa⁰, ᶜρ⁰)
+            @. ᶜa_scalar =
+                (ᶜq_liq⁰ - specific(Y.c.ρq_liq, Y.c.ρ)) * draft_area(ᶜρa⁰, ᶜρ⁰)
             vtt = vertical_transport(
                 ᶜρ⁰,
                 ᶠu³_diff,
@@ -182,7 +184,8 @@ function edmfx_sgs_mass_flux_tendency!(
             )
             @. Yₜ.c.ρq_liq += vtt
 
-            @. ᶜa_scalar = (ᶜq_ice⁰ - ᶜspecific.q_ice) * draft_area(ᶜρa⁰, ᶜρ⁰)
+            @. ᶜa_scalar =
+                (ᶜq_ice⁰ - specific(Y.c.ρq_ice, Y.c.ρ)) * draft_area(ᶜρa⁰, ᶜρ⁰)
             vtt = vertical_transport(
                 ᶜρ⁰,
                 ᶠu³_diff,
@@ -192,7 +195,8 @@ function edmfx_sgs_mass_flux_tendency!(
             )
             @. Yₜ.c.ρq_ice += vtt
 
-            @. ᶜa_scalar = (ᶜq_rai⁰ - ᶜspecific.q_rai) * draft_area(ᶜρa⁰, ᶜρ⁰)
+            @. ᶜa_scalar =
+                (ᶜq_rai⁰ - specific(Y.c.ρq_rai, Y.c.ρ)) * draft_area(ᶜρa⁰, ᶜρ⁰)
             vtt = vertical_transport(
                 ᶜρ⁰,
                 ᶠu³_diff,
@@ -202,7 +206,8 @@ function edmfx_sgs_mass_flux_tendency!(
             )
             @. Yₜ.c.ρq_rai += vtt
 
-            @. ᶜa_scalar = (ᶜq_sno⁰ - ᶜspecific.q_sno) * draft_area(ᶜρa⁰, ᶜρ⁰)
+            @. ᶜa_scalar =
+                (ᶜq_sno⁰ - specific(Y.c.ρq_sno, Y.c.ρ)) * draft_area(ᶜρa⁰, ᶜρ⁰)
             vtt = vertical_transport(
                 ᶜρ⁰,
                 ᶠu³_diff,
@@ -271,11 +276,11 @@ function edmfx_sgs_mass_flux_tendency!(
             for j in 1:n
                 @. ᶠu³_diff = ᶠu³ʲs.:($$j) - ᶠu³
                 # @. ᶜa_scalar =
-                #     (ᶜq_totʲs.:($$j) - ᶜspecific.q_tot) *
+                #     (ᶜq_totʲs.:($$j) - specific(Y.c.ρq_tot, Y.c.ρ)) *
                 #     draft_area(ᶜρaʲs.:($$j), ᶜρʲs.:($$j))
                 # TODO: remove this filter when mass flux is treated implicitly
                 @. ᶜa_scalar =
-                    (ᶜq_totʲs.:($$j) - ᶜspecific.q_tot) * min(
+                    (ᶜq_totʲs.:($$j) - specific(Y.c.ρq_tot, Y.c.ρ)) * min(
                         min(draft_area(ᶜρaʲs.:($$j), ᶜρʲs.:($$j)), a_max),
                         FT(0.02) / max(
                             Geometry.WVector(
@@ -491,7 +496,7 @@ function edmfx_sgs_diffusive_flux_tendency!(
                 bottom = Operators.SetValue(C3(FT(0))),
             )
             @. ᶜρχₜ_diffusion =
-                ᶜdivᵥ_ρq_tot(-(ᶠρaK_h * ᶠgradᵥ(ᶜspecific.q_tot)))
+                ᶜdivᵥ_ρq_tot(-(ᶠρaK_h * ᶠgradᵥ(specific(Y.c.ρq_tot, Y.c.ρ))))
             @. Yₜ.c.ρq_tot -= ᶜρχₜ_diffusion
             @. Yₜ.c.ρ -= ᶜρχₜ_diffusion
         end
diff --git a/src/prognostic_equations/forcing/external_forcing.jl b/src/prognostic_equations/forcing/external_forcing.jl
@@ -363,7 +363,8 @@ function external_forcing_tendency!(
     ᶜdqtdt_nudging = p.scratch.ᶜtemp_scalar_2
     @. ᶜdTdt_nudging =
         -(TD.air_temperature(thermo_params, ᶜts) - ᶜT_nudge) * ᶜinv_τ_scalar
-    @. ᶜdqtdt_nudging = -(ᶜspecific.q_tot - ᶜqt_nudge) * ᶜinv_τ_scalar
+    @. ᶜdqtdt_nudging =
+        -(specific(Y.c.ρq_tot, Y.c.ρ) - ᶜqt_nudge) * ᶜinv_τ_scalar
 
     ᶜdTdt_sum = p.scratch.ᶜtemp_scalar
     ᶜdqtdt_sum = p.scratch.ᶜtemp_scalar_2
@@ -572,7 +573,7 @@ function external_forcing_tendency!(Yₜ, Y, p, t, ::ISDACForcing)
     FT = Spaces.undertype(axes(Y.c))
     (; params) = p
     thermo_params = CAP.thermodynamics_params(params)
-    (; ᶜspecific, ᶜts, ᶜh_tot, ᶜp) = p.precomputed
+    (; ᶜts, ᶜh_tot, ᶜp) = p.precomputed
 
     ᶜinv_τ_scalar = APL.ISDAC_inv_τ_scalar(FT)  # s⁻¹
     ᶜinv_τ_wind = APL.ISDAC_inv_τ_wind(FT)  # s⁻¹
@@ -601,7 +602,8 @@ function external_forcing_tendency!(Yₜ, Y, p, t, ::ISDACForcing)
     @. ᶜdTdt_nudging =
         -(TD.air_temperature(thermo_params, ᶜts) - ta_ISDAC(ᶜp, ᶜz)) *
         ᶜinv_τ_scalar(ᶜz)
-    @. ᶜdqtdt_nudging = -(ᶜspecific.q_tot - q_tot(ᶜz)) * ᶜinv_τ_scalar(ᶜz)
+    @. ᶜdqtdt_nudging =
+        -(specific(Y.c.ρq_tot, Y.c.ρ) - q_tot(ᶜz)) * ᶜinv_τ_scalar(ᶜz)
 
     T_0 = TD.Parameters.T_0(thermo_params)
     Lv_0 = TD.Parameters.LH_v0(thermo_params)
diff --git a/src/prognostic_equations/hyperdiffusion.jl b/src/prognostic_equations/hyperdiffusion.jl
@@ -96,7 +96,8 @@ NVTX.@annotate function prep_hyperdiffusion_tendency!(Yₜ, Y, p, t)
         C123(wgradₕ(divₕ(p.precomputed.ᶜu))) -
         C123(wcurlₕ(C123(curlₕ(p.precomputed.ᶜu))))
 
-    @. ᶜ∇²specific_energy = wdivₕ(gradₕ(ᶜspecific.e_tot + ᶜp / Y.c.ρ))
+    @. ᶜ∇²specific_energy =
+        wdivₕ(gradₕ(specific(Y.c.ρe_tot, Y.c.ρ) + ᶜp / Y.c.ρ))
 
     if diffuse_tke
         (; ᶜtke⁰) = p.precomputed
diff --git a/src/prognostic_equations/implicit/implicit_tendency.jl b/src/prognostic_equations/implicit/implicit_tendency.jl
diff --git a/src/prognostic_equations/implicit/manual_sparse_jacobian.jl b/src/prognostic_equations/implicit/manual_sparse_jacobian.jl
diff --git a/src/utils/variable_manipulations.jl b/src/utils/variable_manipulations.jl
