Slim method arguments

src/parameterized_tendencies/sponge/rayleigh_sponge.jl

@@ -13,8 +13,9 @@ import ClimaCore.Fields as Fields
 β_rayleigh_w(s::RayleighSponge{FT}, z, zmax) where {FT} =
     αₘ(s, z, s.α_w) * ζ_rayleigh(s, z, zmax)
 
-rayleigh_sponge_tendency_uₕ(ᶜuₕ, ᶠz, ᶜz, s::Nothing) = (zero(eltype(ᶜuₕ)),)
-function rayleigh_sponge_tendency_uₕ(ᶜuₕ, ᶠz, ᶜz, s::RayleighSponge)
+rayleigh_sponge_tendency_uₕ(ᶜuₕ, s::Nothing) = (zero(eltype(ᶜuₕ)),)
+function rayleigh_sponge_tendency_uₕ(ᶜuₕ, s::RayleighSponge)
+    (; ᶜz, ᶠz) = z_coordinate_fields(axes(ᶜuₕ))
     zmax = z_max(axes(ᶠz))
     return @lazy @. -β_rayleigh_uₕ(s, ᶜz, zmax) * ᶜuₕ
 end

src/parameterized_tendencies/sponge/sponge_tendencies.jl

@@ -7,19 +7,25 @@ import ClimaCore.Fields as Fields
 import ClimaCore.Geometry as Geometry
 import ClimaCore.Spaces as Spaces
 
+function z_coordinate_fields(space::Spaces.AbstractSpace)
+    ᶜz = Fields.coordinate_field(Spaces.center_space(space)).z
+    ᶠz = Fields.coordinate_field(Spaces.face_space(space)).z
+    return (; ᶜz, ᶠz)
+end
+
+
 function sponge_tendencies!(Yₜ, Y, p, t)
     rs, vs = p.atmos.rayleigh_sponge, p.atmos.viscous_sponge
     (; ᶜh_tot, ᶜspecific) = p.precomputed
     ᶜuₕ = Y.c.uₕ
     ᶠu₃ = Yₜ.f.u₃
     ᶜρ = Y.c.ρ
-    ᶜz = Fields.coordinate_field(Y.c).z
-    ᶠz = Fields.coordinate_field(Y.f).z
+    (; ᶜz, ᶠz) = z_coordinate_fields(axes(ᶜuₕ))
     zmax = z_max(axes(ᶠz))
-    vst_uₕ = viscous_sponge_tendency_uₕ(ᶜuₕ, ᶜz, ᶠz, vs)
-    vst_u₃ = viscous_sponge_tendency_u₃(ᶠu₃, ᶠz, vs)
-    vst_ρe_tot = viscous_sponge_tendency_ρe_tot(ᶜρ, ᶜz, ᶠz, ᶜh_tot, vs)
-    rst_uₕ = rayleigh_sponge_tendency_uₕ(ᶜuₕ, ᶠz, ᶜz, rs)
+    vst_uₕ = viscous_sponge_tendency_uₕ(ᶜuₕ, vs)
+    vst_u₃ = viscous_sponge_tendency_u₃(ᶠu₃, vs)
+    vst_ρe_tot = viscous_sponge_tendency_ρe_tot(ᶜρ, ᶜh_tot, vs)
+    rst_uₕ = rayleigh_sponge_tendency_uₕ(ᶜuₕ, rs)
 
     @. Yₜ.c.uₕ += vst_uₕ + rst_uₕ
     @. Yₜ.f.u₃.components.data.:1 += vst_u₃

src/parameterized_tendencies/sponge/viscous_sponge.jl

@@ -13,8 +13,9 @@ import ClimaCore.Spaces as Spaces
 β_viscous(s::ViscousSponge{FT}, z, zmax) where {FT} =
     αₘ(s, z) * ζ_viscous(s, z, zmax)
 
-viscous_sponge_tendency_uₕ(ᶜuₕ, ᶜz, ᶠz, ::Nothing) = (zero(eltype(ᶜuₕ)),)
-function viscous_sponge_tendency_uₕ(ᶜuₕ, ᶜz, ᶠz, s::ViscousSponge)
+viscous_sponge_tendency_uₕ(ᶜuₕ, ::Nothing) = (zero(eltype(ᶜuₕ)),)
+function viscous_sponge_tendency_uₕ(ᶜuₕ, s::ViscousSponge)
+    (; ᶜz, ᶠz) = z_coordinate_fields(axes(ᶜuₕ))
     zmax = z_max(axes(ᶠz))
     return @lazy @. β_viscous(s, ᶜz, zmax) * (
         wgradₕ(divₕ(ᶜuₕ)) - Geometry.project(
@@ -24,16 +25,17 @@ function viscous_sponge_tendency_uₕ(ᶜuₕ, ᶜz, ᶠz, s::ViscousSponge)
     )
 end
 
-viscous_sponge_tendency_u₃(u₃, ᶠz, ::Nothing) =
+viscous_sponge_tendency_u₃(u₃, ::Nothing) =
     (zero(eltype(u₃.components.data.:1)),)
-function viscous_sponge_tendency_u₃(u₃, ᶠz, s::ViscousSponge)
+function viscous_sponge_tendency_u₃(u₃, s::ViscousSponge)
+    (; ᶠz) = z_coordinate_fields(axes(u₃))
     zmax = z_max(axes(ᶠz))
     return @lazy @. β_viscous(s, ᶠz, zmax) * wdivₕ(gradₕ(u₃.components.data.:1))
 end
 
-viscous_sponge_tendency_ρe_tot(ᶜρ, ᶜz, ᶠz, ᶜh_tot, ::Nothing) =
-    (zero(eltype(ᶜρ)),)
-function viscous_sponge_tendency_ρe_tot(ᶜρ, ᶜz, ᶠz, ᶜh_tot, s::ViscousSponge)
+viscous_sponge_tendency_ρe_tot(ᶜρ, ᶜh_tot, ::Nothing) = (zero(eltype(ᶜρ)),)
+function viscous_sponge_tendency_ρe_tot(ᶜρ, ᶜh_tot, s::ViscousSponge)
+    (; ᶜz, ᶠz) = z_coordinate_fields(axes(ᶜρ))
     zmax = z_max(axes(ᶠz))
     return @lazy @. β_viscous(s, ᶜz, zmax) * wdivₕ(ᶜρ * gradₕ(ᶜh_tot))
 end

test/parameterized_tendencies/sponge/rayleigh_sponge.jl

@@ -36,12 +36,12 @@ using ClimaCore.CommonSpaces
     ### Component test begins here
     rs = CA.RayleighSponge(; zd = FT(0), α_uₕ = FT(1), α_w = FT(1))
     expected = @. sin(FT(π) / 2 * ᶜz / zmax)^2
-    computed = CA.rayleigh_sponge_tendency_uₕ(ᶜuₕ, ᶠz, ᶜz, rs)
+    computed = CA.rayleigh_sponge_tendency_uₕ(ᶜuₕ, rs)
     @test CA.β_rayleigh_uₕ.(rs, ᶜz, zmax) == expected
     @test materialize(computed) == .-expected .* ᶜuₕ
 
     # Test when not using a Rayleigh sponge.
-    computed = CA.rayleigh_sponge_tendency_uₕ(ᶜuₕ, ᶠz, ᶜz, nothing)
+    computed = CA.rayleigh_sponge_tendency_uₕ(ᶜuₕ, nothing)
     expected = @. ᶜuₕ .* 0
     @test eltype(computed) == eltype(expected)
     @. ᶜuₕ += computed # test that it can broadcast