Fixes

Author: charleskawczynski

src/prognostic_equations/hyperdiffusion.jl

@@ -289,25 +289,18 @@ 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(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
+    for (ᶜρχₜ, ᶜ∇²χ, χ_name) in matching_subfields(Yₜ.c, ᶜ∇²specific_tracers)
+        ν₄_scalar = ifelse(
+            χ_name in (:q_rai, :q_sno, :n_rai),
+            α_hyperdiff_tracer * ν₄_scalar,
+            ν₄_scalar,
+        )
+        @. ᶜρχₜ -= ν₄_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
     if turbconv_model isa PrognosticEDMFX

src/prognostic_equations/vertical_diffusion_boundary_layer.jl

@@ -115,7 +115,7 @@ function vertical_diffusion_boundary_layer_tendency!(
             @. ᶜρχₜ -= ᶜρχₜ_diffusion
             # Only add contribution from total water diffusion to mass tendency 
             # (exclude contributions from diffusion of condensate, precipitation) 
-            if χ_name == :q_tot
+            if ρχ_name == :ρq_tot
                 @. Yₜ.c.ρ -= ᶜρχₜ_diffusion
             end
         end

src/utils/variable_manipulations.jl

@@ -42,20 +42,6 @@ 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
@@ -158,6 +144,35 @@ function divide_by_ρa(ρaχ, ρa, ρχ, ρ, turbconv_model)
     return ρa == 0 ? ρχ / ρ : weight * ρaχ / ρa + (1 - weight) * ρχ / ρ
 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
+
 # Helper functions for manipulating symbols in the generated functions:
 has_prefix(symbol, prefix_symbol) =
     startswith(string(symbol), string(prefix_symbol))