Merge pull request #3447 from CliMA/ck/elim_some_cache

Eliminate ViscousSponge cache, fix unit test

src/ClimaAtmos.jl

@@ -3,6 +3,7 @@ module ClimaAtmos
 using NVTX
 import Thermodynamics as TD
 
+include("compat.jl")
 include(joinpath("parameters", "Parameters.jl"))
 import .Parameters as CAP
 

src/cache/cache.jl

@@ -14,8 +14,6 @@ struct AtmosCache{
     SCRA,
     HYPE,
     DSS,
-    RS,
-    VS,
     PR,
     LSAD,
     EXTFORCING,
@@ -76,8 +74,6 @@ struct AtmosCache{
     do_dss::DSS
 
     """Additional parameters used by the various tendencies"""
-    rayleigh_sponge::RS
-    viscous_sponge::VS
     precipitation::PR
     large_scale_advection::LSAD
     external_forcing::EXTFORCING
@@ -183,8 +179,6 @@ function build_cache(Y, atmos, params, surface_setup, sim_info, aerosol_names)
         (start_date, params, atmos.ozone, aerosol_names, atmos.insolation) : ()
 
     hyperdiff = hyperdiffusion_cache(Y, atmos)
-    rayleigh_sponge = rayleigh_sponge_cache(Y, atmos)
-    viscous_sponge = viscous_sponge_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)
@@ -211,8 +205,6 @@ function build_cache(Y, atmos, params, surface_setup, sim_info, aerosol_names)
         scratch,
         hyperdiff,
         do_dss,
-        rayleigh_sponge,
-        viscous_sponge,
         precipitation,
         large_scale_advection,
         external_forcing,

src/compat.jl

@@ -0,0 +1,22 @@
+import ClimaCore
+import ClimaCore: Domains, Spaces, Topologies
+
+# To allow for backwards compatibility of ClimaCore:
+if pkgversion(ClimaCore) < v"0.14.18"
+    """
+        z_max(::ClimaCore.Spaces.AbstractSpace)
+
+    The domain maximum along the z-direction.
+    """
+    function z_max end
+
+    z_max(domain::Domains.IntervalDomain) = domain.coord_max.z
+    function z_max(space::Spaces.AbstractSpace)
+        mesh = Topologies.mesh(Spaces.vertical_topology(space))
+        domain = Topologies.domain(mesh)
+        return z_max(domain)
+    end
+
+else
+    z_max(s::Spaces.AbstractSpace) = Spaces.z_max(s)
+end

src/parameterized_tendencies/sponge/rayleigh_sponge.jl

@@ -4,23 +4,18 @@
 
 import ClimaCore.Fields as Fields
 
-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, ::Nothing) = nothing
 
-rayleigh_sponge_cache(Y, rs::RayleighSponge) = nothing
-
 αₘ(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)
+ζ_rayleigh(s::RayleighSponge{FT}, z, zmax) where {FT} =
+    sin(FT(π) / 2 * (z - s.zd) / (zmax - s.zd))^2
+β_rayleigh_uₕ(s::RayleighSponge{FT}, z, zmax) where {FT} =
+    αₘ(s, z, s.α_uₕ) * ζ_rayleigh(s, z, zmax)
+β_rayleigh_w(s::RayleighSponge{FT}, z, zmax) where {FT} =
+    αₘ(s, z, s.α_w) * ζ_rayleigh(s, z, zmax)
 
 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ₕ
+    zmax = z_max(axes(Y.f))
+    @. Yₜ.c.uₕ -= β_rayleigh_uₕ(s, ᶜz, zmax) * Y.c.uₕ
 end

src/parameterized_tendencies/sponge/viscous_sponge.jl

@@ -6,43 +6,35 @@ import ClimaCore.Fields as Fields
 import ClimaCore.Geometry as Geometry
 import ClimaCore.Spaces as Spaces
 
