Fix column integrals for deep atmosphere

src/Operators/integrals.jl

@@ -5,37 +5,50 @@ import ClimaComms
 """
     column_integral_definite!(ϕ_top, ᶜ∂ϕ∂z, [ϕ_bot])
 
-Sets `ϕ_top```{}= \\int_{z_{bot}}^{z_{top}}\\,```ᶜ∂ϕ∂z```(z)\\,dz +{}```ϕ_bot`,
-where ``z_{bot}`` and ``z_{top}`` are the values of `z` at the bottom and top of
-the domain, respectively. The input `ᶜ∂ϕ∂z` should be a cell-center `Field` or
-`AbstractBroadcasted`, and the output `ϕ_top` should be a horizontal `Field`.
-The default value of `ϕ_bot` is 0.
+Sets `ϕ_top```{}= \\frac{1}{ΔA(z_{bot})}\\int_{z_{bot}}^{z_{top}}\\,
+```ᶜ∂ϕ∂z```(z)\\,ΔA(z)\\,dz +{}```ϕ_bot`, where ``z_{bot}`` and ``z_{top}`` are
+the values of `z` at the bottom and top of the domain, and where `ΔA` is the
+area differential `J/Δz`, with `J` denoting the metric Jacobian. The input
+`ᶜ∂ϕ∂z` should be a cell-center `Field` or `AbstractBroadcasted`, and the output
+`ϕ_top` should be a horizontal `Field`. The default value of `ϕ_bot` is 0.
 """
 function column_integral_definite!(ϕ_top, ᶜ∂ϕ∂z, ϕ_bot = rzero(eltype(ϕ_top)))
-    ᶜΔϕ = Base.Broadcast.broadcasted(⊠, ᶜ∂ϕ∂z, Fields.Δz_field(axes(ᶜ∂ϕ∂z)))
+    ᶜJ = Fields.local_geometry_field(axes(ᶜ∂ϕ∂z)).J
+    f_space = Spaces.face_space(axes(ᶜ∂ϕ∂z))
+    J_bot = Fields.level(Fields.local_geometry_field(f_space).J, half)
+    Δz_bot = Fields.level(Fields.Δz_field(f_space), half)
+    ΔA_bot = Base.broadcasted(/, J_bot, Δz_bot)
+    ᶜΔϕ = Base.broadcasted(⊠, ᶜ∂ϕ∂z, Base.broadcasted(/, ᶜJ, ΔA_bot))
     column_reduce!(⊞, ϕ_top, ᶜΔϕ; init = ϕ_bot)
 end
 
 """
     column_integral_indefinite!(ᶠϕ, ᶜ∂ϕ∂z, [ϕ_bot])
 
-Sets `ᶠϕ```(z) = \\int_{z_{bot}}^z\\,```ᶜ∂ϕ∂z```(z')\\,dz' +{}```ϕ_bot`, where
-``z_{bot}`` is the value of `z` at the bottom of the domain. The input `ᶜ∂ϕ∂z`
-should be a cell-center `Field` or `AbstractBroadcasted`, and the output `ᶠϕ`
-should be a cell-face `Field`. The default value of `ϕ_bot` is 0.
+Sets `ᶠϕ```(z) = \\frac{1}{ΔA(z_{bot})}\\int_{z_{bot}}^z\\,```ᶜ∂ϕ∂z```(z')\\,
+ΔA(z')\\,dz' +{}```ϕ_bot`, where ``z_{bot}`` is the value of `z` at the bottom
+of the domain, and where `ΔA` is the area differential `J/Δz`, with `J` denoting
+the metric Jacobian. The input `ᶜ∂ϕ∂z` should be a cell-center `Field` or
+`AbstractBroadcasted`, and the output `ᶠϕ` should be a cell-face `Field`. The
+default value of `ϕ_bot` is 0.
 
     column_integral_indefinite!(∂ϕ∂z, ᶠϕ, [ϕ_bot], [rtol])
 
-Sets
-`ᶠϕ```(z) = \\int_{z_{bot}}^z\\,```∂ϕ∂z```(```ᶠϕ```(z'), z')\\,dz' +{}```ϕ_bot`,
-where `∂ϕ∂z` can be any scalar-valued two-argument function. The output `ᶠϕ`
-satisfies `ᶜgradᵥ.(ᶠϕ) ≈ ∂ϕ∂z.(ᶜint.(ᶠϕ), ᶜz)`, where `ᶜgradᵥ = GradientF2C()`,
-`ᶜint = InterpolateF2C()`, and `ᶜz = Fields.coordinate_field(ᶜint.(ᶠϕ)).z`, and
-where the approximation is accurate to a relative tolerance of `rtol`. The
+Sets `ᶠϕ```(z) = \\frac{1}{ΔA(z_{bot})}\\int_{z_{bot}}^z\\,
+```∂ϕ∂z```(```ᶠϕ```(z'), z')\\,ΔA(z')\\,dz' +{}```ϕ_bot`, where `∂ϕ∂z` can be
+any scalar-valued two-argument function. When a shallow atmosphere approximation
+is used, `ΔA = ΔA_{bot}` at all values of `z`, and the output `ᶠϕ` satisfies
+`ᶜgradᵥ.(ᶠϕ) ≈ ∂ϕ∂z.(ᶜint.(ᶠϕ), ᶜz)` with a relative tolerance of `rtol`, where
+`ᶜgradᵥ = GradientF2C()` and `ᶜint = InterpolateF2C()`. When a deep atmosphere
+is used, the vertical gradient is replaced with an area-weighted gradient. The
 default value of `ϕ_bot` is 0, and the default value of `rtol` is 0.001.
 """
 function column_integral_indefinite!(ᶠϕ, ᶜ∂ϕ∂z, ϕ_bot = rzero(eltype(ᶠϕ)))
