Reduce cache: lazy velocity computations

src/cache/diagnostic_edmf_precomputed_quantities.jl

@@ -93,20 +93,24 @@ 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
+    (; ᶜp, ᶜh_tot, ᶜ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
+    (; ᶜK⁰) = p.precomputed
+
 
     (; params) = p
     thermo_params = CAP.thermodynamics_params(params)
     turbconv_params = CAP.turbconv_params(params)
 
+    ᶠu³ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃)
+    ᶠu³⁰ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃)
+
     ρ_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))
+    u³_int_halflevel = Fields.field_values(Fields.level(Base.materialize(ᶠu³), half))
     h_tot_int_level = Fields.field_values(Fields.level(ᶜ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))
@@ -194,7 +198,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_bottom_bc!(
 
     ρaʲs_int_level = Fields.field_values(Fields.level(ᶜρaʲs, 1))
     u³ʲs_int_halflevel = Fields.field_values(Fields.level(ᶠu³ʲs, half))
-    u³⁰_int_halflevel = Fields.field_values(Fields.level(ᶠu³⁰, half))
+    u³⁰_int_halflevel = Fields.field_values(Fields.level(Base.materialize(ᶠu³⁰), half))
     K⁰_int_level = Fields.field_values(Fields.level(ᶜK⁰, 1))
     set_diagnostic_edmfx_env_quantities_level!(
         ρ_int_level,
@@ -305,7 +309,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
+    (; ᶜp, ᶜts, ᶜh_tot, ᶜK) = p.precomputed
     (; q_tot) = p.precomputed.ᶜspecific
     (;
         ᶜρaʲs,
@@ -321,7 +325,7 @@ 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
+    (; ᶜK⁰, ᶜtke⁰) = p.precomputed
 
     if precip_model isa Microphysics1Moment
         ᶜq_liqʲs = p.precomputed.ᶜq_liqʲs
@@ -343,14 +347,17 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
     ᶜ∇Φ₃ = p.scratch.ᶜtemp_C3
     @. ᶜ∇Φ₃ = ᶜgradᵥ(ᶠΦ)
 
+    ᶠu³ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃)
+    ᶠu³⁰ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃)
+
     z_sfc_halflevel =
         Fields.field_values(Fields.level(Fields.coordinate_field(Y.f).z, half))
 
     # integral
     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))
