Use LazyBroadcast in set_velocity_quantities

Author: charleskawczynski

src/cache/diagnostic_edmf_precomputed_quantities.jl

@@ -944,8 +944,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
     ᶠu⁰ = p.scratch.ᶠtemp_C123
     @. ᶠu⁰ = C123(ᶠinterp(Y.c.uₕ)) + C123(ᶠu³⁰)
     ᶜstrain_rate = p.scratch.ᶜtemp_UVWxUVW
-    bc_strain_rate = compute_strain_rate_center(ᶠu⁰)
-    @. ᶜstrain_rate = bc_strain_rate
+    ᶜstrain_rate .= compute_strain_rate_center(ᶠu⁰)
     @. ᶜstrain_rate_norm = norm_sqr(ᶜstrain_rate)
 
     ᶜprandtl_nvec = p.scratch.ᶜtemp_scalar

src/cache/precomputed_quantities.jl

@@ -197,8 +197,7 @@ the `turbconv_model` is EDMFX, the `Y.f.sgsʲs` are also modified so that each
 """
 function set_velocity_at_surface!(Y, ᶠuₕ³, turbconv_model)
     sfc_u₃ = Fields.level(Y.f.u₃.components.data.:1, half)
-    bc_sfc_u₃ = surface_velocity(Y.f.u₃, ᶠuₕ³)
-    @. sfc_u₃ = bc_sfc_u₃
+    sfc_u₃ .= surface_velocity(Y.f.u₃, ᶠuₕ³)
     if turbconv_model isa PrognosticEDMFX
         for j in 1:n_mass_flux_subdomains(turbconv_model)
             sfc_u₃ʲ = Fields.level(Y.f.sgsʲs.:($j).u₃.components.data.:1, half)
@@ -240,12 +239,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,15 +382,18 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
     @. ᶜspecific = specific_gs(Y.c)
     ᶜρ = Y.c.ρ
     ᶜuₕ = Y.c.uₕ
-    bc_ᶠuₕ³ = compute_ᶠuₕ³(ᶜuₕ, ᶜρ)
-    @. ᶠuₕ³ = bc_ᶠuₕ³
+    ᶠu₃ = Y.f.u₃
+    ᶠuₕ³ .= compute_ᶠuₕ³(ᶜuₕ, ᶜρ)
 
     # 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³ .= compute_ᶠu³(ᶜuₕ, ᶜρ, ᶠu₃)
+    ᶜu .= compute_ᶜu(ᶜuₕ, ᶠu₃)
+    ᶜK .= compute_kinetic(ᶜuₕ, ᶠu₃)
+
     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 +496,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,
@@ -521,12 +521,10 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
 
     if vert_diff isa DecayWithHeightDiffusion
         (; ᶜK_h) = p.precomputed
-        bc_K_h = compute_eddy_diffusivity_coefficient(ᶜρ, vert_diff)
-        @. ᶜK_h = bc_K_h
+        ᶜK_h .= compute_eddy_diffusivity_coefficient(ᶜρ, vert_diff)
     elseif vert_diff isa VerticalDiffusion
         (; ᶜK_h) = p.precomputed
-        bc_K_h = compute_eddy_diffusivity_coefficient(Y.c.uₕ, ᶜp, vert_diff)
-        @. ᶜK_h = bc_K_h
+        ᶜK_h .= compute_eddy_diffusivity_coefficient(Y.c.uₕ, ᶜp, vert_diff)
     end
 
     # TODO

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,34 @@ 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ₕ³)
+    ᶠu³⁰ .= compute_ᶠu³(ᶜuₕ, ᶜρ, ᶠu₃)
+    ᶜu⁰ .= compute_ᶜu(ᶜuₕ, ᶠu₃⁰)
+    ᶜK⁰ .= compute_kinetic(ᶜuₕ, ᶠu₃⁰)
+
     # @. ᶜ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 +95,10 @@ 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ₕ³)
+        ᶠu³ʲ .= compute_ᶠu³(ᶜuₕ, ᶜρ, ᶠu₃)
+        ᶜuʲ .= compute_ᶜu(ᶜuₕ, ᶠu₃ʲ)
+        ᶜKʲ .= compute_kinetic(ᶜuₕ, ᶠu₃ʲ)
+
         @. ᶠKᵥʲ = (adjoint(CT3(ᶠu₃ʲ)) * ᶠu₃ʲ) / 2
         @. ᶜtsʲ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmseʲ - ᶜΦ, ᶜq_totʲ)
         @. ᶜρʲ = TD.air_density(thermo_params, ᶜtsʲ)
@@ -306,8 +316,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_closures!(
     ᶠu⁰ = p.scratch.ᶠtemp_C123
     @. ᶠu⁰ = C123(ᶠinterp(Y.c.uₕ)) + C123(ᶠu³⁰)
     ᶜstrain_rate = p.scratch.ᶜtemp_UVWxUVW
-    bc_strain_rate = compute_strain_rate_center(ᶠu⁰)
-    @. ᶜstrain_rate = bc_strain_rate
+    ᶜstrain_rate .= compute_strain_rate_center(ᶠu⁰)
     @. ᶜstrain_rate_norm = norm_sqr(ᶜstrain_rate)
 
     ᶜprandtl_nvec = p.scratch.ᶜtemp_scalar

src/initial_conditions/initial_conditions.jl

@@ -459,8 +459,7 @@ function overwrite_initial_conditions!(
     Y.c.uₕ .= C12.(Geometry.UVVector.(vel))
     Y.f.u₃ .= ᶠinterp.(C3.(Geometry.WVector.(vel)))
     e_kin = similar(ᶜT)
-    bc_kinetic = compute_kinetic(Y.c.uₕ, Y.f.u₃)
-    @. e_kin = bc_kinetic
+    e_kin .= compute_kinetic(Y.c.uₕ, Y.f.u₃)
     e_pot = Fields.coordinate_field(Y.c).z .* thermo_params.grav
     Y.c.ρe_tot .= TD.total_energy.(thermo_params, ᶜts, e_kin, e_pot) .* Y.c.ρ
     if hasproperty(Y.c, :ρq_tot)

src/prognostic_equations/edmfx_sgs_flux.jl

@@ -225,8 +225,7 @@ function edmfx_sgs_diffusive_flux_tendency!(
 
         # momentum
         ᶠstrain_rate = p.scratch.ᶠtemp_UVWxUVW
-        bc_strain_rate = compute_strain_rate_face(ᶜu⁰)
-        @. ᶠstrain_rate = bc_strain_rate
+        ᶠstrain_rate .= compute_strain_rate_face(ᶜu⁰)
         @. Yₜ.c.uₕ -= C12(ᶜdivᵥ(-(2 * ᶠρaK_u * ᶠstrain_rate)) / Y.c.ρ)
         # apply boundary condition for momentum flux
         ᶜdivᵥ_uₕ = Operators.DivergenceF2C(
@@ -298,8 +297,7 @@ function edmfx_sgs_diffusive_flux_tendency!(
 
         # momentum
         ᶠstrain_rate = p.scratch.ᶠtemp_UVWxUVW
-        bc_strain_rate = compute_strain_rate_face(ᶜu)
-        @. ᶠstrain_rate = bc_strain_rate
+        ᶠstrain_rate .= compute_strain_rate_face(ᶜu)
         @. Yₜ.c.uₕ -= C12(ᶜdivᵥ(-(2 * ᶠρaK_u * ᶠstrain_rate)) / Y.c.ρ)
         # apply boundary condition for momentum flux
         ᶜdivᵥ_uₕ = Operators.DivergenceF2C(

src/prognostic_equations/gm_sgs_closures.jl

@@ -47,8 +47,7 @@ NVTX.@annotate function compute_gm_mixing_length!(ᶜmixing_length, Y, p)
     ᶠu = p.scratch.ᶠtemp_C123
     @. ᶠu = C123(ᶠinterp(Y.c.uₕ)) + C123(ᶠu³)
     ᶜstrain_rate = p.scratch.ᶜtemp_UVWxUVW
-    bc_strain_rate = compute_strain_rate_center(ᶠu)
-    @. ᶜstrain_rate = bc_strain_rate
+    ᶜstrain_rate .= compute_strain_rate_center(ᶠu)
     @. ᶜstrain_rate_norm = norm_sqr(ᶜstrain_rate)
 
     ᶜprandtl_nvec = p.scratch.ᶜtemp_scalar_2

src/prognostic_equations/vertical_diffusion_boundary_layer.jl

@@ -31,8 +31,7 @@ function vertical_diffusion_boundary_layer_tendency!(
 
     if diffuse_momentum(p.atmos.vert_diff)
         ᶠstrain_rate = p.scratch.ᶠtemp_UVWxUVW
-        bc_strain_rate = compute_strain_rate_face(ᶜu)
-        @. ᶠstrain_rate = bc_strain_rate
+        ᶠstrain_rate .= compute_strain_rate_face(ᶜu)
         @. Yₜ.c.uₕ -= C12(
             ᶜdivᵥ(-2 * ᶠinterp(Y.c.ρ) * ᶠinterp(ᶜK_u) * ᶠstrain_rate) / Y.c.ρ,
         )

src/utils/utilities.jl

@@ -48,8 +48,7 @@ sort_files_by_time(files) =
     permute!(files, sortperm(time_from_filename.(files)))
 
 """
