Fuse sponge operations, using LazyBroadcast

Project.toml

@@ -21,6 +21,7 @@ FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
 Insolation = "e98cc03f-d57e-4e3c-b70c-8d51efe9e0d8"
 Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
 LazyArtifacts = "4af54fe1-eca0-43a8-85a7-787d91b784e3"
+LazyBroadcast = "9dccce8e-a116-406d-9fcc-a88ed4f510c8"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"

examples/Manifest.toml

@@ -1,8 +1,8 @@
 # This file is machine-generated - editing it directly is not advised
 
-julia_version = "1.10.7"
+julia_version = "1.10.8"
 manifest_format = "2.0"
-project_hash = "13bc74579bbfd1b4bab2e5d6cddd05f9a01dba06"
+project_hash = "db1ce831cbee482854057f6b6a17adf9f6688802"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "72af59f5b8f09faee36b4ec48e014a79210f2f4f"
@@ -336,7 +336,7 @@ version = "0.5.12"
     Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 
 [[deps.ClimaAtmos]]
-deps = ["ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Dates", "DiffEqBase", "FastGaussQuadrature", "Insolation", "Interpolations", "LazyArtifacts", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "RRTMGP", "Random", "SciMLBase", "StaticArrays", "Statistics", "SurfaceFluxes", "Thermodynamics", "UnrolledUtilities", "YAML"]
+deps = ["ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Dates", "DiffEqBase", "FastGaussQuadrature", "Insolation", "Interpolations", "LazyArtifacts", "LazyBroadcast", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "RRTMGP", "Random", "SciMLBase", "StaticArrays", "Statistics", "SurfaceFluxes", "Thermodynamics", "UnrolledUtilities", "YAML"]
 path = ".."
 uuid = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
 version = "0.28.2"
@@ -1385,6 +1385,11 @@ version = "0.1.17"
 deps = ["Artifacts", "Pkg"]
 uuid = "4af54fe1-eca0-43a8-85a7-787d91b784e3"
 
+[[deps.LazyBroadcast]]
+git-tree-sha1 = "6d2586dfde8b989198181fe21b8b6d0a42787ab8"
+uuid = "9dccce8e-a116-406d-9fcc-a88ed4f510c8"
+version = "1.0.0"
+
 [[deps.LazyModules]]
 git-tree-sha1 = "a560dd966b386ac9ae60bdd3a3d3a326062d3c3e"
 uuid = "8cdb02fc-e678-4876-92c5-9defec4f444e"

examples/Project.toml

@@ -33,6 +33,7 @@ JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7"
 LazyArtifacts = "4af54fe1-eca0-43a8-85a7-787d91b784e3"
+LazyBroadcast = "9dccce8e-a116-406d-9fcc-a88ed4f510c8"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"

src/ClimaAtmos.jl

@@ -111,6 +111,7 @@ include(
 )
 include(joinpath("parameterized_tendencies", "sponge", "rayleigh_sponge.jl"))
 include(joinpath("parameterized_tendencies", "sponge", "viscous_sponge.jl"))
