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,99 @@
+#=
+
+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, lim::VerticalMassBorrowingLimiter) =
+    apply_limiter!(q, axes(q), lim, ClimaComms.device(axes(q)))
+
+function apply_limiter!(
+    q::Fields.Field,
+    space::Spaces.FiniteDifferenceSpace,
+    lim::VerticalMassBorrowingLimiter,
+    device::ClimaComms.AbstractCPUDevice,
+)
+    cache = (; bmass = lim.bmass, ic = lim.ic, qmin = lim.qmin)
+    columnwise_massborrow_cpu(q, cache)
+end
+
+function apply_limiter!(
+    q::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], qmin = lim.qmin)
+        columnwise_massborrow_cpu(q[colidx], cache)
+    end
+end
+
+function columnwise_massborrow_cpu(q::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,91 @@
+#=
+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
+using Random
+
+@testset "Vertical mass borrowing limiter - column" begin
+    Random.seed!(1234)
+    FT = Float64
+    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())
+    context = ClimaComms.context(device)
+    ArrayType = ClimaComms.array_type(device)
+    tol = 0.01
+    perturb = 0.2
+
+    # Initialize fields
+    coords = Fields.coordinate_field(cspace)
+    q = map(coord -> 0.1, coords)
+    (; z) = coords
+    rand_data = ArrayType(rand(size(parent(q))...)) # [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 # rand_data now in [-0.2:0.2]
+    parent(q) .= parent(q) .+ rand_data # q in [0.1 ± 0.2]
+    sum_q_init = sum(q)
+    # 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 ≤ minimum(q) ≤ -tol
+    limiter = Limiters.VerticalMassBorrowingLimiter(q, (0.0,))
+    Limiters.apply_limiter!(q, limiter)
+    @test 0 ≤ minimum(q)
+    @test isapprox(sum(q), sum_q_init; atol = 0.00000000001)
+    @test isapprox(sum(q), sum_q_init; rtol = 0.01)
+end
+
+@testset "Vertical mass borrowing limiter" begin
+    FT = Float64
+    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())
+    context = ClimaComms.context(device)
+    ArrayType = ClimaComms.array_type(device)
+    tol = 0.01
+    perturb = 0.2
+
+    # Initialize fields
+    coords = Fields.coordinate_field(cspace)
+    q = map(coord -> 0.1, coords)
+
+    rand_data = ArrayType(rand(size(parent(q))...)) # [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 # rand_data now in [-0.2:0.2]
+    parent(q) .= parent(q) .+ rand_data # q in [0.1 ± 0.2]
+    sum_q_init = sum(q)
+
+    # 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 -0.11 ≤ minimum(q) ≤ -0.08
+    limiter = Limiters.VerticalMassBorrowingLimiter(q, (0.0,))
+    Limiters.apply_limiter!(q, limiter)
+    @test 0 ≤ minimum(q)
+    @test isapprox(sum(q), sum_q_init; atol = 0.013)
+    @test isapprox(sum(q), sum_q_init; rtol = 0.01)
+end

test/Limiters/vertical_mass_borrowing_limiter_advection.jl

