Implement and test vertical mass borrowing limiters

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,92 @@
+#=
+
+From E3SM:
+    https://github.com/E3SM-Project/E3SM/blob/2c377c5ec9a5585170524b366ad85074ab1b1a5c/components/eam/src/physics/cam/massborrow.F90
+Described in Appendix B of
+    https://gmd.copernicus.org/articles/11/1971/2018/gmd-11-1971-2018.pdf
+
+=#
+
+import .DataLayouts as DL
+
+struct VerticalMassBorrowingLimiter{F, T}
+    bmass::F
+    ic::F
+    qmin::T
+end
+function VerticalMassBorrowingLimiter(f::Fields.Field, qmin)
+    bmass = similar(Spaces.level(f, 1))
+    ic = similar(Spaces.level(f, 1))
+    return VerticalMassBorrowingLimiter(bmass, ic, qmin)
+end
+
+apply_limiter!(
+    ρq::Fields.Field,
+    ρ::Fields.Field,
+    lim::VerticalMassBorrowingLimiter,
+) = apply_limiter!(ρq, ρ, lim, ClimaComms.device(axes(ρ)))
+
+function apply_limiter!(
+    ρq::Fields.Field,
+    ρ::Fields.Field,
+    lim::VerticalMassBorrowingLimiter,
+    device::ClimaComms.AbstractCPUDevice,
+)
+    Fields.bycolumn(axes(ρ)) do colidx
+        cache =
+            (; bmass = lim.bmass[colidx], ic = lim.ic[colidx], qmin = lim.qmin)
+        columnwise_massborrow_cpu(ρq[colidx], ρ[colidx], cache)
+    end
+end
+
+function columnwise_massborrow_cpu(ρq::Fields.Field, ρ::Fields.Field, cache) # column fields
+    (; bmass, ic, qmin) = cache
+
+    pdel = Fields.field_values(Fields.Δz_field(ρq))
+    #  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
+            pdel_v = DL.getindex_field(pdel, CI)
+            nmass = DL.getindex_field(ρq_vals, CI) + bmass[] / pdel_v
+            if nmass > qmin[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 qmin in the current layer, and save bmass
+                bmass[] = (nmass - qmin[f]) * pdel_v
+                DL.setindex_field!(ρq_vals, qmin[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
+                pdel_v = DL.getindex_field(pdel, CI)
+                # new mass in the current layer
+                nmass = DL.getindex_field(ρq_vals, CI) + bmass[] / pdel_v
+                if nmass > qmin[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 - qmin[f]) * pdel_v
+                    DL.setindex_field!(ρq_vals, qmin[f], CI)
+                end
+            end
+        end
+    end
+
+    return nothing
+end

test/Limiters/vertical_mass_borrowing_limiter.jl

@@ -0,0 +1,43 @@
+#=
+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.Grids
+using ClimaCore.CommonGrids
+using Test
+
+@testset "Vertical mass borrowing limiter" begin
+
+    FT = Float64
+    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())
+
+    # Initialize fields
+    coords = Fields.coordinate_field(cspace)
+    ρ = map(coord -> 1.0, coords)
+    x_scale = FT(1.2)
+    y_scale = FT(1.5)
+    q = similar(ρ)
+    ρq = ρ .⊠ q
+    q_ref = similar(ρ)
+    ρq_ref = ρ .⊠ q_ref
+
+    limiter = Limiters.VerticalMassBorrowingLimiter(ρq, (0.0,))
+
+    # Limiters.compute_bounds!(limiter, ρq_ref, ρ)
+    Limiters.apply_limiter!(ρq, ρ, limiter)
+    q = RecursiveApply.rdiv.(ρq, ρ)
+
+    # @test
+end