Merge pull request #3446 from CliMA/ck/elim_some_cache

Eliminate Rayleigh sponge cache

src/parameterized_tendencies/sponge/rayleigh_sponge.jl

@@ -10,20 +10,17 @@ rayleigh_sponge_cache(Y, atmos::AtmosModel) =
 rayleigh_sponge_cache(Y, ::Nothing) = (;)
 rayleigh_sponge_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing
 
-function rayleigh_sponge_cache(Y, rs::RayleighSponge)
-    FT = Spaces.undertype(axes(Y.c))
-    (; zd, α_uₕ, α_w) = rs
-    ᶜz = Fields.coordinate_field(Y.c).z
-    ᶠz = Fields.coordinate_field(Y.f).z
-    ᶜαₘ_uₕ = @. ifelse(ᶜz > zd, α_uₕ, $(FT(0)))
-    ᶠαₘ_w = @. ifelse(ᶠz > zd, α_w, $(FT(0)))
-    zmax = maximum(ᶠz)
-    ᶜβ_rayleigh_uₕ = @. ᶜαₘ_uₕ * sin($(FT(π)) / 2 * (ᶜz - zd) / (zmax - zd))^2
-    ᶠβ_rayleigh_w = @. ᶠαₘ_w * sin($(FT(π)) / 2 * (ᶠz - zd) / (zmax - zd))^2
-    return (; ᶜβ_rayleigh_uₕ, ᶠβ_rayleigh_w)
-end
+rayleigh_sponge_cache(Y, rs::RayleighSponge) = nothing
 
-function rayleigh_sponge_tendency!(Yₜ, Y, p, t, ::RayleighSponge)
-    (; ᶜβ_rayleigh_uₕ) = p.rayleigh_sponge
-    @. Yₜ.c.uₕ -= ᶜβ_rayleigh_uₕ * Y.c.uₕ
+αₘ(s::RayleighSponge{FT}, z, α) where {FT} = ifelse(z > s.zd, α, FT(0))
+ζ_rayleigh(s::RayleighSponge{FT}, z) where {FT} =
+    sin(FT(π) / 2 * (z - s.zd) / (s.zmax - s.zd))^2
+β_rayleigh_uₕ(s::RayleighSponge{FT}, z) where {FT} =
+    αₘ(s, z, s.α_uₕ) * ζ_rayleigh(s, z)
+β_rayleigh_w(s::RayleighSponge{FT}, z) where {FT} =
+    αₘ(s, z, s.α_w) * ζ_rayleigh(s, z)
+
+function rayleigh_sponge_tendency!(Yₜ, Y, p, t, s::RayleighSponge)
+    ᶜz = Fields.coordinate_field(Y.c).z
+    @. Yₜ.c.uₕ -= β_rayleigh_uₕ(s, ᶜz) * Y.c.uₕ
 end

src/prognostic_equations/implicit/implicit_solver.jl

@@ -451,10 +451,6 @@ NVTX.@annotate function Wfact!(A, Y, p, dtγ, t)
         p.dt,
         p.params,
         p.atmos,
-        (
-            p.atmos.rayleigh_sponge isa RayleighSponge ?
-            (; p.rayleigh_sponge.ᶠβ_rayleigh_w) : (;)
-        )...,
     )
 
     # Convert dtγ from a Float64 to an FT.
@@ -570,11 +566,14 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ)
     I_u₃ = DiagonalMatrixRow(one_C3xACT3)
     @. ∂ᶠu₃_err_∂ᶜuₕ =
         dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(-(ᶜkappa_m) * ᶜρ) ⋅ ∂ᶜK_∂ᶜuₕ
-    if p.atmos.rayleigh_sponge isa RayleighSponge
+    rs = p.atmos.rayleigh_sponge
+    ᶠz = Fields.coordinate_field(Y.f).z
+    if rs isa RayleighSponge
         @. ∂ᶠu₃_err_∂ᶠu₃ =
             dtγ * (
                 ᶠp_grad_matrix ⋅ DiagonalMatrixRow(-(ᶜkappa_m) * ᶜρ) ⋅
-                ∂ᶜK_∂ᶠu₃ + DiagonalMatrixRow(-p.ᶠβ_rayleigh_w * (one_C3xACT3,))
+                ∂ᶜK_∂ᶠu₃ +
+                DiagonalMatrixRow(-β_rayleigh_w(rs, ᶠz) * (one_C3xACT3,))
             ) - (I_u₃,)
     else
         @. ∂ᶠu₃_err_∂ᶠu₃ =