-    bc_kinetic = compute_kinetic(uₕ::Field, uᵥ::Field)
-    @. κ = bc_kinetic
+    κ .= compute_kinetic(uₕ::Field, uᵥ::Field)
 
 Compute the specific kinetic energy at cell centers, resulting in `κ` from
 individual velocity components:
@@ -78,8 +77,7 @@ state.
 compute_kinetic(Y::Fields.FieldVector) = compute_kinetic(Y.c.uₕ, Y.f.u₃)
 
 """
-    bc_ϵ = compute_strain_rate_center(u::Field)
-    @. ϵ = bc_ϵ
+    ϵ .= compute_strain_rate_center(u::Field)
 
 Compute the strain_rate at cell centers from velocity at cell faces.
 """
@@ -93,8 +91,7 @@ function compute_strain_rate_center(u::Fields.Field)
 end
 
 """
-    bc_ϵ = compute_strain_rate_face(u::Field)
-    @. ϵ = bc_ϵ
+    ϵ .= compute_strain_rate_face(u::Field)
 
 Compute the strain_rate at cell faces from velocity at cell centers.
 """

test/utilities.jl

@@ -47,8 +47,7 @@ end
     CT123 = Geometry.Contravariant123Vector
     # Exercise function
     κ = zeros(cent_space)
