Merge pull request #3798 from CliMA/ck/single_line_calls

Inline functions that return broadcasted objects

Author: charleskawczynski

src/cache/diagnostic_edmf_precomputed_quantities.jl

@@ -979,8 +979,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

@@ -270,8 +270,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)
@@ -316,8 +315,7 @@ end
 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
+    ᶜK .= compute_kinetic(ᶜuₕ, ᶠu₃)
     return nothing
 end
 

src/cache/prognostic_edmf_precomputed_quantities.jl

@@ -497,8 +497,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
     ᶠ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

@@ -407,8 +407,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

@@ -356,8 +356,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.ρ)
     end
     return nothing
@@ -419,8 +418,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.ρ)
     end
 

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
@@ -64,6 +63,6 @@ NVTX.@annotate function compute_gm_mixing_length!(ᶜmixing_length, Y, p)
         sqrt(max(ᶜlinear_buoygrad, 0)),   #N_eff
         ᶜdz,
         ᶜprandtl_nvec,
-        norm_sqr(ᶜstrain_rate),
+        ᶜstrain_rate_norm,
     )
 end

src/prognostic_equations/vertical_diffusion_boundary_layer.jl

@@ -24,8 +24,7 @@ function vertical_diffusion_boundary_layer_tendency!(
 
     if !disable_momentum_vertical_diffusion(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:
@@ -80,8 +79,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.
 """
@@ -97,8 +95,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

@@ -48,8 +48,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))
@@ -87,12 +86,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)