@@ -0,0 +1,145 @@
+#=
+julia --project=.buildkite
+using Revise; include(joinpath("test", "Limiters", "vertical_mass_borrowing_limiter_advection.jl"))
+=#
+using Test
+using LinearAlgebra
+import ClimaComms
+ClimaComms.@import_required_backends
+using OrdinaryDiffEqSSPRK: ODEProblem, solve, SSPRK33
+using ClimaTimeSteppers
+
+import ClimaCore:
+    Fields,
+    Domains,
+    Limiters,
+    Topologies,
+    Meshes,
+    DataLayouts,
+    Operators,
+    Geometry,
+    Spaces
+
+
+# Advection Equation, with constant advective velocity (so advection form = flux form)
+# ∂_t y + w ∂_z y  = 0
+# the solution translates to the right at speed w,
+# so at time t, the solution is y(z - w * t)
+
+# visualization artifacts
+ENV["GKSwstype"] = "nul"
+using ClimaCorePlots, Plots
+Plots.GRBackend()
+dir = "vert_mass_borrow_advection"
+path = joinpath(@__DIR__, "output", dir)
+mkpath(path)
+
+function lim!(y, parameters, t, y_ref)
+    (; w, Δt, limiter) = parameters
+    Limiters.apply_limiter!(y.q, limiter)
+    return nothing
+end
+
+function tendency!(yₜ, y, parameters, t)
+    (; w, Δt, limiter) = parameters
+    FT = Spaces.undertype(axes(y.q))
+    bcvel = pulse(-π, t, z₀, zₕ, z₁)
+    divf2c = Operators.DivergenceF2C(
+        bottom = Operators.SetValue(Geometry.WVector(FT(bcvel))),
+        top = Operators.SetValue(Geometry.WVector(FT(0))),
+    )
+    upwind1 = Operators.UpwindBiasedProductC2F(
+        bottom = Operators.Extrapolate(),
+        top = Operators.Extrapolate(),
+    )
+    upwind3 = Operators.Upwind3rdOrderBiasedProductC2F(
+        bottom = Operators.ThirdOrderOneSided(),
+        top = Operators.ThirdOrderOneSided(),
+    )
+    If = Operators.InterpolateC2F()
+    @. yₜ.q =
+    # -divf2c(w * If(y.q))
+        -divf2c(upwind1(w, y.q) * If(y.q))
+    # -divf2c(upwind3(w, y.q) * If(y.q))
+    return nothing
+end
+
+# Define a pulse wave or square wave
+
+FT = Float64
+t₀ = FT(0.0)
+Δt = 0.0001 * 25
+t₁ = 200Δt
+z₀ = FT(0.0)
+zₕ = FT(1.0)
+z₁ = FT(1.0)
+speed = FT(-1.0)
+pulse(z, t, z₀, zₕ, z₁) = abs(z - speed * t) ≤ zₕ ? z₁ : z₀
+
+n = 2 .^ 6
+
+domain = Domains.IntervalDomain(
+    Geometry.ZPoint{FT}(-π),
+    Geometry.ZPoint{FT}(π);
+    boundary_names = (:bottom, :top),
+)
+
+# stretch_fns = (Meshes.Uniform(), Meshes.ExponentialStretching(FT(7.0)))
+stretch_fns = (Meshes.Uniform(),)
+plot_string = ["uniform", "stretched"]
+
+@testset "VerticalMassBorrowingLimiter on advection" begin
+    for (i, stretch_fn) in enumerate(stretch_fns)
+        mesh = Meshes.IntervalMesh(domain, stretch_fn; nelems = n)
+        cent_space = Spaces.CenterFiniteDifferenceSpace(mesh)
+        face_space = Spaces.FaceFiniteDifferenceSpace(cent_space)
+        z = Fields.coordinate_field(cent_space).z
+        O = ones(FT, face_space)
+
+        # Initial condition
+        q_init = pulse.(z, 0.0, z₀, zₕ, z₁)
+        q = q_init
+        y = Fields.FieldVector(; q)
+        limiter = Limiters.VerticalMassBorrowingLimiter(q, (0.0,))
+
+        # Unitary, constant advective velocity
+        w = Geometry.WVector.(speed .* O)
+
+        # Solve the ODE
+        parameters = (; w, Δt, limiter)
+        prob = ODEProblem(
+            ClimaODEFunction(; lim!, T_lim! = tendency!),
+            y,
+            (t₀, t₁),
+            parameters,
+        )
+        sol = solve(
+            prob,
+            ExplicitAlgorithm(SSP33ShuOsher()),
+            dt = Δt,
+            saveat = Δt,
+        )
+
+        q_init = sol.u[1].q
+        q_final = sol.u[end].q
+        q_analytic = pulse.(z, t₁, z₀, zₕ, z₁)
+        err = norm(q_final .- q_analytic)
+        rel_mass_err = norm((sum(q_final) - sum(q_init)) / sum(q_init))
+
+
+        p = plot()
+        Plots.plot!(q_init, label = "init")
+        Plots.plot!(q_final, label = "computed")
+        Plots.plot!(q_analytic, label = "analytic")
+        Plots.plot!(; legend = :topright)
+        Plots.plot!(; xlabel = "q", title = "VerticalMassBorrowingLimiter")
+        f = joinpath(
+            path,
+            "VerticalMassBorrowingLimiter_comparison_$(plot_string[i]).png",
+        )
+        Plots.png(p, f)
+        @test err ≤ 0.25
+        @test rel_mass_err ≤ 10eps()
+        @test minimum(q_final) ≥ 0
+    end
+end