Put zmax into cache, for rayleigh sponge

Author: charleskawczynski

src/parameterized_tendencies/sponge/rayleigh_sponge.jl

@@ -14,7 +14,13 @@ import ClimaCore.Fields as Fields
 
 function rayleigh_sponge_tendency_uₕ(ᶜuₕ, s)
     s isa Nothing && return NullBroadcasted()
-    (; ᶜz, ᶠz) = z_coordinate_fields(axes(ᶜuₕ))
+    ᶜz = Fields.coordinate_field(Spaces.center_space(axes(ᶜuₕ))).z
+    ᶠz = Fields.coordinate_field(Spaces.face_space(axes(ᶜuₕ))).z
     zmax = z_max(axes(ᶠz))
+    return rayleigh_sponge_tendency_uₕ(ᶜuₕ, s, ᶜz, ᶠz, zmax)
+end
+
+function rayleigh_sponge_tendency_uₕ(ᶜuₕ, s, ᶜz, ᶠz, zmax)
+    s isa Nothing && return NullBroadcasted()
     return @. lazy(-β_rayleigh_uₕ(s, ᶜz, zmax) * ᶜuₕ)
 end

src/prognostic_equations/remaining_tendency.jl

@@ -69,10 +69,13 @@ sub_tend(∑tends, t) = lazy.(∑tends .- t)
 function ᶜremaining_tendency(::Val{:uₕ}, ᶜY, ᶠY, p, t)
     :uₕ in propertynames(ᶜY) || return ()
     ∑tendencies = zero(eltype(ᶜY.uₕ))
+    (; zmax) = p
     (; viscous_sponge, rayleigh_sponge) = p.atmos
+    ᶜz = Fields.coordinate_field(ᶜY).z
+    ᶠz = Fields.coordinate_field(ᶠY).z
     ᶜuₕ = ᶜY.uₕ
     ∑tendencies = add_tend(∑tendencies, viscous_sponge_tendency_uₕ(ᶜuₕ, viscous_sponge))
-    ∑tendencies = add_tend(∑tendencies, rayleigh_sponge_tendency_uₕ(ᶜuₕ, rayleigh_sponge))
+    ∑tendencies = add_tend(∑tendencies, rayleigh_sponge_tendency_uₕ(ᶜuₕ, rayleigh_sponge, ᶜz, ᶠz, zmax))
 
     return ∑tendencies
 end
@@ -255,6 +258,7 @@ NVTX.@annotate function remaining_tendency!(Yₜ, Yₜ_lim, Y, p, t)
     end
     (; moisture_model, viscous_sponge, precip_model) = p.atmos
     p_kernel = (;
+        zmax = Spaces.z_max(axes(Y.f)),
         atmos = p.atmos,
         params = p.params,
         dt = p.dt,