Use more foreach_scalar

Author: charleskawczynski

src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl

@@ -156,23 +156,19 @@ function vertical_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::SmagorinskyLilly)
     @. Yₜ.c.ρe_tot -= ᶜdivᵥ_ρe_tot(-(ᶠρ * ᶠD_smag * ᶠgradᵥ(ᶜh_tot)))
 
     ## Tracer diffusion and associated mass changes
-    for ρχ_name in propertynames(Yₜ.c)
-        if is_tracer_var(ρχ_name) && ρχ_name != :ρe_tot
-            ᶜρχ = getproperty(Y.c, ρχ_name)
-            ᶜρχₜ = getproperty(Yₜ.c, ρχ_name)
-            ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
-
-            ᶜdivᵥ_ρχ = Operators.DivergenceF2C(;
-                top = Operators.SetValue(C3(FT(0))),
-                bottom = Operators.SetValue(C3(FT(0))),
-            )
-
-            ᶜ∇ᵥρD∇χₜ = @. lazy(ᶜdivᵥ_ρχ(-(ᶠρ * ᶠD_smag * ᶠgradᵥ(ᶜχ))))
-            @. ᶜρχₜ -= ᶜ∇ᵥρD∇χₜ
-            # Rain and snow does not affect the mass
-            if ρχ_name == :ρq_tot
-                @. Yₜ.c.ρ -= ᶜ∇ᵥρD∇χₜ
-            end
+    foreach_scalar(Yₜ, Y) do ᶜρχₜ, ᶜρχ, is_ρq_tot
+        ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
+
+        ᶜdivᵥ_ρχ = Operators.DivergenceF2C(;
+            top = Operators.SetValue(C3(FT(0))),
+            bottom = Operators.SetValue(C3(FT(0))),
+        )
+
+        ᶜ∇ᵥρD∇χₜ = @. lazy(ᶜdivᵥ_ρχ(-(ᶠρ * ᶠD_smag * ᶠgradᵥ(ᶜχ))))
+        @. ᶜρχₜ -= ᶜ∇ᵥρD∇χₜ
+        # Rain and snow does not affect the mass
+        if is_ρq_tot
+            @. Yₜ.c.ρ -= ᶜ∇ᵥρD∇χₜ
         end
     end
 end

src/prognostic_equations/advection.jl

@@ -213,10 +213,8 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
     ᶜρ = Y.c.ρ
 
     # Full vertical advection of passive tracers (like liq, rai, etc) ...
-    for ρχ_name in propertynames(Yₜ.c)
-        if is_tracer_var(ρχ_name) && !(ρχ_name in (:ρe_tot, :ρq_tot))
-            ᶜρχ = getproperty(Y.c, ρχ_name)
-            ᶜρχₜ = getproperty(Yₜ.c, ρχ_name)
+    foreach_scalar(Yₜ, Y) do ᶜρχₜ, ᶜρχ, is_ρq_tot
+        if !is_ρq_tot
             ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
             vtt = vertical_transport(ᶜρ, ᶠu³, ᶜχ, float(dt), tracer_upwinding)
             @. ᶜρχₜ += vtt

src/prognostic_equations/remaining_tendency.jl

@@ -166,15 +166,11 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t)
     @. Yₜ.c.ρe_tot += vst_ρe_tot
 
     # TODO: can we write this out explicitly?
-    for ρχ_name in propertynames(Yₜ.c)
-        if is_tracer_var(ρχ_name) && ρχ_name != :ρe_tot
-            ᶜρχ = getproperty(Y.c, ρχ_name)
-            ᶜρχₜ = getproperty(Yₜ.c, ρχ_name)
-            ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
-            vst_tracer = viscous_sponge_tendency_tracer(ᶜρ, ᶜχ, viscous_sponge)
-            @. ᶜρχₜ += vst_tracer
-            @. Yₜ.c.ρ += vst_tracer  # TODO: This doesn't look right for all tracers here. Remove?
-        end
+    foreach_scalar(Yₜ, Y) do ᶜρχₜ, ᶜρχ, is_ρq_tot
+        ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
+        vst_tracer = viscous_sponge_tendency_tracer(ᶜρ, ᶜχ, viscous_sponge)
+        @. ᶜρχₜ += vst_tracer
+        @. Yₜ.c.ρ += vst_tracer  # TODO: This doesn't look right for all tracers here. Remove?
     end
 
     # Held Suarez tendencies

src/prognostic_equations/surface_flux.jl

