Use LazyBroadcast for subsidence tendency

Author: charleskawczynski

src/prognostic_equations/forcing/external_forcing.jl

@@ -201,6 +201,7 @@ function external_forcing_tendency!(
         ᶜinv_τ_wind,
         ᶜinv_τ_scalar,
     ) = p.external_forcing
+    ᶜρ = Y.c.ρ
 
     ᶜlg = Fields.local_geometry_field(Y.c)
     ᶜuₕ_nudge = p.scratch.ᶜtemp_C12
@@ -239,19 +240,21 @@ function external_forcing_tendency!(
     ᶠls_subsidence³ = p.scratch.ᶠtemp_CT3
     @. ᶠls_subsidence³ =
         ᶠinterp(ᶜls_subsidence * CT3(unit_basis_vector_data(CT3, ᶜlg)))
-    subsidence!(
-        Yₜ.c.ρe_tot,
-        Y.c.ρ,
+    Yₜ.c.ρe_tot .+= subsidence_tendency(
+        Val(:h_tot),
         ᶠls_subsidence³,
-        ᶜh_tot,
-        Val{:first_order}(),
+        thermo_params,
+        ᶜts,
+        ᶜρ,
+        ᶜρe_tot,
     )
-    subsidence!(
-        Yₜ.c.ρq_tot,
-        Y.c.ρ,
+    Yₜ.c.ρq_tot .+= subsidence_tendency(
+        Val(:q_tot),
         ᶠls_subsidence³,
-        ᶜspecific.q_tot,
-        Val{:first_order}(),
+        thermo_params,
+        ᶜts,
+        ᶜρ,
+        ᶜρe_tot,
     )
 
     # needed to address top boundary condition for forcings. Otherwise upper portion of domain is anomalously cold

src/prognostic_equations/forcing/subsidence.jl

@@ -3,63 +3,42 @@
 #####
 
 import Thermodynamics as TD
-import ClimaCore.Spaces as Spaces
 import ClimaCore.Fields as Fields
-import ClimaCore.Operators as Operators
-
-#####
-##### No subsidence
-#####
-
-subsidence_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing
 
 #####
 ##### Subsidence
 #####
 
-subsidence!(ᶜρχₜ, ᶜρ, ᶠu³, ᶜχ, ::Val{:none}) =
-    @. ᶜρχₜ -= ᶜρ * (ᶜsubdivᵥ(ᶠu³ * ᶠinterp(ᶜχ)) - ᶜχ * ᶜsubdivᵥ(ᶠu³))
-subsidence!(ᶜρχₜ, ᶜρ, ᶠu³, ᶜχ, ::Val{:first_order}) =
-    @. ᶜρχₜ -= ᶜρ * (ᶜsubdivᵥ(ᶠupwind1(ᶠu³, ᶜχ)) - ᶜχ * ᶜsubdivᵥ(ᶠu³))
-subsidence!(ᶜρχₜ, ᶜρ, ᶠu³, ᶜχ, ::Val{:third_order}) =
-    @. ᶜρχₜ -= ᶜρ * (ᶜsubdivᵥ(ᶠupwind3(ᶠu³, ᶜχ)) - ᶜχ * ᶜsubdivᵥ(ᶠu³))
-
-function subsidence_tendency!(Yₜ, Y, p, t, ::Subsidence)
-    (; moisture_model) = p.atmos
-    subsidence_profile = p.atmos.subsidence.prof
-    (; ᶜh_tot, ᶜspecific) = p.precomputed
-
-    ᶠz = Fields.coordinate_field(axes(Y.f)).z
-    ᶠlg = Fields.local_geometry_field(Y.f)
-    ᶠsubsidence³ = p.scratch.ᶠtemp_CT3
-    @. ᶠsubsidence³ =
-        subsidence_profile(ᶠz) * CT3(unit_basis_vector_data(CT3, ᶠlg))
-
-    # LS Subsidence
-    subsidence!(Yₜ.c.ρe_tot, Y.c.ρ, ᶠsubsidence³, ᶜh_tot, Val{:first_order}())
-    subsidence!(
-        Yₜ.c.ρq_tot,
-        Y.c.ρ,
-        ᶠsubsidence³,
-        ᶜspecific.q_tot,
-        Val{:first_order}(),
-    )
-    if moisture_model isa NonEquilMoistModel
-        subsidence!(
-            Yₜ.c.ρq_liq,
-            Y.c.ρ,
-            ᶠsubsidence³,
-            ᶜspecific.q_liq,
-            Val{:first_order}(),
-        )
-        subsidence!(
-            Yₜ.c.ρq_ice,
-            Y.c.ρ,
-            ᶠsubsidence³,
-            ᶜspecific.q_ice,
-            Val{:first_order}(),
-        )
+subsidence_tracer(::Val{:h_tot}, thermo_params, ts, ρ, ρe_tot) =
+    TD.total_specific_enthalpy(thermo_params, ts, ρe_tot / ρ)
+subsidence_tracer(::Val{:q_tot}, thermo_params, ts, ρ, ρe_tot) =
+    TD.total_specific_humidity(thermo_params, ts)
+subsidence_tracer(::Val{:q_liq}, thermo_params, ts, ρ, ρe_tot) =
+    TD.liquid_specific_humidity(thermo_params, ts)
+subsidence_tracer(::Val{:q_ice}, thermo_params, ts, ρ, ρe_tot) =
+    TD.ice_specific_humidity(thermo_params, ts)
+
+function subsidence_tendency(χname, subsidence, thermo_params, ᶜts, ᶜρ, ᶜρe_tot)
+    subsidence isa Nothing && return NullBroadcasted()
+    ᶠspace = Spaces.face_space(axes(ᶜρ))
+    ᶠu³ = if subsidence isa Subsidence
+        ᶠsubsidence³(ᶠspace, subsidence)
+    else
+        subsidence
     end
+    ᶜχ = @. lazy(subsidence_tracer(χname, thermo_params, ᶜts, ᶜρ, ᶜρe_tot))
+    return subsidence_tendency(ᶜρ, ᶜχ, ᶠu³, Val(:first_order))
+end
 
-    return nothing
+function ᶠsubsidence³(ᶠspace, subsidence::Subsidence)
+    ᶠz = Fields.coordinate_field(ᶠspace).z
+    ᶠlg = Fields.local_geometry_field(ᶠspace)
+    (; prof) = subsidence
+    return @. lazy(prof(ᶠz) * CT3(unit_basis_vector_data(CT3, ᶠlg)))
 end
+subsidence_tendency(ᶜρ, ᶜχ, ᶠu³, ::Val{:none}) =
+    @. lazy(- ᶜρ * (ᶜsubdivᵥ(ᶠu³ * ᶠinterp(ᶜχ)) - ᶜχ * ᶜsubdivᵥ(ᶠu³)))
+subsidence_tendency(ᶜρ, ᶜχ, ᶠu³, ::Val{:first_order}) =
+    @. lazy(- ᶜρ * (ᶜsubdivᵥ(ᶠupwind1(ᶠu³, ᶜχ)) - ᶜχ * ᶜsubdivᵥ(ᶠu³)))
+subsidence_tendency(ᶜρ, ᶜχ, ᶠu³, ::Val{:third_order}) =
+    @. lazy(- ᶜρ * (ᶜsubdivᵥ(ᶠupwind3(ᶠu³, ᶜχ)) - ᶜχ * ᶜsubdivᵥ(ᶠu³)))

src/prognostic_equations/remaining_tendency.jl

@@ -40,8 +40,9 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t)
     ᶜuₕ = Y.c.uₕ
     ᶠu₃ = Yₜ.f.u₃
     ᶜρ = Y.c.ρ
+    ᶜρe_tot = Y.c.ρe_tot
     (; forcing_type, moisture_model, rayleigh_sponge, viscous_sponge) = p.atmos
-    (; ls_adv, edmf_coriolis) = p.atmos
+    (; ls_adv, subsidence, edmf_coriolis) = p.atmos
     (; params) = p
     thermo_params = CAP.thermodynamics_params(params)
     (; ᶜp, sfc_conditions, ᶜts) = p.precomputed
@@ -53,6 +54,7 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t)
     hs_args = (ᶜuₕ, ᶜp, params, sfc_conditions.ts, moisture_model, forcing_type)
     hs_tendency_uₕ = held_suarez_forcing_tendency_uₕ(hs_args...)
     hs_tendency_ρe_tot = held_suarez_forcing_tendency_ρe_tot(ᶜρ, hs_args...)
+    subs_args = (subsidence, thermo_params, ᶜts, ᶜρ, ᶜρe_tot)
     edmf_cor_tend_uₕ = edmf_coriolis_tendency_uₕ(ᶜuₕ, edmf_coriolis)
     lsa_args = (ᶜρ, thermo_params, ᶜts, t, ls_adv)
     bc_lsa_tend_ρe_tot = large_scale_advection_tendency_ρe_tot(lsa_args...)
@@ -76,12 +78,17 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t)
     @. Yₜ.c.uₕ += hs_tendency_uₕ
     @. Yₜ.c.ρe_tot += hs_tendency_ρe_tot
 
-    subsidence_tendency!(Yₜ, Y, p, t, p.atmos.subsidence)
+    if moisture_model isa AbstractMoistModel
+        Yₜ.c.ρq_tot .+= subsidence_tendency(Val(:q_tot), subs_args...)
+    end
+    if moisture_model isa NonEquilMoistModel
+        Yₜ.c.ρq_liq .+= subsidence_tendency(Val(:q_liq), subs_args...)
+        Yₜ.c.ρq_ice .+= subsidence_tendency(Val(:q_ice), subs_args...)
+    end
 
     @. Yₜ.c.ρe_tot += bc_lsa_tend_ρe_tot
     if moisture_model isa AbstractMoistModel
-        bc_lsa_tend_ρq_tot = large_scale_advection_tendency_ρq_tot(lsa_args...)
-        @. Yₜ.c.ρq_tot += bc_lsa_tend_ρq_tot
+        Yₜ.c.ρq_tot .+= large_scale_advection_tendency_ρq_tot(lsa_args...)
     end
 
     @. Yₜ.c.uₕ += edmf_cor_tend_uₕ