Remove more use of matching_subfields

Author: charleskawczynski

src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl

@@ -107,14 +107,18 @@ 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 CA.matching_subfields(Yₜ.c, ᶜspecific)
-        χ_name == :e_tot && continue
-        ᶜρχₜ_diffusion =
-            @. p.scratch.ᶜtemp_scalar = wdivₕ(Y.c.ρ * ᶜD_smag * gradₕ(ᶜχ))
-        @. ᶜρχₜ += ᶜρχₜ_diffusion
-        # Rain and snow does not affect the mass
-        if χ_name ∉ (:q_rai, :q_sno)
-            @. Yₜ.c.ρ += ᶜρχₜ_diffusion
+    UU.unrolled_foreach(filter(is_tracer_var, propertynames(Yₜ.c))) do ρχ_name
+        if ρχ_name != :ρe_tot
+            ᶜρχ = getproperty(Y.c, ρχ_name)
+            ᶜρχₜ = getproperty(Yₜ.c, ρχ_name)
+            ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
+            ᶜρχₜ_diffusion =
+                @. p.scratch.ᶜtemp_scalar = wdivₕ(Y.c.ρ * ᶜD_smag * gradₕ(ᶜχ))
+            @. ᶜρχₜ += ᶜρχₜ_diffusion
+            # Rain and snow does not affect the mass
+            if ρχ_name ∉ (:ρq_rai, :ρq_sno)
+                @. Yₜ.c.ρ += ᶜρχₜ_diffusion
+            end
         end
     end
 

src/prognostic_equations/advection.jl

@@ -213,10 +213,14 @@ 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 matching_subfields(Yₜ.c, ᶜspecific)
-        χ_name in (:e_tot, :q_tot) && continue
-        vtt = vertical_transport(ᶜρ, ᶠu³, ᶜχ, float(dt), tracer_upwinding)
-        @. ᶜρχₜ += vtt
+    UU.unrolled_foreach(filter(is_tracer_var, propertynames(Yₜ.c))) do ρχ_name
+        if !(ρχ_name in (:ρe_tot, :ρq_tot))
+            ᶜρχ = getproperty(Y.c, ρχ_name)
+            ᶜρχₜ = getproperty(Yₜ.c, ρχ_name)
+            ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
+            vtt = vertical_transport(ᶜρ, ᶠu³, ᶜχ, float(dt), tracer_upwinding)
+            @. ᶜρχₜ += vtt
+        end
     end
     # ... and upwinding correction of energy and total water.
     # (The central advection of energy and total water is done implicitly.)

src/prognostic_equations/hyperdiffusion.jl

@@ -289,17 +289,22 @@ 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.
-    for (ᶜρχₜ, ᶜ∇²χ, χ_name) in matching_subfields(Yₜ.c, ᶜ∇²specific_tracers)
+    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),
+            ρχ_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
+        if ρχ_name == :ρq_tot
             @. Yₜ.c.ρ -= ν₄_scalar * wdivₕ(Y.c.ρ * gradₕ(ᶜ∇²χ))
         end
     end

src/prognostic_equations/remaining_tendency.jl

@@ -166,11 +166,15 @@ 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 matching_subfields(Yₜ.c, ᶜspecific)
-        χ_name == :e_tot && continue
-        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?
+    UU.unrolled_foreach(filter(is_tracer_var, propertynames(Yₜ.c))) do ρχ_name
+        if ρχ_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
     end
 
     # Held Suarez tendencies

src/prognostic_equations/surface_flux.jl

@@ -115,17 +115,21 @@ 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 matching_subfields(Yₜ.c, ᶜspecific)
-        χ_name == :e_tot && continue
-        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
+    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 == :ρ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
         end
     end
 end

src/prognostic_equations/vertical_diffusion_boundary_layer.jl

@@ -46,7 +46,7 @@ This function is dispatched based on the type of the vertical diffusion model
       for scalars. Zero-flux boundary conditions are explicitly applied.
     - **Note on mass conservation for `q_tot` diffusion**: The current implementation
       also modifies the tendency of total moist air density `Yₜ.c.ρ` based on the