@@ -115,21 +115,17 @@ function surface_flux_tendency!(Yₜ, Y, p, t)
     btt = boundary_tendency_scalar(ᶜh_tot, sfc_conditions.ρ_flux_h_tot)
     @. Yₜ.c.ρe_tot -= btt
     ρ_flux_χ = p.scratch.sfc_temp_C3
-    for ρχ_name in propertynames(Yₜ.c)
-        if is_tracer_var(ρχ_name) && ρχ_name != :ρe_tot
-            ᶜρχ = getproperty(Y.c, ρχ_name)
-            ᶜρχₜ = getproperty(Yₜ.c, ρχ_name)
-            ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
-            if ρχ_name == :ρq_tot
-                @. ρ_flux_χ = sfc_conditions.ρ_flux_q_tot
-            else
-                @. ρ_flux_χ = C3(FT(0))
-            end
-            btt = boundary_tendency_scalar(ᶜχ, ρ_flux_χ)
-            @. ᶜρχₜ -= btt
-            if ρχ_name == :ρq_tot
-                @. Yₜ.c.ρ -= btt
-            end
+    foreach_scalar(Yₜ, Y) do ᶜρχₜ, ᶜρχ, is_ρq_tot
+        ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
+        if is_ρq_tot
+            @. ρ_flux_χ = sfc_conditions.ρ_flux_q_tot
+        else
+            @. ρ_flux_χ = C3(FT(0))
+        end
+        btt = boundary_tendency_scalar(ᶜχ, ρ_flux_χ)
+        @. ᶜρχₜ -= btt
+        if is_ρq_tot
+            @. Yₜ.c.ρ -= btt
         end
     end
 end

src/prognostic_equations/vertical_diffusion_boundary_layer.jl

@@ -96,28 +96,24 @@ function vertical_diffusion_boundary_layer_tendency!(
 
     ᶜρχₜ_diffusion = p.scratch.ᶜtemp_scalar
     ᶜK_h_scaled = p.scratch.ᶜtemp_scalar_2
-    for ρχ_name in propertynames(Yₜ.c)
-        if is_tracer_var(ρχ_name) && ρχ_name != :ρe_tot
-            ᶜρχ = getproperty(Y.c, ρχ_name)
-            ᶜρχₜ = getproperty(Yₜ.c, ρχ_name)
-            ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
-            if ρχ_name in (:ρq_rai, :ρq_sno, :ρn_rai)
-                @. ᶜK_h_scaled = α_vert_diff_tracer * ᶜK_h
-            else
-                @. ᶜK_h_scaled = ᶜK_h
-            end
-            ᶜdivᵥ_ρχ = Operators.DivergenceF2C(
-                top = Operators.SetValue(C3(0)),
-                bottom = Operators.SetValue(C3(0)),
-            )
-            @. ᶜρχₜ_diffusion =
-                ᶜdivᵥ_ρχ(-(ᶠinterp(Y.c.ρ) * ᶠinterp(ᶜK_h_scaled) * ᶠgradᵥ(ᶜχ)))
-            @. ᶜρχₜ -= ᶜρχₜ_diffusion
-            # Only add contribution from total water diffusion to mass tendency 
-            # (exclude contributions from diffusion of condensate, precipitation) 
-            if ρχ_name == :ρq_tot
-                @. Yₜ.c.ρ -= ᶜρχₜ_diffusion
-            end
+    foreach_scalar(Yₜ, Y) do ᶜρχₜ, ᶜρχ, is_ρq_tot
+        ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
+        if ρχ_name in (:ρq_rai, :ρq_sno, :ρn_rai)
+            @. ᶜK_h_scaled = α_vert_diff_tracer * ᶜK_h
+        else
+            @. ᶜK_h_scaled = ᶜK_h
+        end
+        ᶜdivᵥ_ρχ = Operators.DivergenceF2C(
+            top = Operators.SetValue(C3(0)),
+            bottom = Operators.SetValue(C3(0)),
+        )
+        @. ᶜρχₜ_diffusion =
+            ᶜdivᵥ_ρχ(-(ᶠinterp(Y.c.ρ) * ᶠinterp(ᶜK_h_scaled) * ᶠgradᵥ(ᶜχ)))
+        @. ᶜρχₜ -= ᶜρχₜ_diffusion
+        # Only add contribution from total water diffusion to mass tendency
+        # (exclude contributions from diffusion of condensate, precipitation)
+        if is_ρq_tot
+            @. Yₜ.c.ρ -= ᶜρχₜ_diffusion
         end
     end
 end