Make more compile-time friendly

Author: charleskawczynski

src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl

@@ -107,8 +107,8 @@ function horizontal_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::SmagorinskyLill
     @. Yₜ.c.ρe_tot += wdivₕ(Y.c.ρ * ᶜD_smag * gradₕ(ᶜh_tot))
 
     ## Tracer diffusion and associated mass changes
-    UU.unrolled_foreach(filter(is_tracer_var, propertynames(Yₜ.c))) do ρχ_name
-        if ρχ_name != :ρe_tot
+    UU.unrolled_foreach(propertynames(Yₜ.c)) do ρχ_name
+        if is_tracer_var(ρχ_name) && ρχ_name != :ρe_tot
             ᶜρχ = getproperty(Y.c, ρχ_name)
             ᶜρχₜ = getproperty(Yₜ.c, ρχ_name)
             ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
@@ -161,8 +161,8 @@ 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
-    UU.unrolled_foreach(filter(is_tracer_var, propertynames(Yₜ.c))) do ρχ_name
-        if ρχ_name != :ρe_tot
+    UU.unrolled_foreach(propertynames(Yₜ.c)) do ρχ_name
+        if is_tracer_var(ρχ_name) && ρχ_name != :ρe_tot
             ᶜρχ = getproperty(Y.c, ρχ_name)
             ᶜρχₜ = getproperty(Yₜ.c, ρχ_name)
             ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))

src/prognostic_equations/advection.jl

@@ -213,8 +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) ...
-    UU.unrolled_foreach(filter(is_tracer_var, propertynames(Yₜ.c))) do ρχ_name
-        if !(ρχ_name in (:ρe_tot, :ρq_tot))
+    UU.unrolled_foreach(propertynames(Yₜ.c)) do ρχ_name
+        if is_tracer_var(ρχ_name) && !(ρχ_name in (:ρe_tot, :ρq_tot))
             ᶜρχ = getproperty(Y.c, ρχ_name)
             ᶜρχₜ = getproperty(Yₜ.c, ρχ_name)
             ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))

src/prognostic_equations/hyperdiffusion.jl

@@ -289,23 +289,25 @@ NVTX.@annotate function apply_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t)
     # by the limiter. Is this a significant problem?
     # TODO: Figure out why caching the duplicated tendencies in ᶜtemp_scalar
     # triggers allocations.
-    UU.unrolled_foreach(filter(is_tracer_var, propertynames(Yₜ.c))) do ρχ_name
-        ᶜρχ = getproperty(Y.c, ρχ_name)
-        ᶜρχₜ = getproperty(Yₜ.c, ρχ_name)
-        ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
-
-        ν₄_scalar = ifelse(
-            ρχ_name in (:ρq_rai, :ρq_sno, :ρn_rai),
-            α_hyperdiff_tracer * ν₄_scalar,
-            ν₄_scalar,
-        )
-        ᶜ∇²χ = getproperty(ᶜ∇²specific_tracers, specific_name(ρχ_name))
-        @. ᶜρχₜ -= ν₄_scalar * wdivₕ(Y.c.ρ * gradₕ(ᶜ∇²χ))
-
-        # Exclude contributions from hyperdiffusion of condensate, 
-        # precipitating species from mass tendency. 
-        if ρχ_name == :ρq_tot
-            @. Yₜ.c.ρ -= ν₄_scalar * wdivₕ(Y.c.ρ * gradₕ(ᶜ∇²χ))
+    UU.unrolled_foreach(propertynames(Yₜ.c)) do ρχ_name
+        if is_tracer_var(ρχ_name)
+            ᶜρχ = getproperty(Y.c, ρχ_name)
+            ᶜρχₜ = getproperty(Yₜ.c, ρχ_name)
+            ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
+
+            ν₄_scalar = ifelse(
+                ρχ_name in (:ρq_rai, :ρq_sno, :ρn_rai),
+                α_hyperdiff_tracer * ν₄_scalar,
+                ν₄_scalar,
+            )
+            ᶜ∇²χ = getproperty(ᶜ∇²specific_tracers, specific_name(ρχ_name))
+            @. ᶜρχₜ -= ν₄_scalar * wdivₕ(Y.c.ρ * gradₕ(ᶜ∇²χ))
+
+            # Exclude contributions from hyperdiffusion of condensate, 
+            # precipitating species from mass tendency. 
+            if ρχ_name == :ρq_tot
+                @. Yₜ.c.ρ -= ν₄_scalar * wdivₕ(Y.c.ρ * gradₕ(ᶜ∇²χ))
+            end
         end
     end
     if turbconv_model isa PrognosticEDMFX

src/prognostic_equations/remaining_tendency.jl

@@ -166,8 +166,8 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t)
     @. Yₜ.c.ρe_tot += vst_ρe_tot
 
     # TODO: can we write this out explicitly?
-    UU.unrolled_foreach(filter(is_tracer_var, propertynames(Yₜ.c))) do ρχ_name
-        if ρχ_name != :ρe_tot
+    UU.unrolled_foreach(propertynames(Yₜ.c)) do ρχ_name
+        if is_tracer_var(ρχ_name) && ρχ_name != :ρe_tot
             ᶜρχ = getproperty(Y.c, ρχ_name)
             ᶜρχₜ = getproperty(Yₜ.c, ρχ_name)
             ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))

src/prognostic_equations/surface_flux.jl

@@ -115,11 +115,11 @@ 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
-    UU.unrolled_foreach(filter(is_tracer_var, propertynames(Yₜ.c))) do ρχ_name
-        ᶜρχ = getproperty(Y.c, ρχ_name)
-        ᶜρχₜ = getproperty(Yₜ.c, ρχ_name)
-        ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
-        if ρχ_name != :ρe_tot
+    UU.unrolled_foreach(propertynames(Yₜ.c)) do ρχ_name
+        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

src/prognostic_equations/vertical_diffusion_boundary_layer.jl

@@ -96,12 +96,12 @@ function vertical_diffusion_boundary_layer_tendency!(
 
     ᶜρχₜ_diffusion = p.scratch.ᶜtemp_scalar
     ᶜK_h_scaled = p.scratch.ᶜtemp_scalar_2
-    UU.unrolled_foreach(filter(is_tracer_var, propertynames(Yₜ.c))) do ρχ_name
-        ᶜρχ = getproperty(Y.c, ρχ_name)
-        ᶜρχₜ = getproperty(Yₜ.c, ρχ_name)
-        ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
-        if ρχ_name != :ρe_tot
-            if χ_name in (:q_rai, :q_sno, :n_rai)
+    UU.unrolled_foreach(propertynames(Yₜ.c)) do ρχ_name
+        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