-    bc_kinetic = CA.compute_kinetic(uₕ, uᵥ)
-    @. κ = bc_kinetic
+    κ .= CA.compute_kinetic(uₕ, uᵥ)
     ᶜκ_exact = @. 1 // 2 *
        cos(z)^2 *
        ((sin(x)^2) * (cos(y)^2) + (cos(x)^2) * (sin(y)^2))
@@ -86,12 +85,8 @@ end
     ᶠu = @. UVW(Geometry.UVector(ᶠu)) +
        UVW(Geometry.VVector(ᶠv)) +
        UVW(Geometry.WVector(ᶠw))
-    bc_strain_rate =
-        CA.compute_strain_rate_center(Geometry.Covariant123Vector.(ᶠu))
-    @. ᶜϵ = bc_strain_rate
-    bc_strain_rate =
-        CA.compute_strain_rate_face(Geometry.Covariant123Vector.(ᶜu))
-    @. ᶠϵ = bc_strain_rate
+    ᶜϵ .= CA.compute_strain_rate_center(Geometry.Covariant123Vector.(ᶠu))
+    ᶠϵ .= CA.compute_strain_rate_face(Geometry.Covariant123Vector.(ᶜu))
 
     # Center valued strain rate
     @test ᶜϵ.components.data.:1 == ᶜϵ.components.data.:1 .* FT(0)