-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(Y, viscous_sponge::ViscousSponge)
-    (; κ₂, zd) = viscous_sponge
-    FT = Spaces.undertype(axes(Y.c))
-    ᶜz = Fields.coordinate_field(Y.c).z
-    ᶠz = Fields.coordinate_field(Y.f).z
-    ᶜαₘ = @. ifelse(ᶜz > zd, κ₂, FT(0))
-    ᶠαₘ = @. ifelse(ᶠz > zd, κ₂, FT(0))
-    zmax = maximum(ᶠz)
-    ᶜβ_viscous = @. ᶜαₘ * sin(FT(π) / 2 * (ᶜz - zd) / (zmax - zd))^2
-    ᶠβ_viscous = @. ᶠαₘ * sin(FT(π) / 2 * (ᶠz - zd) / (zmax - zd))^2
-    return (; ᶜβ_viscous, ᶠβ_viscous)
-end
+αₘ(s::ViscousSponge{FT}, z) where {FT} = ifelse(z > s.zd, s.κ₂, FT(0))
+ζ_viscous(s::ViscousSponge{FT}, z, zmax) where {FT} =
+    sin(FT(π) / 2 * (z - s.zd) / (zmax - s.zd))^2
+β_viscous(s::ViscousSponge{FT}, z, zmax) where {FT} =
+    αₘ(s, z) * ζ_viscous(s, z, zmax)
 
-function viscous_sponge_tendency!(Yₜ, Y, p, t, ::ViscousSponge)
-    (; ᶜβ_viscous, ᶠβ_viscous) = p.viscous_sponge
+function viscous_sponge_tendency!(Yₜ, Y, p, t, s::ViscousSponge)
     (; ᶜh_tot, ᶜspecific) = p.precomputed
     ᶜuₕ = Y.c.uₕ
+    ᶜz = Fields.coordinate_field(Y.c).z
+    ᶠz = Fields.coordinate_field(Y.f).z
+    zmax = z_max(axes(ᶠz))
     @. Yₜ.c.uₕ +=
-        ᶜβ_viscous * (
+        β_viscous(s, ᶜz, zmax) * (
             wgradₕ(divₕ(ᶜuₕ)) - Geometry.project(
                 Geometry.Covariant12Axis(),
                 wcurlₕ(Geometry.project(Geometry.Covariant3Axis(), curlₕ(ᶜuₕ))),
             )
         )
     @. Yₜ.f.u₃.components.data.:1 +=
-        ᶠβ_viscous * wdivₕ(gradₕ(Y.f.u₃.components.data.:1))
+        β_viscous(s, ᶠz, zmax) * wdivₕ(gradₕ(Y.f.u₃.components.data.:1))
 
-    @. Yₜ.c.ρe_tot += ᶜβ_viscous * wdivₕ(Y.c.ρ * gradₕ(ᶜh_tot))
+    @. Yₜ.c.ρe_tot += β_viscous(s, ᶜz, zmax) * wdivₕ(Y.c.ρ * gradₕ(ᶜh_tot))
     for (ᶜρχₜ, ᶜχ, χ_name) in matching_subfields(Yₜ.c, ᶜspecific)
         χ_name == :e_tot && continue
-        @. ᶜρχₜ += ᶜβ_viscous * wdivₕ(Y.c.ρ * gradₕ(ᶜχ))
-        @. Yₜ.c.ρ += ᶜβ_viscous * wdivₕ(Y.c.ρ * gradₕ(ᶜχ))
+        @. ᶜρχₜ += β_viscous(s, ᶜz, zmax) * wdivₕ(Y.c.ρ * gradₕ(ᶜχ))
+        @. Yₜ.c.ρ += β_viscous(s, ᶜz, zmax) * wdivₕ(Y.c.ρ * gradₕ(ᶜχ))
     end
 end

src/prognostic_equations/implicit/implicit_solver.jl

@@ -568,12 +568,13 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ)
         dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(-(ᶜkappa_m) * ᶜρ) ⋅ ∂ᶜK_∂ᶜuₕ
     rs = p.atmos.rayleigh_sponge
     ᶠz = Fields.coordinate_field(Y.f).z
+    zmax = z_max(axes(Y.f))
     if rs isa RayleighSponge
         @. ∂ᶠu₃_err_∂ᶠu₃ =
             dtγ * (
                 ᶠp_grad_matrix ⋅ DiagonalMatrixRow(-(ᶜkappa_m) * ᶜρ) ⋅
                 ∂ᶜK_∂ᶠu₃ +
-                DiagonalMatrixRow(-β_rayleigh_w(rs, ᶠz) * (one_C3xACT3,))
+                DiagonalMatrixRow(-β_rayleigh_w(rs, ᶠz, zmax) * (one_C3xACT3,))
             ) - (I_u₃,)
     else
         @. ∂ᶠu₃_err_∂ᶠu₃ =
