Add tracer name to foreach_tracer

Author: charleskawczynski

src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl

@@ -107,7 +107,7 @@ 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
-    foreach_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, is_ρq_tot
+    foreach_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name, is_ρq_tot
         ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
         ᶜρχₜ_diffusion = @. lazy(wdivₕ(Y.c.ρ * ᶜD_smag * gradₕ(ᶜχ)))
         @. ᶜρχₜ += ᶜρχₜ_diffusion
@@ -161,7 +161,7 @@ function vertical_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::SmagorinskyLilly)
         bottom = Operators.SetValue(C3(FT(0))),
     )
 
-    foreach_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, is_ρq_tot
+    foreach_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name, is_ρq_tot
         ᶜ∇ᵥρD∇χₜ = @. lazy(ᶜdivᵥ_ρχ(-(ᶠρ * ᶠD_smag * ᶠgradᵥ(specific(ᶜρχ, Y.c.ρ)))))
         @. ᶜρχₜ -= ᶜ∇ᵥρD∇χₜ
         # Rain and snow does not affect the mass

src/prognostic_equations/advection.jl

@@ -213,7 +213,7 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
     ᶜρ = Y.c.ρ
 
     # Full vertical advection of passive tracers (like liq, rai, etc) ...
-    foreach_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, is_ρq_tot
+    foreach_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name, is_ρq_tot
         if !is_ρq_tot
             ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
             vtt = vertical_transport(ᶜρ, ᶠu³, ᶜχ, float(dt), tracer_upwinding)

src/prognostic_equations/remaining_tendency.jl

@@ -166,7 +166,7 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t)
     @. Yₜ.c.ρe_tot += vst_ρe_tot
 
     # TODO: can we write this out explicitly?
-    foreach_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, is_ρq_tot
+    foreach_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name, is_ρq_tot
         ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
         vst_tracer = viscous_sponge_tendency_tracer(ᶜρ, ᶜχ, viscous_sponge)
         @. ᶜρχₜ += vst_tracer

src/prognostic_equations/surface_flux.jl

@@ -115,7 +115,7 @@ 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
-    foreach_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, is_ρq_tot
+    foreach_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name, is_ρq_tot
         ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
         if is_ρq_tot
             @. ρ_flux_χ = sfc_conditions.ρ_flux_q_tot

src/prognostic_equations/vertical_diffusion_boundary_layer.jl

@@ -100,8 +100,8 @@ function vertical_diffusion_boundary_layer_tendency!(
         top = Operators.SetValue(C3(0)),
         bottom = Operators.SetValue(C3(0)),
     )
-    foreach_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, is_ρq_tot
-        if ρχ_name in (:ρq_rai, :ρq_sno, :ρn_rai)
+    foreach_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name, is_ρq_tot
+        if ρχ_name in (@name(ρq_rai), @name(ρq_sno), @name(ρn_rai))
             @. ᶜK_h_scaled = α_vert_diff_tracer * ᶜK_h
         else
             @. ᶜK_h_scaled = ᶜK_h

src/utils/variable_manipulations.jl

@@ -85,16 +85,16 @@ 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`.
+- `f`: A function to apply to each scalar. It must have the signature `f
+  (ᶜρχₜ, ᶜρχ, ρχ_name, 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_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, is_ρq_tot
+foreach_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name, is_ρq_tot
     # Apply some operation, e.g., a sponge layer
     @. ᶜρχₜ += some_sponge_function(ᶜρχ)
     if is_ρq_tot
@@ -108,7 +108,7 @@ foreach_tracer(f::F, Yₜ, Y) where {F} =
         ᶜρχₜ = 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)
+        f(ᶜρχₜ, ᶜρχ, scalar_name, is_ρq_tot)
     end
 
 """