Fix broadcasting

Author: charleskawczynski

src/cache/precipitation_precomputed_quantities.jl

@@ -34,31 +34,31 @@ function compute_precip_velocities(ᶜY, ᶠY, p, t)
     cmp = CAP.microphysics_1m_params(p.params)
 
     # compute the precipitation terminal velocity [m/s]
-    ᶜwᵣ = CM1.terminal_velocity(
-        cmp.pr,
-        cmp.tv.rain,
-        ᶜY.ρ,
-        max(zero(ᶜY.ρ), ᶜY.ρq_rai / ᶜY.ρ),
-    )
-    ᶜwₛ = CM1.terminal_velocity(
+    ᶜwᵣ = @. lazy(CM1.terminal_velocity(
+            cmp.pr,
+            cmp.tv.rain,
+            ᶜY.ρ,
+            max(zero(ᶜY.ρ), ᶜY.ρq_rai / ᶜY.ρ),
+        ))
+    ᶜwₛ = @. lazy(CM1.terminal_velocity(
         cmp.ps,
         cmp.tv.snow,
         ᶜY.ρ,
         max(zero(ᶜY.ρ), ᶜY.ρq_sno / ᶜY.ρ),
-    )
+    ))
     # compute sedimentation velocity for cloud condensate [m/s]
-    ᶜwₗ = CMNe.terminal_velocity(
+    ᶜwₗ = @. lazy(CMNe.terminal_velocity(
         cmc.liquid,
         cmc.Ch2022.rain,
         ᶜY.ρ,
         max(zero(ᶜY.ρ), ᶜY.ρq_liq / ᶜY.ρ),
-    )
-    ᶜwᵢ = CMNe.terminal_velocity(
+    ))
+    ᶜwᵢ = @. lazy(CMNe.terminal_velocity(
         cmc.ice,
         cmc.Ch2022.small_ice,
         ᶜY.ρ,
         max(zero(ᶜY.ρ), ᶜY.ρq_ice / ᶜY.ρ),
-    )
+    ))
     return (ᶜwᵣ, ᶜwₛ, ᶜwₗ, ᶜwᵢ)
 end
 
@@ -73,12 +73,12 @@ function compute_ᶜwₜqₜ(ᶜY, ᶠY, p, t,
 )
     (ᶜwᵣ, ᶜwₛ, ᶜwₗ, ᶜwᵢ) = compute_precip_velocities(ᶜY, ᶠY, p, t)
     # compute their contributions to energy and total water advection
-    return Geometry.WVector(
+    return @. lazy(Geometry.WVector(
             ᶜwₗ * ᶜY.ρq_liq +
             ᶜwᵢ * ᶜY.ρq_ice +
             ᶜwᵣ * ᶜY.ρq_rai +
             ᶜwₛ * ᶜY.ρq_sno,
-        ) / ᶜY.ρ
+        ) / ᶜY.ρ)
 end
 
 compute_ᶜwₕhₜ(ᶜY, ᶠY, p, t) =

src/prognostic_equations/remaining_tendency.jl

@@ -10,12 +10,20 @@ NVTX.@annotate function hyperdiffusion_tendency!(Yₜ, Yₜ_lim, Y, p, t)
     apply_hyperdiffusion_tendency!(Yₜ, Y, p, t)
 end
 
-function prognostic_nt(::Val{names}; tends...) where {names}
-    nt_ordered = NamedTuple{names}(rzero(values(tends)))
-    nt_values = NamedTuple{keys(tends)}(values(tends))
-    return merge(nt_ordered, nt_values)
+@generated function sorted_nt(::Val{snames}, ::Val{unames}, vals...) where {snames,unames}
+    svals_exprs = []
+    for sn in snames
+        i = findfirst(un -> un == sn, unames)::Int
+        push!(svals_exprs, :(getfield(vals, $i)))
+    end
+    return quote
+        NamedTuple{snames}(($(svals_exprs...),))
+    end
 end
 
+prognostic_nt(::Val{names}, ::Val{K}, vals...) where {names, K} =
+    sorted_nt(Val(names), Val(K), vals...)
+
 function ᶜremaining_tendency(ᶜY, ᶠY, p, t)
     names = propertynames(ᶜY)
     tends = (;
@@ -30,22 +38,23 @@ function ᶜremaining_tendency(ᶜY, ᶠY, p, t)
         ᶜremaining_tendency_sgs⁰(ᶜY, ᶠY, p, t)...,
         ᶜremaining_tendency_sgsʲs(ᶜY, ᶠY, p, t)...,
     )
-    return lazy.(prognostic_nt.(Val(names); tends...))
+    return lazy.(prognostic_nt.(Val(names), Val(keys(tends)), values(tends)...))
 end
 function ᶠremaining_tendency(ᶜY, ᶠY, p, t)
     names = propertynames(ᶠY)
     tends = (;
         ᶠremaining_tendency_u₃(ᶜY, ᶠY, p, t)...,
         ᶠremaining_tendency_sgsʲs(ᶜY, ᶠY, p, t)...,
     )
-    return lazy.(prognostic_nt.(Val(names); tends...))
+    return lazy.(prognostic_nt.(Val(names), Val(keys(tends)), values(tends)...))
 end
 using ClimaCore.RecursiveApply: rzero
 function ᶜremaining_tendency_ρ(ᶜY, ᶠY, p, t)
     :ρ in propertynames(ᶜY) || return ()
     ∑tendencies = zero(eltype(ᶜY.ρ))
     ᶜJ = Fields.local_geometry_field(ᶜY).J
     ᶠJ = Fields.local_geometry_field(ᶠY).J
+    ᶜρ = ᶜY.ρ
 
     if !(p.atmos.moisture_model isa DryModel)
         ᶜwₜqₜ = compute_ᶜwₜqₜ(ᶜY, ᶠY, p, t)
@@ -178,7 +187,8 @@ function ᶠremaining_tendency_sgsʲs(ᶜY, ᶠY, p, t)
 end
 
 
-water_adv(ᶜρ, ᶜJ, ᶠJ, ᶜχ) =
+water_adv(ᶜρ, ᶜJ, ᶠJ, ᶜχ::Real) = zero(eltype(ᶜρ))
+water_adv(ᶜρ, ᶜJ, ᶠJ, ᶜχ) = # only valid when ᶜχ is a field
     @. lazy(ᶜprecipdivᵥ(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ * ᶠright_bias(-(ᶜχ))))
 
 function surface_velocity_full(ᶠu₃, ᶠuₕ³)