-    ᶜΔϕ = Base.Broadcast.broadcasted(⊠, ᶜ∂ϕ∂z, Fields.Δz_field(axes(ᶜ∂ϕ∂z)))
+    ᶜJ = Fields.local_geometry_field(axes(ᶜ∂ϕ∂z)).J
+    J_bot = Fields.level(Fields.local_geometry_field(ᶠϕ).J, half)
+    Δz_bot = Fields.level(Fields.Δz_field(ᶠϕ), half)
+    ΔA_bot = Base.broadcasted(/, J_bot, Δz_bot)
+    ᶜΔϕ = Base.broadcasted(⊠, ᶜ∂ϕ∂z, Base.broadcasted(/, ᶜJ, ΔA_bot))
     column_accumulate!(⊞, ᶠϕ, ᶜΔϕ; init = ϕ_bot)
 end
 function column_integral_indefinite!(
@@ -45,26 +58,23 @@ function column_integral_indefinite!(
     rtol = eltype(ᶠϕ)(0.001),
 ) where {F <: Function}
     device = ClimaComms.device(ᶠϕ)
-    face_space = axes(ᶠϕ)
-    center_space = if face_space isa Spaces.FaceFiniteDifferenceSpace
-        Spaces.CenterFiniteDifferenceSpace(face_space)
-    elseif face_space isa Spaces.FaceExtrudedFiniteDifferenceSpace
-        Spaces.CenterExtrudedFiniteDifferenceSpace(face_space)
-    else
-        error("output of column_integral_indefinite! must be on cell faces")
-    end
-    ᶜz = Fields.coordinate_field(center_space).z
-    ᶜΔz = Fields.Δz_field(center_space)
-    ᶜz_and_Δz = Base.Broadcast.broadcasted(tuple, ᶜz, ᶜΔz)
-    column_accumulate!(ᶠϕ, ᶜz_and_Δz; init = ϕ_bot) do ϕ_prev, (z, Δz)
-        residual(ϕ_new) = (ϕ_new - ϕ_prev) / Δz - ∂ϕ∂z((ϕ_prev + ϕ_new) / 2, z)
+    c_space = Spaces.center_space(axes(ᶠϕ))
+    ᶜz = Fields.coordinate_field(c_space).z
+    ᶜJ = Fields.local_geometry_field(c_space).J
+    J_bot = Fields.level(Fields.local_geometry_field(ᶠϕ).J, half)
+    Δz_bot = Fields.level(Fields.Δz_field(ᶠϕ), half)
+    ΔA_bot = Base.broadcasted(/, J_bot, Δz_bot)
+    ᶜz_and_Δz = Base.broadcasted(tuple, ᶜz, Base.broadcasted(/, ᶜJ, ΔA_bot))
+    column_accumulate!(ᶠϕ, ᶜz_and_Δz; init = ϕ_bot) do ϕ_prev, (z, weighted_Δz)
+        residual(ϕ_new) =
+            (ϕ_new - ϕ_prev) / weighted_Δz - ∂ϕ∂z((ϕ_prev + ϕ_new) / 2, z)
         (; converged, root) = RootSolvers.find_zero(
             residual,
             RootSolvers.NewtonsMethodAD(ϕ_prev),
             RootSolvers.CompactSolution(),
             RootSolvers.RelativeSolutionTolerance(rtol),
         )
-        ClimaComms.@assert device converged "∂ϕ∂z could not be integrated over \
+        ClimaComms.@assert device converged "unable to integrate through \
                                              z = $z with rtol set to $rtol"
         return root
     end

src/Spaces/extruded.jl

@@ -5,7 +5,8 @@
     ExtrudedFiniteDifferenceSpace(
         horizontal_space::AbstractSpace,
         vertical_space::FiniteDifferenceSpace,
-        hypsography::Grids.HypsographyAdaption = Grids.Flat(),
+        hypsography::Grids.HypsographyAdaption = Grids.Flat();
+        deep::Bool = false,
     )
 
 An extruded finite-difference space,
@@ -68,12 +69,14 @@ ExtrudedFiniteDifferenceSpace{S}(
 function ExtrudedFiniteDifferenceSpace(
     horizontal_space::AbstractSpace,
     vertical_space::FiniteDifferenceSpace,
-    hypsography::Grids.HypsographyAdaption = Grids.Flat(),
+    hypsography::Grids.HypsographyAdaption = Grids.Flat();
+    deep = false,
 )
     grid_space = Grids.ExtrudedFiniteDifferenceGrid(
         grid(horizontal_space),
         grid(vertical_space),
-        hypsography,
+        hypsography;
+        deep,
     )
     return ExtrudedFiniteDifferenceSpace(grid_space, vertical_space.staggering)
 end

test/Operators/integrals.jl

@@ -1,5 +1,6 @@
 using Test
 using JET
+import Random
 import CUDA # explicitly required due to JET
 import ClimaComms
 ClimaComms.@import_required_backends
@@ -152,18 +153,55 @@ function test_column_reduce_and_accumulate!(center_space)
     end
 end
 
-@testset "Integral operations unit tests" begin
+# @testset "Integral operations unit tests" begin
+#     device = ClimaComms.device()
+#     context = ClimaComms.SingletonCommsContext(device)
+#     for FT in (Float32, Float64)
+#         for space in (
+#             TU.ColumnCenterFiniteDifferenceSpace(FT; context),
+#             TU.CenterExtrudedFiniteDifferenceSpace(FT; context),
+#         )
+#             test_column_integral_definite!(space)
+#             test_column_integral_indefinite!(space)
+#             test_column_integral_indefinite_fn!(space)
+#             test_column_reduce_and_accumulate!(space)
+#         end
+#     end
+# end
+
+@testset "Shallow and deep integral consistency" begin
     device = ClimaComms.device()
     context = ClimaComms.SingletonCommsContext(device)
     for FT in (Float32, Float64)
-        for space in (
-            TU.ColumnCenterFiniteDifferenceSpace(FT; context),
-            TU.CenterExtrudedFiniteDifferenceSpace(FT; context),
-        )
-            test_column_integral_definite!(space)
-            test_column_integral_indefinite!(space)
-            test_column_integral_indefinite_fn!(space)
-            test_column_reduce_and_accumulate!(space)
+        kwargs = (; context, topography = true)
+        shallow_space =
+            TU.CenterExtrudedFiniteDifferenceSpace(FT; kwargs..., deep = false)
+        deep_space =
+            TU.CenterExtrudedFiniteDifferenceSpace(FT; kwargs..., deep = true)
+        shallow_field = zeros(shallow_space)
+        deep_field = zeros(deep_space)
+        shallow_face_field = zeros(center_to_face_space(shallow_space))
+        deep_face_field = zeros(center_to_face_space(deep_space))
+        shallow_surface_field = Fields.level(shallow_face_field, Fields.half)
+        deep_surface_field = Fields.level(deep_face_field, Fields.half)
+        shallow_surface_Δz = Fields.Δz_field(shallow_surface_field) ./ 2
+        deep_surface_Δz = Fields.Δz_field(deep_surface_field) ./ 2
+
+        Random.seed!(1) # ensures reproducibility
+        parent(shallow_field) .= rand.(FT)
+        parent(deep_field) .= parent(shallow_field)
+        column_integral_definite!(shallow_surface_field, shallow_field)
+        column_integral_definite!(deep_surface_field, deep_field)
+        shallow_split_sum = sum(shallow_surface_field ./ shallow_surface_Δz)
+        deep_split_sum = sum(deep_surface_field ./ deep_surface_Δz)
+        shallow_sum = sum(shallow_field)
+        deep_sum = sum(deep_field)
+
+        @test abs(shallow_split_sum - shallow_sum) / shallow_sum < 10 * eps(FT)
+        if FT == Float64
+            @test deep_split_sum == deep_sum # exact consistency with Float64!
+        else
+            @test abs(deep_split_sum - deep_sum) / deep_sum < 10 * eps(FT)
         end
     end
 end

test/TestUtilities/TestUtilities.jl

@@ -20,6 +20,7 @@ import ClimaCore.Meshes
 import ClimaCore.Spaces
 import ClimaCore.Topologies
 import ClimaCore.Domains
+import ClimaCore.Hypsography
 
 function PointSpace(
     ::Type{FT};
@@ -110,17 +111,20 @@ function CenterExtrudedFiniteDifferenceSpace(
     context = ClimaComms.SingletonCommsContext(),
     helem = 4,
     Nq = 4,
+    deep = false,
+    topography = false,
     horizontal_layout_type = DataLayouts.IJFH,
 ) where {FT}
     radius = FT(128)
-    zlim = (0, 1)
-    vertdomain = Domains.IntervalDomain(
-        Geometry.ZPoint{FT}(zlim[1]),
-        Geometry.ZPoint{FT}(zlim[2]);
+    zlim = (FT(0), FT(1))
+
+    vdomain = Domains.IntervalDomain(
+        Geometry.ZPoint(zlim[1]),
+        Geometry.ZPoint(zlim[2]);
         boundary_names = (:bottom, :top),
     )
-    vertmesh = Meshes.IntervalMesh(vertdomain, nelems = zelem)
-    vtopology = Topologies.IntervalTopology(context, vertmesh)
+    vmesh = Meshes.IntervalMesh(vdomain, nelems = zelem)
+    vtopology = Topologies.IntervalTopology(context, vmesh)
     vspace = Spaces.CenterFiniteDifferenceSpace(vtopology)
 
     hdomain = Domains.SphereDomain(radius)
@@ -129,16 +133,27 @@ function CenterExtrudedFiniteDifferenceSpace(
     quad = Quadratures.GLL{Nq}()
     hspace =
         Spaces.SpectralElementSpace2D(htopology, quad; horizontal_layout_type)
-    return Spaces.ExtrudedFiniteDifferenceSpace(hspace, vspace)
+
+    hypsography = if topography
+        # some non-trivial function of latitude and longitude
+        H = (zlim[2] - zlim[1]) / zelem
+        (; lat, long) = Fields.coordinate_field(hspace)
+        surface_elevation =
+            @. Geometry.ZPoint(H * (cosd(lat) + cosd(long) + 1))
+        Hypsography.LinearAdaption(surface_elevation)
+    else
+        Hypsography.Flat()
+    end
+    return Spaces.ExtrudedFiniteDifferenceSpace(
+        hspace,
+        vspace,
+        hypsography;
+        deep,
+    )
 end
 
-function FaceExtrudedFiniteDifferenceSpace(
-    ::Type{FT};
-    zelem = 10,
-    helem = 4,
-    context = ClimaComms.SingletonCommsContext(),
-) where {FT}
-    cspace = CenterExtrudedFiniteDifferenceSpace(FT; zelem, context, helem)
+function FaceExtrudedFiniteDifferenceSpace(::Type{FT}; kwargs...) where {FT}
+    cspace = CenterExtrudedFiniteDifferenceSpace(FT; kwargs...)
     return Spaces.FaceExtrudedFiniteDifferenceSpace(cspace)
 end