Use LazyBroadcast in set_velocity_quantities

Author: charleskawczynski

src/cache/precomputed_quantities.jl

@@ -240,12 +240,10 @@ 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₃)
-    bc_kinetic = compute_kinetic(ᶜuₕ, ᶠu₃)
-    @. ᶜK = bc_kinetic
-    return nothing
+compute_ᶜu(ᶜuₕ, ᶠu₃) = @lazy @. C123(ᶜuₕ) + ᶜinterp(C123(ᶠu₃))
+function compute_ᶠu³(ᶜuₕ, ᶜρ, ᶠu₃)
+    ᶠuₕ³ = compute_ᶠuₕ³(ᶜuₕ, ᶜρ)
+    return @lazy @. ᶠuₕ³ + CT3(ᶠu₃)
 end
 
 function set_sgs_ᶠu₃!(w_function, ᶠu₃, Y, turbconv_model)
@@ -385,6 +383,7 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
     @. ᶜspecific = specific_gs(Y.c)
     ᶜρ = Y.c.ρ
     ᶜuₕ = Y.c.uₕ
+    ᶠu₃ = Y.f.u₃
     bc_ᶠuₕ³ = compute_ᶠuₕ³(ᶜuₕ, ᶜρ)
     @. ᶠuₕ³ = bc_ᶠuₕ³
 
@@ -393,7 +392,13 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
     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ₕ³)
+    bc_ᶠu³ = compute_ᶠu³(ᶜuₕ, ᶜρ, ᶠu₃)
+    bc_ᶜu = compute_ᶜu(ᶜuₕ, ᶠu₃)
+    bc_ᶜK = compute_kinetic(ᶜuₕ, ᶠu₃)
+    @. ᶠu³ = bc_ᶠu³
+    @. ᶜu = bc_ᶜu
+    @. ᶜK = bc_ᶜK
+
     if n > 0
         # TODO: In the following increments to ᶜK, we actually need to add
         # quantities of the form ᶜρaχ⁰ / ᶜρ⁰ and ᶜρaχʲ / ᶜρʲ to ᶜK, rather than
@@ -496,8 +501,8 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
     end
 
     if turbconv_model isa PrognosticEDMFX
-        set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, ᶠuₕ³, t)
-        set_prognostic_edmf_precomputed_quantities_environment!(Y, p, ᶠuₕ³, t)
+        set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, t)
+        set_prognostic_edmf_precomputed_quantities_environment!(Y, p, t)
         set_prognostic_edmf_precomputed_quantities_closures!(Y, p, t)
         set_prognostic_edmf_precomputed_quantities_precipitation!(
             Y,

src/cache/prognostic_edmf_precomputed_quantities.jl

@@ -13,11 +13,13 @@ Updates the edmf environment precomputed quantities stored in `p` for edmfx.
 NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!(
     Y,
     p,
-    ᶠuₕ³,
     t,
 )
     @assert !(p.atmos.moisture_model isa DryModel)
 
+    ᶜuₕ = Y.c.uₕ
+    ᶠu₃ = Y.f.u₃
+    ᶜρ = Y.c.ρ
     thermo_params = CAP.thermodynamics_params(p.params)
     (; turbconv_model) = p.atmos
     (; ᶜΦ,) = p.core
@@ -42,29 +44,37 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!(
         turbconv_model,
     )
     set_sgs_ᶠu₃!(u₃⁰, ᶠu₃⁰, Y, turbconv_model)
-    set_velocity_quantities!(ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶠu₃⁰, Y.c.uₕ, ᶠuₕ³)
+    bc_ᶠu³ = compute_ᶠu³(ᶜuₕ, ᶜρ, ᶠu₃)
+    bc_ᶜu = compute_ᶜu(ᶜuₕ, ᶠu₃⁰)
+    bc_ᶜK = compute_kinetic(ᶜuₕ, ᶠu₃⁰)
+    @. ᶠu³⁰ = bc_ᶠu³
+    @. ᶜu⁰ = bc_ᶜu
+    @. ᶜK⁰ = bc_ᶜK
+
     # @. ᶜK⁰ += ᶜtke⁰
     @. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmse⁰ - ᶜΦ, ᶜq_tot⁰)
     @. ᶜρ⁰ = TD.air_density(thermo_params, ᶜts⁰)
     return nothing
 end
 
 """
-    set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, ᶠuₕ³, t)
+    set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, t)
 
 Updates the draft thermo state and boundary conditions
 precomputed quantities stored in `p` for edmfx.
 """
 NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_draft_and_bc!(
     Y,
     p,
-    ᶠuₕ³,
     t,
 )
     (; moisture_model, turbconv_model) = p.atmos
     #EDMFX BCs only support total energy as state variable
     @assert !(moisture_model isa DryModel)
 
+    ᶜuₕ = Y.c.uₕ
+    ᶠu₃ = Y.f.u₃
+    ᶜρ = Y.c.ρ
     FT = Spaces.undertype(axes(Y.c))
     n = n_mass_flux_subdomains(turbconv_model)
 
@@ -88,7 +98,13 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_draft_and_bc!
         ᶜmseʲ = Y.c.sgsʲs.:($j).mse
         ᶜq_totʲ = Y.c.sgsʲs.:($j).q_tot
 
-        set_velocity_quantities!(ᶜuʲ, ᶠu³ʲ, ᶜKʲ, ᶠu₃ʲ, Y.c.uₕ, ᶠuₕ³)
+        bc_ᶠu³ʲ = compute_ᶠu³(ᶜuₕ, ᶜρ, ᶠu₃)
+        bc_ᶜuʲ = compute_ᶜu(ᶜuₕ, ᶠu₃ʲ)
+        bc_ᶜKʲ = compute_kinetic(ᶜuₕ, ᶠu₃ʲ)
+        @. ᶠu³ʲ = bc_ᶠu³ʲ
+        @. ᶜuʲ = bc_ᶜuʲ
+        @. ᶜKʲ = bc_ᶜKʲ
+
         @. ᶠKᵥʲ = (adjoint(CT3(ᶠu₃ʲ)) * ᶠu₃ʲ) / 2
         @. ᶜtsʲ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmseʲ - ᶜΦ, ᶜq_totʲ)
         @. ᶜρʲ = TD.air_density(thermo_params, ᶜtsʲ)