@@ -849,7 +850,7 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ)
                         ᶠtridiagonal_matrix_c3 ⋅
                         DiagonalMatrixRow(adjoint(CT3(Y.f.sgsʲs.:(1).u₃))) -
                         DiagonalMatrixRow(
-                            β_rayleigh_w(rs, ᶠz) * (one_C3xACT3,),
+                            β_rayleigh_w(rs, ᶠz, zmax) * (one_C3xACT3,),
                         )
                     ) - (I_u₃,)
             else
@@ -862,8 +863,9 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ)
                 matrix[@name(f.sgsʲs.:(1).u₃), @name(f.sgsʲs.:(1).u₃)]
             @. ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ =
                 dtγ *
-                -DiagonalMatrixRow(β_rayleigh_w(rs, ᶠz) * (one_C3xACT3,)) -
-                (I_u₃,)
+                -DiagonalMatrixRow(
+                    β_rayleigh_w(rs, ᶠz, zmax) * (one_C3xACT3,),
+                ) - (I_u₃,)
         end
     end
 

src/prognostic_equations/implicit/implicit_tendency.jl

@@ -166,12 +166,13 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t)
 
     if rayleigh_sponge isa RayleighSponge
         ᶠz = Fields.coordinate_field(Y.f).z
+        zmax = z_max(axes(Y.f))
         rs = rayleigh_sponge
-        @. Yₜ.f.u₃ -= β_rayleigh_w(rs, ᶠz) * Y.f.u₃
+        @. Yₜ.f.u₃ -= β_rayleigh_w(rs, ᶠz, zmax) * Y.f.u₃
         if turbconv_model isa PrognosticEDMFX
             for j in 1:n
                 @. Yₜ.f.sgsʲs.:($$j).u₃ -=
-                    β_rayleigh_w(rs, ᶠz) * Y.f.sgsʲs.:($$j).u₃
+                    β_rayleigh_w(rs, ᶠz, zmax) * Y.f.sgsʲs.:($$j).u₃
             end
         end
     end

src/solver/model_getters.jl

@@ -158,14 +158,13 @@ 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}(; zmax, zd, α_uₕ, α_w)
+        RayleighSponge{FT}(; zd, α_uₕ, α_w)
     else
         error("Uncaught rayleigh sponge model `$rs_name`.")
     end

src/solver/types.jl

@@ -120,18 +120,17 @@ diffuse_momentum(::FriersonDiffusion{DM}) where {DM} = DM
 diffuse_momentum(::Nothing) = false
 
 abstract type AbstractSponge end
+Base.Broadcast.broadcastable(x::AbstractSponge) = tuple(x)
 Base.@kwdef struct ViscousSponge{FT} <: AbstractSponge
     zd::FT
     κ₂::FT
 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/coupler_compatibility.jl

@@ -80,8 +80,6 @@ const T2 = 290
         p.scratch,
         p.hyperdiff,
         p.do_dss,
-        p.rayleigh_sponge,
-        p.viscous_sponge,
         p.precipitation,
         p.large_scale_advection,
         p.external_forcing,

test/parameterized_tendencies/sponge/rayleigh_sponge.jl

@@ -1,14 +1,18 @@
+#=
+julia --project=examples
+using Revise; include("test/parameterized_tendencies/sponge/rayleigh_sponge.jl")
+=#
 using ClimaComms
 ClimaComms.@import_required_backends
 import ClimaAtmos as CA
 import SurfaceFluxes as SF
 import ClimaAtmos.Parameters as CAP
 import ClimaCore as CC
+using Test
 
 include("../../test_helpers.jl")
 ### Common Objects ###