-      diffusion tendency of total specific humidity `ρq_tot`: 
+      diffusion tendency of total specific humidity `ρq_tot`:
       `Yₜ.c.ρ -= ᶜρχₜ_diffusion_for_q_tot`.
 
 Arguments for all methods:
@@ -58,7 +58,7 @@ Arguments for all methods:
 - `t`: Current simulation time (not directly used in diffusion calculations).
 - `vert_diff_model` (for dispatched methods): The specific vertical diffusion model instance.
 
-Modifies components of tendency vector `Yₜ.c` (e.g., `Yₜ.c.uₕ`, `Yₜ.c.ρe_tot`, `Yₜ.c.ρ`, and 
+Modifies components of tendency vector `Yₜ.c` (e.g., `Yₜ.c.uₕ`, `Yₜ.c.ρe_tot`, `Yₜ.c.ρ`, and
 various tracer fields such as `Yₜ.c.ρq_tot`).
 """
 
@@ -96,24 +96,28 @@ function vertical_diffusion_boundary_layer_tendency!(
 
     ᶜρχₜ_diffusion = p.scratch.ᶜtemp_scalar
     ᶜK_h_scaled = p.scratch.ᶜtemp_scalar_2
-    for (ᶜρχₜ, ᶜχ, χ_name) in matching_subfields(Yₜ.c, ᶜspecific)
-        χ_name == :e_tot && continue
-        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
+    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)
+                @. ᶜ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
         end
     end
 end

src/utils/variable_manipulations.jl

@@ -42,6 +42,20 @@ Arguments:
 """
 @inline specific(ρχ, ρ) = ρχ / ρ
 
+#! format: off
+@inline specific_name(ρχ_name::Symbol) =
+        if ρχ_name == :ρe_tot; return :e_tot
+    elseif ρχ_name == :ρq_tot; return :q_tot
+    elseif ρχ_name == :ρq_liq; return :q_liq
+    elseif ρχ_name == :ρq_ice; return :q_ice
+    elseif ρχ_name == :ρq_rai; return :q_rai
+    elseif ρχ_name == :ρn_liq; return :n_liq
+    elseif ρχ_name == :ρn_rai; return :q_rai
+    elseif ρχ_name == :ρq_sno; return :q_sno
+    else; error("Uncaught name: $ρχ_name")
+    end
+#! format: on
+
 @inline function specific(ρaχ, ρa, ρχ, ρ, turbconv_model)
     # TODO: Replace turbconv_model struct by parameters, and include a_half in
     # parameters, not in config
@@ -202,36 +216,6 @@ value is assumed to be equal to the value of `ρaχ` in `sgs`.
     return :(NamedTuple{$specific_sgs_names}(($(specific_sgs_values...),)))
 end
 
-"""
-    matching_subfields(tendency_field, specific_field)
-
-Given a field that contains the tendencies of variables of the form `ρχ` or
-`ρaχ` and another field that contains the values of specific variables `χ`,
-returns all tuples `(tendency_field.<ρχ or ρaχ>, specific_field.<χ>, :<χ>)`.
-Variables in `tendency_field` that do not have matching variables in
-`specific_field` are omitted, as are variables in `specific_field` that do not
-have matching variables in `tendency_field`. This function is needed to avoid
-allocations due to failures in type inference, which are triggered when the
-`propertynames` of these fields are manipulated during runtime in order to pick
-out the matching subfields (as of Julia 1.8).
-"""
-@generated function matching_subfields(tendency_field, specific_field)
-    tendency_names = Base._nt_names(eltype(tendency_field))
-    specific_names = Base._nt_names(eltype(specific_field))
-    prefix = :ρa in tendency_names ? :ρa : :ρ
-    relevant_specific_names =
-        filter(name -> Symbol(prefix, name) in tendency_names, specific_names)
-    subfield_tuples = map(
-        name -> :((
-            tendency_field.$(Symbol(prefix, name)),
-            specific_field.$name,
-            $(QuoteNode(name)),
-        )),
-        relevant_specific_names,
-    )
-    return :(($(subfield_tuples...),))
-end
-
 """
     ρa⁺(gs)