Implement and test vertical mass borrowing limiters

NEWS.md

@@ -4,6 +4,9 @@ ClimaCore.jl Release Notes
 main
 -------
 
+ - A new limiter, `VerticalMassBorrowingLimiter`, was added. Which redistributes all negative values from a given field while preserving mass.
+  PR [2084](https://github.com/CliMA/ClimaCore.jl/pull/2084)
+
  - We've added new convenience constructors for spaces PR [2082](https://github.com/CliMA/ClimaCore.jl/pull/2082). Here are links to the new constructors:
    - [ExtrudedCubedSphereSpace]()
    - [CubedSphereSpace]()

docs/refs.bib

@@ -257,3 +257,14 @@ @article{Wiin1967
   year = {1967},
   publisher = {Wiley Online Library}
 }
+
+@article{zhang2018impact,
+  title={Impact of numerical choices on water conservation in the E3SM Atmosphere Model version 1 (EAMv1)},
+  author={Zhang, Kai and Rasch, Philip J and Taylor, Mark A and Wan, Hui and Leung, Ruby and Ma, Po-Lun and Golaz, Jean-Christophe and Wolfe, Jon and Lin, Wuyin and Singh, Balwinder and others},
+  journal={Geoscientific Model Development},
+  volume={11},
+  number={5},
+  pages={1971--1988},
+  year={2018},
+  publisher={Copernicus Publications G{\"o}ttingen, Germany}
+}

docs/src/api.md

@@ -356,14 +356,15 @@ Hypsography.ref_z_to_physical_z
 The limiters supertype is
 ```@docs
 Limiters.AbstractLimiter
+Limiters.QuasiMonotoneLimiter
+Limiters.VerticalMassBorrowingLimiter
 ```
 
 This class of flux-limiters is applied only in the horizontal direction (on spectral advection operators).
 
 
 ### Interfaces
 ```@docs
-Limiters.QuasiMonotoneLimiter
 Limiters.compute_bounds!
 Limiters.apply_limiter!
 ```

src/Limiters/Limiters.jl

@@ -19,5 +19,6 @@ abstract type AbstractLimiter end
 
 # implementations
 include("quasimonotone.jl")
+include("vertical_mass_borrowing_limiter.jl")
 
 end # end module

src/Limiters/vertical_mass_borrowing_limiter.jl

@@ -0,0 +1,133 @@
+import .DataLayouts as DL
+
+"""
+    VerticalMassBorrowingLimiter(f::Fields.Field, q_min)
+
+A vertical-only mass borrowing limiter.
+
+The mass borrower borrows tracer mass from an adjacent, lower layer.
+It conserves the total tracer mass and can avoid negative tracers.
+
+At level k, it will first borrow the mass from the layer k+1 (lower level).
+If the mass is not sufficient in layer k+1, it will borrow mass from
+layer k+2. The borrower will proceed this process until the bottom layer.
+If the tracer mass in the bottom layer goes negative, it will repeat the
+process from the bottom to the top. In this way, the borrower works for
+any shape of mass profiles.
+
+This code was adapted from [E3SM](https://github.com/E3SM-Project/E3SM/blob/2c377c5ec9a5585170524b366ad85074ab1b1a5c/components/eam/src/physics/cam/massborrow.F90)
+
+References:
+ - [zhang2018impact](@cite)
+"""
+struct VerticalMassBorrowingLimiter{F, T}
+    bmass::F
+    ic::F
+    q_min::T
+end
+function VerticalMassBorrowingLimiter(f::Fields.Field, q_min)
+    bmass = similar(Spaces.level(f, 1))
+    ic = similar(Spaces.level(f, 1))
+    return VerticalMassBorrowingLimiter(bmass, ic, q_min)
+end
+
+
+"""
+    apply_limiter!(q::Fields.Field, ρ::Fields.Field, lim::VerticalMassBorrowingLimiter)
+
+Apply the vertical mass borrowing
+limiter `lim` to field `q`, given
+density `ρ`.
+"""
+apply_limiter!(
+    q::Fields.Field,
+    ρ::Fields.Field,
+    lim::VerticalMassBorrowingLimiter,
+) = apply_limiter!(q, ρ, axes(q), lim, ClimaComms.device(axes(q)))
+
+function apply_limiter!(
+    q::Fields.Field,
+    ρ::Fields.Field,
+    space::Spaces.FiniteDifferenceSpace,
+    lim::VerticalMassBorrowingLimiter,
+    device::ClimaComms.AbstractCPUDevice,
+)
+    cache = (; bmass = lim.bmass, ic = lim.ic, q_min = lim.q_min)
+    columnwise_massborrow_cpu(q, ρ, cache)
+end
+
+function apply_limiter!(
+    q::Fields.Field,
+    ρ::Fields.Field,
+    space::Spaces.ExtrudedFiniteDifferenceSpace,
+    lim::VerticalMassBorrowingLimiter,
+    device::ClimaComms.AbstractCPUDevice,
+)
+    Fields.bycolumn(axes(q)) do colidx
+        cache = (;
+            bmass = lim.bmass[colidx],
+            ic = lim.ic[colidx],
+            q_min = lim.q_min,
+        )
+        columnwise_massborrow_cpu(q[colidx], ρ[colidx], cache)
+    end
+end
+
+# TODO: can we improve the performance?
+# `bycolumn` on the CPU may be better here since we could multithread it.
+function columnwise_massborrow_cpu(q::Fields.Field, ρ::Fields.Field, cache) # column fields
+    (; bmass, ic, q_min) = cache
+
+    Δz = Fields.Δz_field(q)
+    Δz_vals = Fields.field_values(Δz)
+    ρ_vals = Fields.field_values(ρ)
+    #  loop over tracers
+    nlevels = Spaces.nlevels(axes(q))
+    @. ic = 0
+    @. bmass = 0
+    q_vals = Fields.field_values(q)
+    # top to bottom
+    for f in 1:DataLayouts.ncomponents(q_vals)
+        for v in 1:nlevels
+            CI = CartesianIndex(1, 1, f, v, 1)
+            # new mass in the current layer
+            ρΔz_v =
+                DL.getindex_field(Δz_vals, CI) * DL.getindex_field(ρ_vals, CI)
+            nmass = DL.getindex_field(q_vals, CI) + bmass[] / ρΔz_v
+            if nmass > q_min[f]
+                #  if new mass in the current layer is positive, don't borrow mass any more
+                DL.setindex_field!(q_vals, nmass, CI)
+                bmass[] = 0
+            else
+                #  set mass to q_min in the current layer, and save bmass
+                bmass[] = (nmass - q_min[f]) * ρΔz_v
+                DL.setindex_field!(q_vals, q_min[f], CI)
+                ic[] = ic[] + 1
+            end
+        end
+
+        #  bottom to top
+        for v in nlevels:-1:1
+            CI = CartesianIndex(1, 1, f, v, 1)
+            # if the surface layer still needs to borrow mass
+            if bmass[] < 0
+                ρΔz_v =
+                    DL.getindex_field(Δz_vals, CI) *
+                    DL.getindex_field(ρ_vals, CI)
+                # new mass in the current layer
+                nmass = DL.getindex_field(q_vals, CI) + bmass[] / ρΔz_v
+                if nmass > q_min[f]
+                    # if new mass in the current layer is positive, don't borrow mass any more
+                    DL.setindex_field!(q_vals, nmass, CI)
+                    bmass[] = 0
+                else
+                    # if new mass in the current layer is negative, continue to borrow mass
+                    bmass[] = (nmass - q_min[f]) * ρΔz_v
+                    DL.setindex_field!(q_vals, q_min[f], CI)
+                end
+            end
+        end
+    end
+
+    return nothing
+end

test/Limiters/vertical_mass_borrowing_limiter.jl

@@ -0,0 +1,151 @@
+#=
+julia --project
+using Revise; include(joinpath("test", "Limiters", "vertical_mass_borrowing_limiter.jl"))
+=#
+using ClimaComms
+ClimaComms.@import_required_backends
+using ClimaCore: Fields, Spaces, Limiters
+using ClimaCore.RecursiveApply
+using ClimaCore.Geometry
+using ClimaCore.Grids
+using ClimaCore.CommonGrids
+using Test
+using Random
+
+function perturb_field!(f::Fields.Field; perturb_radius)
+    device = ClimaComms.device(f)
+    ArrayType = ClimaComms.array_type(device)
+    rand_data = ArrayType(rand(size(parent(f))...)) # [0-1]
+    rand_data = rand_data .- sum(rand_data) / length(rand_data) # make centered about 0 ([-0.5:0.5])
+    rand_data = (rand_data ./ maximum(rand_data)) .* perturb_radius # rand_data now in [-perturb_radius:perturb_radius]
+    parent(f) .= parent(f) .+ rand_data # f in [f ± perturb_radius]
+    return nothing
+end
+
+@testset "Vertical mass borrowing limiter - column" begin
+    Random.seed!(1234)
+    z_elem = 10
+    z_min = 0
+    z_max = 1
+    device = ClimaComms.device()
+    grid = ColumnGrid(; z_elem, z_min, z_max, device)
+    cspace = Spaces.FiniteDifferenceSpace(grid, Grids.CellCenter())
+    tol = 0.01
+    perturb_q = 0.3
+    perturb_ρ = 0.2
+
+    # Initialize fields
+    coords = Fields.coordinate_field(cspace)
+    ρ = map(coord -> 1.0, coords)
+    q = map(coord -> 0.1, coords)
+    (; z) = coords
+    perturb_field!(q; perturb_radius = perturb_q)
+    perturb_field!(ρ; perturb_radius = perturb_ρ)
+    ρq_init = ρ .⊠ q
+    sum_ρq_init = sum(ρq_init)
+
+    # Test that the minimum is below 0
+    @test length(parent(q)) == z_elem == 10
+    @test 0.3 ≤ count(x -> x < 0, parent(q)) / 10 ≤ 0.5 # ensure reasonable percentage of points are negative
+
+    @test -2 * perturb_q ≤ minimum(q) ≤ -tol
+    limiter = Limiters.VerticalMassBorrowingLimiter(q, (0.0,))
+    Limiters.apply_limiter!(q, ρ, limiter)
+    @test 0 ≤ minimum(q)
+    ρq = ρ .⊠ q
+    @test isapprox(sum(ρq), sum_ρq_init; atol = 1e-15)
+    @test isapprox(sum(ρq), sum_ρq_init; rtol = 1e-10)
+    # @show sum(ρq)     # 0.07388931313511024
+    # @show sum_ρq_init # 0.07388931313511025
+end
+
+@testset "Vertical mass borrowing limiter - sphere" begin
+    z_elem = 10
+    z_min = 0
+    z_max = 1
+    radius = 10
+    h_elem = 10
+    n_quad_points = 4
+    grid = ExtrudedCubedSphereGrid(;
+        z_elem,
+        z_min,
+        z_max,
+        radius,
+        h_elem,
+        n_quad_points,
+    )
+    cspace = Spaces.ExtrudedFiniteDifferenceSpace(grid, Grids.CellCenter())
+    tol = 0.01
+    perturb_q = 0.2
+    perturb_ρ = 0.1
+
+    # Initialize fields
+    coords = Fields.coordinate_field(cspace)
+    ρ = map(coord -> 1.0, coords)
+    q = map(coord -> 0.1, coords)
+
+    perturb_field!(q; perturb_radius = perturb_q)
+    perturb_field!(ρ; perturb_radius = perturb_ρ)
+    ρq_init = ρ .⊠ q
+    sum_ρq_init = sum(ρq_init)
+
+    # Test that the minimum is below 0
+    @test length(parent(q)) == z_elem * h_elem^2 * 6 * n_quad_points^2 == 96000
+    @test 0.10 ≤ count(x -> x < 0.0001, parent(q)) / 96000 ≤ 1 # ensure 10% of points are negative
+
+    @test -2 * perturb_q ≤ minimum(q) ≤ -tol
+    limiter = Limiters.VerticalMassBorrowingLimiter(q, (0.0,))
+    Limiters.apply_limiter!(q, ρ, limiter)
+    @test 0 ≤ minimum(q)
+    ρq = ρ .⊠ q
+    @test isapprox(sum(ρq), sum_ρq_init; atol = 0.029)
+    @test isapprox(sum(ρq), sum_ρq_init; rtol = 0.00023)
+    # @show sum(ρq)     # 125.5483524237572
+    # @show sum_ρq_init # 125.52052986152977
+end
+
+@testset "Vertical mass borrowing limiter - deep atmosphere" begin
+    z_elem = 10
+    z_min = 0
+    z_max = 1
+    radius = 10
+    h_elem = 10
+    n_quad_points = 4
+    grid = ExtrudedCubedSphereGrid(;
+        z_elem,
+        z_min,
+        z_max,
+        radius,
+        global_geometry = Geometry.DeepSphericalGlobalGeometry(radius),
+        h_elem,
+        n_quad_points,
+    )
+    cspace = Spaces.ExtrudedFiniteDifferenceSpace(grid, Grids.CellCenter())
+    tol = 0.01
+    perturb_q = 0.2
+    perturb_ρ = 0.1
+
+    # Initialize fields
+    coords = Fields.coordinate_field(cspace)
+    ρ = map(coord -> 1.0, coords)
+    q = map(coord -> 0.1, coords)
+
+    perturb_field!(q; perturb_radius = perturb_q)
+    perturb_field!(ρ; perturb_radius = perturb_ρ)
+    ρq_init = ρ .⊠ q
+    sum_ρq_init = sum(ρq_init)
+
+    # Test that the minimum is below 0
+    @test length(parent(q)) == z_elem * h_elem^2 * 6 * n_quad_points^2 == 96000
+    @test 0.10 ≤ count(x -> x < 0.0001, parent(q)) / 96000 ≤ 1 # ensure 10% of points are negative
+
+    @test -2 * perturb_q ≤ minimum(q) ≤ -tol
+    limiter = Limiters.VerticalMassBorrowingLimiter(q, (0.0,))
+    Limiters.apply_limiter!(q, ρ, limiter)
+    @test 0 ≤ minimum(q)
+    ρq = ρ .⊠ q
+    @test isapprox(sum(ρq), sum_ρq_init; atol = 0.269)
+    @test isapprox(sum(ρq), sum_ρq_init; rtol = 0.00199)
+    # @show sum(ρq)     # 138.90494931641584
+    # @show sum_ρq_init # 139.1735731377394
+end