Use LazyBroadcast for large-scale advection tendency

src/cache/cache.jl

@@ -10,7 +10,6 @@ struct AtmosCache{
     SCRA,
     HYPE,
     PR,
-    LSAD,
     EXTFORCING,
     NONGW,
     ORGW,
@@ -53,7 +52,6 @@ struct AtmosCache{
 
     """Additional parameters used by the various tendencies"""
     precipitation::PR
-    large_scale_advection::LSAD
     external_forcing::EXTFORCING
     non_orographic_gravity_wave::NONGW
     orographic_gravity_wave::ORGW
@@ -156,7 +154,6 @@ function build_cache(Y, atmos, params, surface_setup, sim_info, aerosol_names)
 
     hyperdiff = hyperdiffusion_cache(Y, atmos)
     precipitation = precipitation_cache(Y, atmos)
-    large_scale_advection = large_scale_advection_cache(Y, atmos)
     external_forcing = external_forcing_cache(Y, atmos, params)
     non_orographic_gravity_wave = non_orographic_gravity_wave_cache(Y, atmos)
     orographic_gravity_wave = orographic_gravity_wave_cache(Y, atmos)
@@ -175,7 +172,6 @@ function build_cache(Y, atmos, params, surface_setup, sim_info, aerosol_names)
         scratch,
         hyperdiff,
         precipitation,
-        large_scale_advection,
         external_forcing,
         non_orographic_gravity_wave,
         orographic_gravity_wave,

src/prognostic_equations/forcing/large_scale_advection.jl

@@ -6,49 +6,44 @@ import Thermodynamics as TD
 import ClimaCore.Spaces as Spaces
 import ClimaCore.Fields as Fields
 
-large_scale_advection_cache(Y, atmos::AtmosModel) =
-    large_scale_advection_cache(Y, atmos.ls_adv)
-
-large_scale_advection_cache(Y, ls_adv::Nothing) = (;)
-function large_scale_advection_cache(Y, ls_adv::LargeScaleAdvection)
-    FT = Spaces.undertype(axes(Y.c))
-    return (; ᶜdqtdt_hadv = similar(Y.c, FT), ᶜdTdt_hadv = similar(Y.c, FT))
+function large_scale_advection_tendency_ρq_tot(
+    ᶜρ,
+    thermo_params,
+    ᶜts,
+    t,
+    ls_adv,
+)
+    ls_adv isa LargeScaleAdvection || return NullBroadcasted()
+    (; prof_dTdt, prof_dqtdt) = ls_adv
+    ᶜz = Fields.coordinate_field(axes(ᶜρ)).z
+    ᶜdqtdt_hadv = @lazy @. prof_dqtdt(thermo_params, ᶜts, t, ᶜz)
+    return @lazy @. ᶜρ * ᶜdqtdt_hadv
 end
 
-large_scale_advection_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing
-function large_scale_advection_tendency!(
-    Yₜ,
-    Y,
-    p,
+function large_scale_advection_tendency_ρe_tot(
+    ᶜρ,
+    thermo_params,
+    ᶜts,
     t,
-    ls_adv::LargeScaleAdvection,
+    ls_adv,
 )
+    ls_adv isa LargeScaleAdvection || return NullBroadcasted()
     (; prof_dTdt, prof_dqtdt) = ls_adv
-
-    thermo_params = CAP.thermodynamics_params(p.params)
-    ᶜts = p.precomputed.ᶜts
-    ᶜdqtdt_hadv = p.large_scale_advection.ᶜdqtdt_hadv
-    ᶜdTdt_hadv = p.large_scale_advection.ᶜdTdt_hadv
-    z = Fields.coordinate_field(axes(ᶜdqtdt_hadv)).z
-
-    @. ᶜdTdt_hadv = prof_dTdt(thermo_params, ᶜts, t, z)
-    @. ᶜdqtdt_hadv = prof_dqtdt(thermo_params, ᶜts, t, z)
-
+    z = Fields.coordinate_field(axes(ᶜρ)).z
+    ᶜdTdt_hadv = @lazy @. prof_dTdt(thermo_params, ᶜts, t, z)
+    ᶜdqtdt_hadv = @lazy @. prof_dqtdt(thermo_params, ᶜts, t, z)
     T_0 = TD.Parameters.T_0(thermo_params)
     Lv_0 = TD.Parameters.LH_v0(thermo_params)
     cv_v = TD.Parameters.cv_v(thermo_params)
     R_v = TD.Parameters.R_v(thermo_params)
-
-    @. Yₜ.c.ρq_tot += Y.c.ρ * ᶜdqtdt_hadv
     # TODO: should `hv` be a thermo function?
     #     (hv = cv_v * (ᶜT - T_0) + Lv_0 - R_v * T_0)
-    @. Yₜ.c.ρe_tot +=
-        Y.c.ρ * (
-            TD.cv_m(thermo_params, ᶜts) * ᶜdTdt_hadv +
-            (
-                cv_v * (TD.air_temperature(thermo_params, ᶜts) - T_0) + Lv_0 -
-                R_v * T_0
-            ) * ᶜdqtdt_hadv
-        )
+    return @lazy @. ᶜρ * (
+        TD.cv_m(thermo_params, ᶜts) * ᶜdTdt_hadv +
+        (
+            cv_v * (TD.air_temperature(thermo_params, ᶜts) - T_0) + Lv_0 -
+            R_v * T_0
+        ) * ᶜdqtdt_hadv
+    )
     return nothing
 end

src/prognostic_equations/remaining_tendency.jl

@@ -42,9 +42,10 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t)
     ᶠu₃ = Yₜ.f.u₃
     ᶜρ = Y.c.ρ
     (; forcing_type, moisture_model, rayleigh_sponge, viscous_sponge) = p.atmos
-    (; edmf_coriolis) = p.atmos
+    (; ls_adv, edmf_coriolis) = p.atmos
     (; params) = p
-    (; ᶜp, sfc_conditions) = p.precomputed
+    thermo_params = CAP.thermodynamics_params(params)
+    (; ᶜp, sfc_conditions, ᶜts) = p.precomputed
 
     vst_uₕ = viscous_sponge_tendency_uₕ(ᶜuₕ, viscous_sponge)
     vst_u₃ = viscous_sponge_tendency_u₃(ᶠu₃, viscous_sponge)
@@ -54,6 +55,8 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t)
     hs_tendency_uₕ = held_suarez_forcing_tendency_uₕ(hs_args...)
     hs_tendency_ρe_tot = held_suarez_forcing_tendency_ρe_tot(ᶜρ, hs_args...)
     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...)
 
     # TODO: fuse, once we fix
     #       https://github.com/CliMA/ClimaCore.jl/issues/2165
@@ -63,13 +66,11 @@ 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
@@ -78,9 +79,14 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t)
 
     subsidence_tendency!(Yₜ, Y, p, t, p.atmos.subsidence)
 
+    @. 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
+    end
+
     @. Yₜ.c.uₕ += edmf_cor_tend_uₕ
 
-    large_scale_advection_tendency!(Yₜ, Y, p, t, p.atmos.ls_adv)
     external_forcing_tendency!(Yₜ, Y, p, t, p.atmos.external_forcing)
 
     if p.atmos.sgs_adv_mode == Explicit()

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

test/coupler_compatibility.jl

@@ -76,7 +76,6 @@ const T2 = 290
         p.scratch,
         p.hyperdiff,
         p.precipitation,
-        p.large_scale_advection,
         p.external_forcing,
         p.non_orographic_gravity_wave,
         p.orographic_gravity_wave,