+include(joinpath("parameterized_tendencies", "sponge", "sponge_tendencies.jl"))
 include(
     joinpath(
         "parameterized_tendencies",

src/parameterized_tendencies/sponge/rayleigh_sponge.jl

@@ -2,10 +2,9 @@
 ##### Rayleigh sponge
 #####
 
+import LazyBroadcast: @lazy
 import ClimaCore.Fields as Fields
 
-rayleigh_sponge_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing
-
 αₘ(s::RayleighSponge{FT}, z, α) where {FT} = ifelse(z > s.zd, α, FT(0))
 ζ_rayleigh(s::RayleighSponge{FT}, z, zmax) where {FT} =
     sin(FT(π) / 2 * (z - s.zd) / (zmax - s.zd))^2
@@ -14,8 +13,8 @@ rayleigh_sponge_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing
 β_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
-    zmax = z_max(axes(Y.f))
-    @. Yₜ.c.uₕ -= β_rayleigh_uₕ(s, ᶜz, zmax) * Y.c.uₕ
+rayleigh_sponge_tendency_uₕ(ᶜuₕ, ᶠz, ᶜz, s::Nothing) = (zero(eltype(ᶜuₕ)),)
+function rayleigh_sponge_tendency_uₕ(ᶜuₕ, ᶠz, ᶜz, s::RayleighSponge)
+    zmax = z_max(axes(ᶠz))
+    return @lazy @. -β_rayleigh_uₕ(s, ᶜz, zmax) * ᶜuₕ
 end

src/parameterized_tendencies/sponge/sponge_tendencies.jl

@@ -0,0 +1,34 @@
+#####
+##### Sponge tendencies
+#####
+
+import LazyBroadcast: @lazy
+import ClimaCore.Fields as Fields
+import ClimaCore.Geometry as Geometry
+import ClimaCore.Spaces as Spaces
+
+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
+    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)
+
+    @. Yₜ.c.uₕ += vst_uₕ + rst_uₕ
+    @. Yₜ.f.u₃.components.data.:1 += vst_u₃
+    @. Yₜ.c.ρe_tot += vst_ρe_tot
+
+    # TODO: can we write this out explicitly?
+    for (ᶜρχₜ, ᶜχ, χ_name) in matching_subfields(Yₜ.c, ᶜspecific)
+        χ_name == :e_tot && continue
+        @. ᶜρχₜ += β_viscous(vs, ᶜz, zmax) * wdivₕ(ᶜρ * gradₕ(ᶜχ))
+        @. Yₜ.c.ρ += β_viscous(vs, ᶜz, zmax) * wdivₕ(ᶜρ * gradₕ(ᶜχ))
+    end
+end

src/parameterized_tendencies/sponge/viscous_sponge.jl

@@ -2,39 +2,38 @@
 ##### Viscous sponge
 #####
 
+import LazyBroadcast: @lazy
 import ClimaCore.Fields as Fields
 import ClimaCore.Geometry as Geometry
 import ClimaCore.Spaces as Spaces
 
-
-viscous_sponge_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing
-
 αₘ(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, 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
+viscous_sponge_tendency_uₕ(ᶜuₕ, ᶜz, ᶠz, ::Nothing) = (zero(eltype(ᶜuₕ)),)
+function viscous_sponge_tendency_uₕ(ᶜuₕ, ᶜz, ᶠz, s::ViscousSponge)
     zmax = z_max(axes(ᶠz))
-    @. Yₜ.c.uₕ +=
-        β_viscous(s, ᶜz, zmax) * (
-            wgradₕ(divₕ(ᶜuₕ)) - Geometry.project(
-                Geometry.Covariant12Axis(),
-                wcurlₕ(Geometry.project(Geometry.Covariant3Axis(), curlₕ(ᶜuₕ))),
-            )
+    return @lazy @. β_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(s, ᶠz, zmax) * wdivₕ(gradₕ(Y.f.u₃.components.data.:1))
+    )
+end
 
-    @. 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(s, ᶜz, zmax) * wdivₕ(Y.c.ρ * gradₕ(ᶜχ))
-        @. Yₜ.c.ρ += β_viscous(s, ᶜz, zmax) * wdivₕ(Y.c.ρ * gradₕ(ᶜχ))
-    end
+viscous_sponge_tendency_u₃(u₃, ᶠz, ::Nothing) =
+    (zero(eltype(u₃.components.data.:1)),)
+function viscous_sponge_tendency_u₃(u₃, ᶠz, s::ViscousSponge)
+    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)
+    zmax = z_max(axes(ᶠz))
+    return @lazy @. β_viscous(s, ᶜz, zmax) * wdivₕ(ᶜρ * gradₕ(ᶜh_tot))
 end

src/prognostic_equations/remaining_tendency.jl

@@ -23,10 +23,8 @@ NVTX.@annotate function remaining_tendency!(Yₜ, Yₜ_lim, Y, p, t)
 end
 
 NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t)
-    viscous_sponge_tendency!(Yₜ, Y, p, t, p.atmos.viscous_sponge)
-
+    sponge_tendencies!(Yₜ, Y, p, t)
     # Vertical tendencies
-    rayleigh_sponge_tendency!(Yₜ, Y, p, t, p.atmos.rayleigh_sponge)
     forcing_tendency!(Yₜ, Y, p, t, p.atmos.forcing_type)
     subsidence_tendency!(Yₜ, Y, p, t, p.atmos.subsidence)
     edmf_coriolis_tendency!(Yₜ, Y, p, t, p.atmos.edmf_coriolis)

