Try foreach_scalar

Author: charleskawczynski

src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl

@@ -107,17 +107,13 @@ 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
-    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.ρ))
-            ᶜρχₜ_diffusion = @. lazy(wdivₕ(Y.c.ρ * ᶜD_smag * gradₕ(ᶜχ)))
-            @. ᶜρχₜ += ᶜρχₜ_diffusion
-            # Rain and snow does not affect the mass
-            if ρχ_name == :ρq_tot
-                @. Yₜ.c.ρ += ᶜρχₜ_diffusion
-            end
+    foreach_scalar(Yₜ, Y) do ᶜρχₜ, ᶜρχ, is_ρq_tot
+        ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
+        ᶜρχₜ_diffusion = @. lazy(wdivₕ(Y.c.ρ * ᶜD_smag * gradₕ(ᶜχ)))
+        @. ᶜρχₜ += ᶜρχₜ_diffusion
+        # Rain and snow does not affect the mass
+        if is_ρq_tot
+            @. Yₜ.c.ρ += ᶜρχₜ_diffusion
         end
     end
 

src/utils/variable_manipulations.jl

@@ -51,6 +51,64 @@ Arguments:
     return ρa == 0 ? ρχ / ρ : weight * ρaχ / ρa + (1 - weight) * ρχ / ρ
 end
 
+"""
+    scalar_names(field)
+
+Filters and returns the names of the variables from a given state
+vector component, excluding `ρ`, `ρe_tot`, and `uₕ`
+
+Arguments:
+
+- `field`: A component of the state vector `Y.c`.
+
+Returns:
+
+- A `Tuple` of `ClimaCore.MatrixFields.FieldName`s corresponding to the tracers.
+"""
+scalar_names(field) =
+    unrolled_filter(MatrixFields.top_level_names(field)) do name
+        !(name in (@name(ρ), @name(ρe_tot), @name(uₕ)))
+    end
+
+"""
+    foreach_scalar(f::F, Yₜ, Y) where {F}
+
+Applies a given function `f` to each scalar variable (except `ρ` and  `ρe_tot`)
+in the state `Y` and its corresponding tendency `Yₜ`.
+This utility abstracts the process of iterating over all scalars. It uses
+`scalar_names` to identify the relevant variables and `unrolled_foreach` to
+ensure a performant loop. For each tracer, it calls the provided function `f` 
+with the tendency field, the state field, and a boolean flag indicating if 
+the current tracer is `ρq_tot` (to allow for special handling).
+
+Arguments:
+
+- `f`: A function to apply to each scalar. It must have the signature
+     `f(ᶜρχₜ, ᶜρχ, is_ρq_tot)`, where `ᶜρχₜ` is the tendency field, `ᶜρχ` is
+     the state field, and `is_ρq_tot` is a `Bool`.
+- `Yₜ`: The tendency state vector.
+- `Y`: The current state vector.
+
+# Example
+
+```julia
+foreach_scalar(Yₜ, Y) do ᶜρχₜ, ᶜρχ, is_ρq_tot
+    # Apply some operation, e.g., a sponge layer
+    @. ᶜρχₜ += some_sponge_function(ᶜρχ)
+    if is_ρq_tot
+        # Perform an additional operation only for ρq_tot
+    end
+end
+```
+"""
+foreach_scalar(f::F, Yₜ, Y) where {F} =
+    unrolled_foreach(scalar_names(Y.c)) do scalar_name
+        ᶜρχₜ = MatrixFields.get_field(Yₜ.c, scalar_name)
+        ᶜρχ = MatrixFields.get_field(Y.c, scalar_name)
+        is_ρq_tot = scalar_name == @name(ρq_tot)
+        f(ᶜρχₜ, ᶜρχ, is_ρq_tot)
+    end
+
 """
     sgs_weight_function(a, a_half)