Merge #2281

2281: Cache: clean up viscous and rayleigh r=Sbozzolo a=Sbozzolo

- Establish the pattern: `cache(Y, atmos)`
- Clean up a couple of things

Co-authored-by: Gabriele Bozzola <gbozzola@caltech.edu>

src/cache/cache.jl

@@ -127,8 +127,10 @@ function additional_cache(
     end
 
     return merge(
-        rayleigh_sponge_cache(atmos.rayleigh_sponge, Y),
-        viscous_sponge_cache(atmos.viscous_sponge, Y),
+        (;
+            rayleigh_sponge = rayleigh_sponge_cache(Y, atmos),
+            viscous_sponge = viscous_sponge_cache(Y, atmos),
+        ),
         precipitation_cache(Y, precip_model),
         subsidence_cache(Y, atmos.subsidence),
         large_scale_advection_cache(Y, atmos.ls_adv),

src/cache/temporary_quantities.jl

@@ -1,9 +1,6 @@
 using ClimaCore: Fields
 using ClimaCore.Utilities: half
 
-# The model also depends on f_plane_coriolis_frequency(params)
-# This is a constant Coriolis frequency that is only used if space is flat
-
 # Fields used to store variables that only need to be used in a single function
 # but cannot be computed on the fly. Unlike the precomputed quantities, these
 # can be modified at any point, so they should never be assumed to be unchanged

src/parameterized_tendencies/sponge/rayleigh_sponge.jl

@@ -4,10 +4,13 @@
 
 import ClimaCore.Fields as Fields
 
-rayleigh_sponge_cache(::Nothing, Y) = NamedTuple()
+rayleigh_sponge_cache(Y, atmos::AtmosModel) =
+    rayleigh_sponge_cache(Y, atmos.rayleigh_sponge)
+
+rayleigh_sponge_cache(Y, ::Nothing) = (;)
 rayleigh_sponge_tendency!(Yₜ, Y, p, t, colidx, ::Nothing) = nothing
 
-function rayleigh_sponge_cache(rs::RayleighSponge, Y)
+function rayleigh_sponge_cache(Y, rs::RayleighSponge)
     FT = Spaces.undertype(axes(Y.c))
     (; zd, α_uₕ, α_w) = rs
     ᶜz = Fields.coordinate_field(Y.c).z
@@ -21,6 +24,6 @@ function rayleigh_sponge_cache(rs::RayleighSponge, Y)
 end
 
 function rayleigh_sponge_tendency!(Yₜ, Y, p, t, colidx, ::RayleighSponge)
-    (; ᶜβ_rayleigh_uₕ) = p
+    (; ᶜβ_rayleigh_uₕ) = p.rayleigh_sponge
     @. Yₜ.c.uₕ[colidx] -= ᶜβ_rayleigh_uₕ[colidx] * Y.c.uₕ[colidx]
 end

src/parameterized_tendencies/sponge/viscous_sponge.jl

@@ -6,10 +6,13 @@ import ClimaCore.Fields as Fields
 import ClimaCore.Geometry as Geometry
 import ClimaCore.Spaces as Spaces
 
-viscous_sponge_cache(::Nothing, Y) = NamedTuple()
+viscous_sponge_cache(Y, atmos::AtmosModel) =
+    viscous_sponge_cache(Y, atmos.viscous_sponge)
+
+viscous_sponge_cache(Y, ::Nothing) = (;)
 viscous_sponge_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing
 
-function viscous_sponge_cache(viscous_sponge::ViscousSponge, Y)
+function viscous_sponge_cache(Y, viscous_sponge::ViscousSponge)
     (; κ₂, zd) = viscous_sponge
     FT = Spaces.undertype(axes(Y.c))
     ᶜz = Fields.coordinate_field(Y.c).z
@@ -22,16 +25,34 @@ function viscous_sponge_cache(viscous_sponge::ViscousSponge, Y)
     return (; ᶜβ_viscous, ᶠβ_viscous)
 end
 
-function viscous_sponge_tendency!(Yₜ, Y, p, t, ::ViscousSponge)
-    (; ᶜβ_viscous, ᶠβ_viscous) = p
+add_viscous_sponge_energy_tendency!(Yₜ, Y, p, t) =
+    add_viscous_sponge_energy_tendency!(Yₜ, Y, p, t, p.atmos.energy_form)
+
+function add_viscous_sponge_energy_tendency!(Yₜ, Y, p, t, ::TotalEnergy)
+    (; ᶜβ_viscous) = p.viscous_sponge
+    (; ᶜh_tot) = p
+    ᶜρ = Y.c.ρ
+    @. Yₜ.c.ρe_tot += ᶜβ_viscous * wdivₕ(ᶜρ * gradₕ(ᶜh_tot))
+end
+
+function add_viscous_sponge_energy_tendency!(
+    Yₜ,
+    Y,
+    p,
+    t,
+    ::PotentialTemperature,
+)
+    (; ᶜβ_viscous) = p.viscous_sponge
     ᶜρ = Y.c.ρ
+    @. Yₜ.c.ρθ += ᶜβ_viscous * wdivₕ(ᶜρ * gradₕ(Y.c.ρθ / ᶜρ))
+end
+
+function viscous_sponge_tendency!(Yₜ, Y, p, t, ::ViscousSponge)
+    (; ᶜβ_viscous, ᶠβ_viscous) = p.viscous_sponge
     ᶜuₕ = Y.c.uₕ
-    if :ρθ in propertynames(Y.c)
-        @. Yₜ.c.ρθ += ᶜβ_viscous * wdivₕ(ᶜρ * gradₕ(Y.c.ρθ / ᶜρ))
-    elseif :ρe_tot in propertynames(Y.c)
-        (; ᶜh_tot) = p
-        @. Yₜ.c.ρe_tot += ᶜβ_viscous * wdivₕ(ᶜρ * gradₕ(ᶜh_tot))
-    end
+
+    add_viscous_sponge_energy_tendency!(Yₜ, Y, p, t)
+
     @. Yₜ.c.uₕ +=
         ᶜβ_viscous * (
             wgradₕ(divₕ(ᶜuₕ)) - Geometry.project(

src/prognostic_equations/edmfx_sgs_flux.jl

@@ -13,12 +13,11 @@ function edmfx_sgs_mass_flux_tendency!(
     turbconv_model::PrognosticEDMFX,
 )
 
-    FT = Spaces.undertype(axes(Y.c))
     n = n_mass_flux_subdomains(turbconv_model)
     (; edmfx_upwinding) = p.atmos.numerics
     (; ᶠu³, ᶜh_tot, ᶜspecific) = p
     (; ᶠu³ʲs) = p
-    (; ᶜρa⁰, ᶠu³⁰, ᶜu⁰, ᶜh_tot⁰, ᶜq_tot⁰) = p
+    (; ᶜρa⁰, ᶠu³⁰, ᶜh_tot⁰, ᶜq_tot⁰) = p
     (; dt) = p.simulation
     ᶜJ = Fields.local_geometry_field(Y.c).J
 
@@ -99,13 +98,10 @@ function edmfx_sgs_mass_flux_tendency!(
     FT = Spaces.undertype(axes(Y.c))
     n = n_mass_flux_subdomains(turbconv_model)
     (; edmfx_upwinding) = p.atmos.numerics
-    (; sfc_conditions) = p
-    (; ᶠu³, ᶜu, ᶜh_tot, ᶜspecific) = p
+    (; ᶠu³, ᶜh_tot, ᶜspecific) = p
     (; ᶜρaʲs, ᶠu³ʲs, ᶜh_totʲs, ᶜq_totʲs) = p
-    (; ᶜK_u, ᶜK_h) = p
     (; dt) = p.simulation
     ᶜJ = Fields.local_geometry_field(Y.c).J
-    ᶠgradᵥ = Operators.GradientC2F()
 
     if p.atmos.edmfx_sgs_mass_flux
         # energy
@@ -223,14 +219,9 @@ function edmfx_sgs_diffusive_flux_tendency!(
 )
 
     FT = Spaces.undertype(axes(Y.c))
-    n = n_mass_flux_subdomains(turbconv_model)
-    (; edmfx_upwinding) = p.atmos.numerics
     (; sfc_conditions) = p
-    (; ᶠu³, ᶜu, ᶜh_tot, ᶜspecific) = p
-    (; ᶜρaʲs, ᶠu³ʲs, ᶜh_totʲs, ᶜq_totʲs) = p
+    (; ᶜu, ᶜh_tot, ᶜspecific) = p
     (; ᶜK_u, ᶜK_h) = p
-    (; dt) = p.simulation
-    ᶜJ = Fields.local_geometry_field(Y.c).J
     ᶠgradᵥ = Operators.GradientC2F()
 
     if p.atmos.edmfx_sgs_diffusive_flux

src/prognostic_equations/implicit/implicit_solver.jl

@@ -221,12 +221,12 @@ function Wfact!(A, Y, p, dtγ, t)
             p.params,
             p.atmos,
             (energy_form isa TotalEnergy ? (; p.ᶜh_tot) : (;))...,
-            (rayleigh_sponge isa RayleighSponge ? (; p.ᶠβ_rayleigh_w) : (;))...,
+            (
+                rayleigh_sponge isa RayleighSponge ?
+                (; p.rayleigh_sponge.ᶠβ_rayleigh_w) : (;)
+            )...,
         )
 
-        (; energy_upwinding, tracer_upwinding, density_upwinding) =
-            p.atmos.numerics
-
         # Convert dtγ from a Float64 to an FT.
         FT = Spaces.undertype(axes(Y.c))
         dtγ′ = FT(dtγ)

src/prognostic_equations/implicit/implicit_tendency.jl

@@ -111,7 +111,7 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t, colidx)
         ) / ᶠinterp(Y.c.ρ[colidx])
 
     if rayleigh_sponge isa RayleighSponge
-        (; ᶠβ_rayleigh_w) = p
+        (; ᶠβ_rayleigh_w) = p.rayleigh_sponge
         @. Yₜ.f.u₃[colidx] -= ᶠβ_rayleigh_w[colidx] * Y.f.u₃[colidx]
         if turbconv_model isa PrognosticEDMFX
             for j in 1:n