test/parameterized_tendencies/sponge/rayleigh_sponge.jl

@@ -5,37 +5,44 @@ 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 ClimaCore.CommonSpaces
+using ClimaCore: Spaces, Fields, Geometry, ClimaCore
 using Test
+using Base.Broadcast: materialize
+
+pkgversion(ClimaCore) < v"0.14.20" && exit() # CommonSpaces
+using ClimaCore.CommonSpaces
 
-include("../../test_helpers.jl")
 ### Common Objects ###
 @testset "Rayleigh-sponge functions" begin
-    ### Boilerplate default integrator objects
-    config = CA.AtmosConfig(
-        Dict("initial_condition" => "DryBaroclinicWave");
-        job_id = "sponge1",
+    FT = Float64
+    ᶜspace = ExtrudedCubedSphereSpace(
+        FT;
+        z_elem = 10,
+        z_min = 0,
+        z_max = 1,
+        radius = 10,
+        h_elem = 10,
+        n_quad_points = 4,
+        staggering = CellCenter(),
     )
-    (; Y) = generate_test_simulation(config)
-    zmax = maximum(CC.Fields.coordinate_field(Y.f).z)
-    z = CC.Fields.coordinate_field(Y.c).z
-    Y.c.uₕ.components.data.:1 .= ones(axes(Y.c))
-    Y.c.uₕ.components.data.:2 .= ones(axes(Y.c))
-    FT = eltype(Y)
-    ᶜYₜ = zero(Y)
+    ᶠspace = Spaces.face_space(ᶜspace)
+    ᶠz = Fields.coordinate_field(ᶠspace).z
+    ᶜz = Fields.coordinate_field(ᶜspace).z
+    zmax = maximum(ᶠz)
+    ᶜuₕ = map(z -> zero(Geometry.Contravariant12Vector{eltype(z)}), ᶜz)
+    @. ᶜuₕ.components.data.:1 = 1
+    @. ᶜuₕ.components.data.:2 = 1
     ### Component test begins here
     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))
-        @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 .* (CA.β_rayleigh_uₕ.(rs, z, zmax))
-    @test ᶜYₜ.c.uₕ.components.data.:2 == -1 .* (CA.β_rayleigh_uₕ.(rs, z, zmax))
+    expected = @. sin(FT(π) / 2 * ᶜz / zmax)^2
+    computed = CA.rayleigh_sponge_tendency_uₕ(ᶜuₕ, ᶠz, ᶜz, 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)
+    expected = @. ᶜuₕ .* 0
+    @test eltype(computed) == eltype(expected)
+    @. ᶜuₕ += computed # test that it can broadcast
 end

test/parameterized_tendencies/sponge/viscous_sponge.jl

@@ -5,30 +5,33 @@ using Revise; include("test/parameterized_tendencies/sponge/viscous_sponge.jl")
 using ClimaComms
 ClimaComms.@import_required_backends
 import ClimaAtmos as CA
-import ClimaCore
-using ClimaCore: Spaces, Grids, Fields
-if pkgversion(ClimaCore) ≥ v"0.14.20"
-    using ClimaCore.CommonGrids
-    using Test
+using ClimaCore.CommonSpaces
+using ClimaCore: Spaces, Fields, Geometry, ClimaCore
+using Test
+using Base.Broadcast: materialize
 
-    ### 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
+pkgversion(ClimaCore) < v"0.14.20" && exit() # CommonSpaces
+using ClimaCore.CommonSpaces
+
+### Common Objects ###
+@testset "Viscous-sponge functions" begin
+    FT = Float64
+    ᶜspace = ExtrudedCubedSphereSpace(
+        FT;
+        z_elem = 10,
+        z_min = 0,
+        z_max = 1,
+        radius = 10,
+        h_elem = 10,
+        n_quad_points = 4,
+        staggering = CellCenter(),
+    )
+    ᶠspace = Spaces.face_space(ᶜspace)
+    ᶜz = Fields.coordinate_field(ᶜspace).z
+    ᶠz = Fields.coordinate_field(ᶠspace).z
+    zmax = maximum(ᶠz)
+    ### 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