Use LazyBroadcast for subsidence tendency

src/prognostic_equations/forcing/subsidence.jl

@@ -7,59 +7,35 @@ 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}(),
-        )
-    end
+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()
+    ᶠu³ = ᶠsubsidence³(axes(ᶜρ), subsidence)
+    ᶜχ = @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,10 +40,15 @@ 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
-    (; edmf_coriolis) = p.atmos
+    (; subsidence, edmf_coriolis) = p.atmos
     (; params) = p
-    (; ᶜp, sfc_conditions) = p.precomputed
+    thermo_params = CAP.thermodynamics_params(params)
+    (; ᶜp, sfc_conditions, ᶜts) = p.precomputed
+    ᶠspace = axes(Y.f)
+    ᶠz = Fields.coordinate_field(ᶠspace).z
+    ᶠlg = Fields.local_geometry_field(ᶠspace)
 
     vst_uₕ = viscous_sponge_tendency_uₕ(ᶜuₕ, viscous_sponge)
     vst_u₃ = viscous_sponge_tendency_u₃(ᶠu₃, viscous_sponge)
@@ -52,6 +57,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)
 
     # TODO: fuse, once we fix
@@ -62,20 +68,27 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t)
     @. Yₜ.c.ρe_tot += vst_ρe_tot
 
     # TODO: can we write this out explicitly?
-    if viscous_sponge isa ViscousSponge
-        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
-        end
+    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
     end
 
     # Held Suarez tendencies
     @. 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
+        bc_subsidence_ρq_tot = subsidence_tendency(Val(:q_tot), subs_args...)
+        @. Yₜ.c.ρq_tot += bc_subsidence_ρq_tot
+    end
+    if moisture_model isa NonEquilMoistModel
+        bc_subsidence_ρq_liq = subsidence_tendency(Val(:q_liq), subs_args...)
+        bc_subsidence_ρq_ice = subsidence_tendency(Val(:q_ice), subs_args...)
+        @. Yₜ.c.ρq_liq += bc_subsidence_ρq_liq
+        @. Yₜ.c.ρq_ice += bc_subsidence_ρq_ice
+    end
 
     @. Yₜ.c.uₕ += edmf_cor_tend_uₕ
 

src/solver/types.jl

@@ -7,9 +7,10 @@ import ClimaUtilities.ClimaArtifacts: @clima_artifact
 import LazyArtifacts
 
 abstract type AbstractMoistureModel end
+abstract type AbstractMoistModel <: AbstractMoistureModel end
 struct DryModel <: AbstractMoistureModel end
-struct EquilMoistModel <: AbstractMoistureModel end
-struct NonEquilMoistModel <: AbstractMoistureModel end
+struct EquilMoistModel <: AbstractMoistModel end
+struct NonEquilMoistModel <: AbstractMoistModel end
 
 abstract type AbstractPrecipitationModel end
 struct NoPrecipitation <: AbstractPrecipitationModel end