+        u³_halflevel = Fields.field_values(Fields.level(Base.materialize(ᶠ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))
@@ -372,9 +379,9 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
         ∇Φ₃_prev_level = Fields.field_values(Fields.level(ᶜ∇Φ₃, i - 1))
         ∇Φ₃_data_prev_level = ∇Φ₃_prev_level.components.data.:1
         ρ_prev_level = Fields.field_values(Fields.level(Y.c.ρ, i - 1))
-        u³_prev_halflevel = Fields.field_values(Fields.level(ᶠu³, i - 1 - half))
+        u³_prev_halflevel = Fields.field_values(Fields.level(Base.materialize(ᶠu³), i - 1 - half))
         u³⁰_prev_halflevel =
-            Fields.field_values(Fields.level(ᶠu³⁰, i - 1 - half))
+        Fields.field_values(Fields.level(Base.materialize(ᶠ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))
@@ -853,7 +860,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
         end
         ρaʲs_level = Fields.field_values(Fields.level(ᶜρaʲs, i))
         u³ʲs_halflevel = Fields.field_values(Fields.level(ᶠu³ʲs, i - half))
-        u³⁰_halflevel = Fields.field_values(Fields.level(ᶠu³⁰, i - half))
+        u³⁰_halflevel = Fields.field_values(Fields.level(Base.materialize(ᶠu³⁰), i - half))
         K⁰_level = Fields.field_values(Fields.level(ᶜK⁰, i))
         set_diagnostic_edmfx_env_quantities_level!(
             ρ_level,
@@ -883,13 +890,14 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_top_bc!(
 )
     n = n_mass_flux_subdomains(p.atmos.turbconv_model)
     (; ᶜentrʲs, ᶜdetrʲs, ᶜturb_entrʲs) = p.precomputed
-    (; ᶠu³⁰, ᶠu³ʲs, ᶜuʲs, ᶠnh_pressure³_buoyʲs, ᶠnh_pressure³_dragʲs) =
+    (; ᶠu³ʲs, ᶜuʲs, ᶠnh_pressure³_buoyʲs, ᶠnh_pressure³_dragʲs) =
         p.precomputed
     (; precip_model) = p.atmos
 
     # set values for the top level
     i_top = Spaces.nlevels(axes(Y.c))
-    u³⁰_halflevel = Fields.field_values(Fields.level(ᶠu³⁰, i_top + half))
+    ᶠu³⁰ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃)
+    u³⁰_halflevel = Fields.field_values(Fields.level(Base.materialize(ᶠu³⁰), i_top + half))
     @. u³⁰_halflevel = CT3(0)
 
     for j in 1:n
@@ -960,7 +968,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
+    (; ᶜp, ᶜts) = p.precomputed
     (; ustar, obukhov_length) = p.precomputed.sfc_conditions
     (; ᶜtke⁰) = p.precomputed
     (;
@@ -974,13 +982,16 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
     thermo_params = CAP.thermodynamics_params(params)
     ᶜlg = Fields.local_geometry_field(Y.c)
 
+    ᶜu = ᶜu_lazy(Y.c.uₕ, Y.f.u₃)
+    ᶠu³⁰ = ᶠu³_lazy(Y.c.uₕ,Y.c.ρ, Y.f.u₃)
+
     if p.atmos.turbconv_model isa DiagnosticEDMFX
-        (; ᶜρaʲs, ᶠu³ʲs, ᶜdetrʲs, ᶠu³⁰, ᶜu⁰) = p.precomputed
+        (; ᶜρaʲs, ᶠu³ʲs, ᶜdetrʲs) = p.precomputed
     elseif p.atmos.turbconv_model isa EDOnlyEDMFX
-        ᶠu³⁰ = p.precomputed.ᶠu³
         ᶜu⁰ = ᶜu
     end
-    @. ᶜu⁰ = C123(Y.c.uₕ) + ᶜinterp(C123(ᶠu³⁰))  # Set here, but used in a different function
+
+    ᶜu⁰ = ᶜu_lazy(Y.c.uₕ, ᶠu³⁰) 
 
     @. ᶜlinear_buoygrad = buoyancy_gradients(
         BuoyGradMean(),
@@ -1039,7 +1050,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
 
     ρatke_flux_values = Fields.field_values(ρatke_flux)
     ρ_int_values = Fields.field_values(Fields.level(Y.c.ρ, 1))
-    u_int_values = Fields.field_values(Fields.level(ᶜu, 1))
+    u_int_values = Fields.field_values(Fields.level(Base.materialize(ᶜu), 1))
     ustar_values = Fields.field_values(ustar)
     int_local_geometry_values =
         Fields.field_values(Fields.level(Fields.local_geometry_field(Y.c), 1))

src/cache/precipitation_precomputed_quantities.jl

@@ -48,12 +48,14 @@ function set_precipitation_velocities!(
     moisture_model::NonEquilMoistModel,
     precip_model::Microphysics1Moment,
 )
-    (; ᶜwₗ, ᶜwᵢ, ᶜwᵣ, ᶜwₛ, ᶜwₜqₜ, ᶜwₕhₜ, ᶜts, ᶜu) = p.precomputed
+    (; ᶜwₗ, ᶜwᵢ, ᶜwᵣ, ᶜwₛ, ᶜwₜqₜ, ᶜwₕhₜ, ᶜts) = p.precomputed
     (; ᶜΦ) = p.core
     cmc = CAP.microphysics_cloud_params(p.params)
     cmp = CAP.microphysics_1m_params(p.params)
     thp = CAP.thermodynamics_params(p.params)
 
+    ᶜu = ᶜu_lazy(Y.c.uₕ, Y.f.u₃)
+
     # compute the precipitation terminal velocity [m/s]
     @. ᶜwᵣ = CM1.terminal_velocity(
         cmp.pr,
@@ -104,14 +106,15 @@ function set_precipitation_velocities!(
     moisture_model::NonEquilMoistModel,
     precip_model::Microphysics2Moment,
 )
-    (; ᶜwₗ, ᶜwᵢ, ᶜwᵣ, ᶜwₛ, ᶜwnₗ, ᶜwnᵣ, ᶜwₜqₜ, ᶜwₕhₜ, ᶜts, ᶜu) = p.precomputed
+    (; ᶜwₗ, ᶜwᵢ, ᶜwᵣ, ᶜwₛ, ᶜwnₗ, ᶜwnᵣ, ᶜwₜqₜ, ᶜwₕhₜ, ᶜts) = p.precomputed
     (; ᶜΦ) = p.core
 
     cm1c = CAP.microphysics_cloud_params(p.params)
     cm1p = CAP.microphysics_1m_params(p.params)
     cm2p = CAP.microphysics_2m_params(p.params)
     thp = CAP.thermodynamics_params(p.params)
 
+    ᶜu = ᶜu_lazy(Y.c.uₕ, Y.f.u₃)
     # compute the precipitation terminal velocity [m/s]
     # TODO sedimentation of snow is based on the 1M scheme
     @. ᶜwnᵣ = getindex(
@@ -280,7 +283,7 @@ function set_precipitation_cache!(
 end
 function set_precipitation_cache!(Y, p, ::Microphysics1Moment, _)
     (; dt) = p
-    (; ᶜts, ᶜwᵣ, ᶜwₛ, ᶜu) = p.precomputed
+    (; ᶜts, ᶜwᵣ, ᶜwₛ) = p.precomputed
     (; ᶜSqₗᵖ, ᶜSqᵢᵖ, ᶜSqᵣᵖ, ᶜSqₛᵖ) = p.precomputed
 
     (; q_tot, q_liq, q_ice, q_rai, q_sno) = p.precomputed.ᶜspecific

src/cache/precomputed_quantities.jl

@@ -43,8 +43,6 @@ function implicit_precomputed_quantities(Y, atmos)
     n = n_mass_flux_subdomains(turbconv_model)
     gs_quantities = (;
         ᶜspecific = Base.materialize(ᶜspecific_gs_tracers(Y)),
-        ᶜu = similar(Y.c, C123{FT}),
-        ᶠu³ = similar(Y.f, CT3{FT}),
         ᶠu = similar(Y.f, CT123{FT}),
         ᶜK = similar(Y.c, FT),
         ᶜts = similar(Y.c, TST),
@@ -72,8 +70,6 @@ function implicit_precomputed_quantities(Y, atmos)
             ᶜmse⁰ = similar(Y.c, FT),
             ᶜq_tot⁰ = similar(Y.c, FT),
             ᶠu₃⁰ = similar(Y.f, C3{FT}),
-            ᶜu⁰ = similar(Y.c, C123{FT}),
-            ᶠu³⁰ = similar(Y.f, CT3{FT}),
             ᶜK⁰ = similar(Y.c, FT),
             ᶜts⁰ = similar(Y.c, TST),
             ᶜρ⁰ = similar(Y.c, FT),
@@ -221,8 +217,6 @@ function precomputed_quantities(Y, atmos)
             ᶜturb_entrʲs = similar(Y.c, NTuple{n, FT}),
             ᶠnh_pressure³_buoyʲs = similar(Y.f, NTuple{n, CT3{FT}}),
             ᶠnh_pressure³_dragʲs = similar(Y.f, NTuple{n, CT3{FT}}),
-            ᶠu³⁰ = similar(Y.f, CT3{FT}),
-            ᶜu⁰ = similar(Y.c, C123{FT}),
             ᶜK⁰ = similar(Y.c, FT),
             ᶜmixing_length_tuple = similar(Y.c, MixingLength{FT}),
             ᶜK_u = similar(Y.c, FT),
@@ -267,10 +261,13 @@ function precomputed_quantities(Y, atmos)
     )
 end
 
-# Interpolates the third contravariant component of Y.c.uₕ to cell faces.
-function compute_ᶠuₕ³(ᶜuₕ, ᶜρ)
-    ᶜJ = Fields.local_geometry_field(ᶜρ).J
-    return @. lazy(ᶠwinterp(ᶜρ * ᶜJ, CT3(ᶜuₕ)))
+# This is used to set the grid-scale velocity quantities ᶜu, ᶠu³, ᶜK based on
+# ᶠu₃, and it is also used to set the SGS quantities based on ᶠu₃⁰ and ᶠu₃ʲ.
+function set_velocity_quantities!(ᶠu³, ᶜK, ᶠu₃, ᶜuₕ, ᶠuₕ³, ρ)
+    ᶜu = ᶜu_lazy(ᶜuₕ, ᶠu₃)
+    ᶠu³ = ᶠu³_lazy(ᶜuₕ, ρ, ᶠu₃)
+    ᶜK .= compute_kinetic(ᶜuₕ, ᶠu₃)
+    return nothing
 end
 
 """
@@ -323,14 +320,6 @@ function set_velocity_at_top!(Y, turbconv_model)
     return nothing
 end
 
-# This is used to set the grid-scale velocity quantities ᶜu, ᶠu³, ᶜK based on
-# ᶠu₃, and it is also used to set the SGS quantities based on ᶠu₃⁰ and ᶠu₃ʲ.
-function set_velocity_quantities!(ᶜu, ᶠu³, ᶜK, ᶠu₃, ᶜuₕ, ᶠuₕ³)
-    @. ᶜu = C123(ᶜuₕ) + ᶜinterp(C123(ᶠu₃))
-    @. ᶠu³ = ᶠuₕ³ + CT3(ᶠu₃)
-    ᶜK .= compute_kinetic(ᶜuₕ, ᶠu₃)
-    return nothing
-end
 
 function set_sgs_ᶠu₃!(w_function, ᶠu₃, Y, turbconv_model)
     ρaʲs(sgsʲs) = map(sgsʲ -> sgsʲ.ρa, sgsʲs)
@@ -455,21 +444,26 @@ 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, ᶜK, ᶜts, ᶜp, ᶜh_tot) = p.precomputed
     ᶠuₕ³ = p.scratch.ᶠtemp_CT3
     n = n_mass_flux_subdomains(turbconv_model)
     thermo_params = CAP.thermodynamics_params(p.params)
     thermo_args = (thermo_params, moisture_model, precip_model)
 
     ᶜspecific .= ᶜspecific_gs_tracers(Y)
-    @. ᶠuₕ³ = $compute_ᶠuₕ³(Y.c.uₕ, Y.c.ρ)
+    #@. ᶠuₕ³ = $compute_ᶠuₕ³(Y.c.uₕ, Y.c.ρ)
+    @. ᶠuₕ³ = $ᶠuₕ³_lazy(Y.c.uₕ, Y.c.ρ)
 
     # TODO: We might want to move this to dss! (and rename dss! to something
     # like enforce_constraints!).
     set_velocity_at_surface!(Y, ᶠuₕ³, turbconv_model)
     set_velocity_at_top!(Y, turbconv_model)
 
-    set_velocity_quantities!(ᶜu, ᶠu³, ᶜK, Y.f.u₃, Y.c.uₕ, ᶠuₕ³)
+    ᶜu = ᶜu_lazy(Y.c.uₕ, Y.f.u₃)
+    ᶠu³ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃)
+
+    set_velocity_quantities!(ᶠu³, ᶜK, Y.f.u₃, Y.c.uₕ, ᶠuₕ³, Y.c.ρ)
+
     ᶜJ = Fields.local_geometry_field(Y.c).J
     @. ᶠu = CT123(ᶠwinterp(Y.c.ρ * ᶜJ, CT12(ᶜu))) + CT123(ᶠu³)
     if n > 0
@@ -517,7 +511,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
+    (; ᶜts, ᶜp) = p.precomputed
     ᶠuₕ³ = p.scratch.ᶠtemp_CT3 # updated in set_implicit_precomputed_quantities!
     thermo_params = CAP.thermodynamics_params(p.params)
 

src/cache/prognostic_edmf_precomputed_quantities.jl

@@ -22,7 +22,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!(
     (; turbconv_model) = p.atmos
     (; ᶜΦ,) = p.core
     (; ᶜp, ᶜh_tot, ᶜK) = p.precomputed
-    (; ᶜtke⁰, ᶜρa⁰, ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜρ⁰, ᶜmse⁰, ᶜq_tot⁰) =
+    (; ᶜtke⁰, ᶜρa⁰, ᶠu₃⁰, ᶜK⁰, ᶜts⁰, ᶜρ⁰, ᶜmse⁰, ᶜq_tot⁰) =
         p.precomputed
     if p.atmos.moisture_model isa NonEquilMoistModel &&
        p.atmos.precip_model isa Microphysics1Moment
@@ -76,8 +76,11 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!(
             turbconv_model,
         )
     end
-    set_sgs_ᶠu₃!(u₃⁰, ᶠu₃⁰, Y, turbconv_model)
-    set_velocity_quantities!(ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶠu₃⁰, Y.c.uₕ, ᶠuₕ³)
+
+    set_sgs_ᶠu₃!(u₃⁰, ᶠu₃⁰, Y, turbconv_model) # Sets \ᶠu₃⁰
+    ᶠu³⁰ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, ᶠu₃⁰)
+    set_velocity_quantities!(ᶠu³⁰, ᶜK⁰, ᶠu₃⁰, Y.c.uₕ, ᶠuₕ³, Y.c.ρ)
+
     # @. ᶜK⁰ += ᶜtke⁰
     if p.atmos.moisture_model isa NonEquilMoistModel &&
        p.atmos.precip_model isa Microphysics1Moment
@@ -132,7 +135,9 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_draft!(
             ᶜq_snoʲ = Y.c.sgsʲs.:($j).q_sno
         end
 
-        set_velocity_quantities!(ᶜuʲ, ᶠu³ʲ, ᶜKʲ, ᶠu₃ʲ, Y.c.uₕ, ᶠuₕ³)
+        ᶠu³ʲ = ᶠu³_lazy(Y.c.uₕ, ᶜρʲ, ᶠu₃ʲ)
+        set_velocity_quantities!(ᶠu³ʲ, ᶜKʲ, ᶠu₃ʲ, Y.c.uₕ, ᶠuₕ³, ᶜρʲ)
+
         @. ᶠKᵥʲ = (adjoint(CT3(ᶠu₃ʲ)) * ᶠu₃ʲ) / 2
         if p.atmos.moisture_model isa NonEquilMoistModel &&
            p.atmos.precip_model isa Microphysics1Moment
@@ -367,7 +372,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⁰, ᶜp, ᶜρa⁰, ᶜts⁰, ᶜq_tot⁰) = p.precomputed
     (;
         ᶜmixing_length_tuple,
         ᶜmixing_length,
@@ -398,6 +403,9 @@ 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
+
+    ᶜu = Base.materialize(ᶜu_lazy(Y.c.uₕ, Y.f.u₃))
+
     for j in 1:n
         # entrainment/detrainment
         @. ᶜentrʲs.:($$j) = entrainment(
@@ -496,6 +504,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
     # TODO: Make strain_rate_norm calculation a function in eddy_diffusion_closures
     # TODO: Currently the shear production only includes vertical gradients
     ᶠu⁰ = p.scratch.ᶠtemp_C123
+    ᶠu³⁰ = ᶠu³_lazy(Y.c.uₕ, Y.c.ρ, Y.f.u₃)
     @. ᶠu⁰ = C123(ᶠinterp(Y.c.uₕ)) + C123(ᶠu³⁰)
     ᶜstrain_rate = p.scratch.ᶜtemp_UVWxUVW
     ᶜstrain_rate .= compute_strain_rate_center(ᶠu⁰)

src/diagnostics/Diagnostics.jl

@@ -58,6 +58,7 @@ import ..PrognosticSurfaceTemperature
 import ..draft_area
 import ..compute_gm_mixing_length!
 import ..horizontal_integral_at_boundary
+import ..ᶜu_lazy, ..ᶠu³_lazy
 
 # We need the abbreviations for symbols like curl, grad, and so on
 include(joinpath("..", "utils", "abbreviations.jl"))