@@ -844,23 +843,26 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ)
             @. bdmr = ifelse(ᶜu₃ʲ.components.data.:1 > 0, bdmr_l, bdmr_r)
 
             @. ᶠtridiagonal_matrix_c3 = -(ᶠgradᵥ_matrix()) ⋅ bdmr
-            if p.atmos.rayleigh_sponge isa RayleighSponge
+            if rs isa RayleighSponge
                 @. ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ =
                     dtγ * (
                         ᶠtridiagonal_matrix_c3 ⋅
                         DiagonalMatrixRow(adjoint(CT3(Y.f.sgsʲs.:(1).u₃))) -
-                        DiagonalMatrixRow(p.ᶠβ_rayleigh_w * (one_C3xACT3,))
+                        DiagonalMatrixRow(
+                            β_rayleigh_w(rs, ᶠz) * (one_C3xACT3,),
+                        )
                     ) - (I_u₃,)
             else
                 @. ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ =
                     dtγ * ᶠtridiagonal_matrix_c3 ⋅
                     DiagonalMatrixRow(adjoint(CT3(Y.f.sgsʲs.:(1).u₃))) - (I_u₃,)
             end
-        elseif p.atmos.rayleigh_sponge isa RayleighSponge
+        elseif rs isa RayleighSponge
             ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ =
                 matrix[@name(f.sgsʲs.:(1).u₃), @name(f.sgsʲs.:(1).u₃)]
             @. ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ =
-                dtγ * -DiagonalMatrixRow(p.ᶠβ_rayleigh_w * (one_C3xACT3,)) -
+                dtγ *
+                -DiagonalMatrixRow(β_rayleigh_w(rs, ᶠz) * (one_C3xACT3,)) -
                 (I_u₃,)
         end
     end

src/prognostic_equations/implicit/implicit_tendency.jl

@@ -165,11 +165,13 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t)
     @. Yₜ.f.u₃ -= ᶠgradᵥ(ᶜp) / ᶠinterp(Y.c.ρ) + ᶠgradᵥ_ᶜΦ
 
     if rayleigh_sponge isa RayleighSponge
-        (; ᶠβ_rayleigh_w) = p.rayleigh_sponge
-        @. Yₜ.f.u₃ -= ᶠβ_rayleigh_w * Y.f.u₃
+        ᶠz = Fields.coordinate_field(Y.f).z
+        rs = rayleigh_sponge
+        @. Yₜ.f.u₃ -= β_rayleigh_w(rs, ᶠz) * Y.f.u₃
         if turbconv_model isa PrognosticEDMFX
             for j in 1:n
-                @. Yₜ.f.sgsʲs.:($$j).u₃ -= ᶠβ_rayleigh_w * Y.f.sgsʲs.:($$j).u₃
+                @. Yₜ.f.sgsʲs.:($$j).u₃ -=
+                    β_rayleigh_w(rs, ᶠz) * Y.f.sgsʲs.:($$j).u₃
             end
         end
     end

src/solver/model_getters.jl

@@ -158,13 +158,14 @@ end
 
 function get_rayleigh_sponge_model(parsed_args, params, ::Type{FT}) where {FT}
     rs_name = parsed_args["rayleigh_sponge"]
+    zmax = parsed_args["z_max"]
     return if rs_name in ("false", false)
         nothing
     elseif rs_name in ("true", true, "RayleighSponge")
         zd = params.zd_rayleigh
         α_uₕ = params.alpha_rayleigh_uh
         α_w = params.alpha_rayleigh_w
-        RayleighSponge{FT}(; zd, α_uₕ, α_w)
+        RayleighSponge{FT}(; zmax, zd, α_uₕ, α_w)
     else
         error("Uncaught rayleigh sponge model `$rs_name`.")
     end

src/solver/types.jl

@@ -126,10 +126,12 @@ Base.@kwdef struct ViscousSponge{FT} <: AbstractSponge
 end
 
 Base.@kwdef struct RayleighSponge{FT} <: AbstractSponge