-@testset begin
-    "Rayleigh-sponge functions"
+@testset "Rayleigh-sponge functions" begin
     ### Boilerplate default integrator objects
     config = CA.AtmosConfig(
         Dict("initial_condition" => "DryBaroclinicWave");
@@ -22,8 +26,8 @@ include("../../test_helpers.jl")
     FT = eltype(Y)
     ᶜYₜ = zero(Y)
     ### Component test begins here
-    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
+    rs = CA.RayleighSponge(; zd = FT(0), α_uₕ = FT(1), α_w = FT(1))
+    @test CA.β_rayleigh_uₕ.(rs, z, zmax) == @. 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))
@@ -32,6 +36,6 @@ include("../../test_helpers.jl")
     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 .* (CA.β_rayleigh_uₕ.(rs, z))
-    @test ᶜYₜ.c.uₕ.components.data.:2 == -1 .* (CA.β_rayleigh_uₕ.(rs, z))
+    @test ᶜYₜ.c.uₕ.components.data.:1 == -1 .* (CA.β_rayleigh_uₕ.(rs, z, zmax))
+    @test ᶜYₜ.c.uₕ.components.data.:2 == -1 .* (CA.β_rayleigh_uₕ.(rs, z, zmax))
 end

test/parameterized_tendencies/sponge/viscous_sponge.jl

@@ -1,27 +1,34 @@
+#=
+julia --project=examples
+using Revise; include("test/parameterized_tendencies/sponge/viscous_sponge.jl")
+=#
 using ClimaComms
 ClimaComms.@import_required_backends
 import ClimaAtmos as CA
-import SurfaceFluxes as SF
-import ClimaAtmos.Parameters as CAP
-import ClimaCore as CC
-include("../../test_helpers.jl")
+import ClimaCore
+using ClimaCore: Spaces, Grids, Fields
+if pkgversion(ClimaCore) ≥ v"0.14.18"
+    using ClimaCore.CommonGrids
+    using Test
 
-### Common Objects ###
-@testset begin
-    "Rayleigh-sponge functions"
-    ### Boilerplate default integrator objects
-    config = CA.AtmosConfig(Dict("initial_condition" => "DryBaroclinicWave"))
-    (; Y) = generate_test_simulation(config)
-    z = CC.Fields.coordinate_field(Y.c).z
-    zmax = maximum(CC.Fields.coordinate_field(Y.f).z)
-    Y.c.uₕ.components.data.:1 .= ones(axes(Y.c))
-    Y.c.uₕ.components.data.:2 .= ones(axes(Y.c))
-    ᶜYₜ = Y .* FT(0)
-    FT = eltype(Y)
-    ### Component test begins here
-    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))
+    ### Common Objects ###
+    @testset "Viscous-sponge functions" begin
+        grid = ExtrudedCubedSphereGrid(;
+            z_elem = 10,
+            z_min = 0,
+            z_max = 1,
+            radius = 10,
+            h_elem = 10,
+            n_quad_points = 4,
+        )
+        cspace = Spaces.ExtrudedFiniteDifferenceSpace(grid, Grids.CellCenter())
+        fspace = Spaces.FaceExtrudedFiniteDifferenceSpace(cspace)
+        z = Fields.coordinate_field(cspace).z
+        zmax = maximum(Fields.coordinate_field(fspace).z)
+        FT = typeof(zmax)
+        ### Component test begins here
+        s = CA.ViscousSponge{FT}(; zd = 0, κ₂ = 1)
+        @test CA.β_viscous.(s, z, zmax) == @. ifelse(z > s.zd, s.κ₂, FT(0)) *
+                 sin(FT(π) / 2 * (z - s.zd) / (zmax - s.zd))^2
+    end
 end

test/runtests.jl

@@ -18,6 +18,7 @@ using Test
 @safetestset "surface albedo tests" begin @time include("surface_albedo.jl") end
 @safetestset "Radiation interface tests" begin @time include("rrtmgp_interface.jl") end
 @safetestset "Sponge interface tests" begin @time include("parameterized_tendencies/sponge/rayleigh_sponge.jl") end
+@safetestset "Sponge interface tests" begin @time include("parameterized_tendencies/sponge/viscous_sponge.jl") end
 @safetestset "Precipitation interface tests" begin @time include("parameterized_tendencies/microphysics/precipitation.jl") end
 @safetestset "Model getters" begin @time include("solver/model_getters.jl") end
 @safetestset "Topography tests" begin @time include("topography.jl") end