+    zmax::FT
     zd::FT
     α_uₕ::FT
     α_w::FT
 end
+Base.Broadcast.broadcastable(x::RayleighSponge) = tuple(x)
 
 abstract type AbstractGravityWave end
 Base.@kwdef struct NonOrographyGravityWave{FT} <: AbstractGravityWave

test/parameterized_tendencies/sponge/rayleigh_sponge.jl

@@ -22,18 +22,16 @@ include("../../test_helpers.jl")
     FT = eltype(Y)
     ᶜYₜ = zero(Y)
     ### Component test begins here
-    rs = CA.RayleighSponge(; zd = FT(0), α_uₕ = FT(1), α_w = FT(1))
-    (; ᶜβ_rayleigh_uₕ) = CA.rayleigh_sponge_cache(Y, rs)
-    @test ᶜβ_rayleigh_uₕ == @. sin(FT(π) / 2 * z / zmax)^2
-    test_cache = (; rayleigh_sponge = (; ᶜβ_rayleigh_uₕ = ᶜβ_rayleigh_uₕ))
-    CA.rayleigh_sponge_tendency!(ᶜYₜ, Y, test_cache, FT(0), rs)
+    rs = CA.RayleighSponge(; zmax, zd = FT(0), α_uₕ = FT(1), α_w = FT(1))
+    @test CA.β_rayleigh_uₕ.(rs, z) == @. sin(FT(π) / 2 * z / zmax)^2
+    CA.rayleigh_sponge_tendency!(ᶜYₜ, Y, nothing, FT(0), rs)
     # Test that only required tendencies are affected
     for (var_name) in filter(x -> (x != :uₕ), propertynames(Y.c))
         @test ᶜYₜ.c.:($var_name) == zeros(axes(Y.c))
     end
     for (var_name) in propertynames(Y.f)
         @test ᶜYₜ.f.:($var_name) == zeros(axes(Y.f))
     end
-    @test ᶜYₜ.c.uₕ.components.data.:1 == -1 .* (ᶜβ_rayleigh_uₕ)
-    @test ᶜYₜ.c.uₕ.components.data.:2 == -1 .* (ᶜβ_rayleigh_uₕ)
+    @test ᶜYₜ.c.uₕ.components.data.:1 == -1 .* (CA.β_rayleigh_uₕ.(rs, z))
+    @test ᶜYₜ.c.uₕ.components.data.:2 == -1 .* (CA.β_rayleigh_uₕ.(rs, z))
 end

test/parameterized_tendencies/sponge/viscous_sponge.jl

@@ -19,11 +19,9 @@ include("../../test_helpers.jl")
     ᶜYₜ = Y .* FT(0)
     FT = eltype(Y)
     ### Component test begins here
-    rs = CA.RayleighSponge(; zd = FT(0), α_uₕ = FT(1), α_w = FT(1))
-    (; ᶜβ_rayleigh_uₕ) = CA.viscous_sponge_cache(Y, rs)
-    @test ᶜβ_rayleigh_uₕ == @. sin(FT(π) / 2 * z / zmax)^2
-    test_cache = (; viscous_sponge = (; ᶜβ_rayleigh_uₕ = ᶜβ_rayleigh_uₕ))
-    CA.viscous_sponge_tendency!(ᶜYₜ, Y, test_cache, FT(0), rs)
-    @test ᶜYₜ.c.uₕ.components.data.:1 == -1 .* (ᶜβ_rayleigh_uₕ)
-    @test ᶜYₜ.c.uₕ.components.data.:2 == -1 .* (ᶜβ_rayleigh_uₕ)
+    rs = CA.RayleighSponge(; zmax, zd = FT(0), α_uₕ = FT(1), α_w = FT(1))
+    @test CA.β_rayleigh_uₕ.(rs, z) == @. sin(FT(π) / 2 * z / zmax)^2
+    CA.viscous_sponge_tendency!(ᶜYₜ, Y, nothing, FT(0), rs)
+    @test ᶜYₜ.c.uₕ.components.data.:1 == -1 .* (CA.β_rayleigh_uₕ.(rs, z))
+    @test ᶜYₜ.c.uₕ.components.data.:2 == -1 .* (CA.β_rayleigh_uₕ.(rs, z))
 end