Maximum-principle-satisfying WENO scheme #2557
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
| @@ -0,0 +1,94 @@ | ||||||
| using Oceananigans.Grids: AbstractGrid | ||||||
|
|
||||||
| const ω̂₁ = 5/18 | ||||||
| const ω̂ₙ = 5/18 | ||||||
| const ε₂ = 1e-20 | ||||||
|
|
||||||
| # Here in the future we can easily add UpwindBiasedFifthOrder | ||||||
| const BoundPreservingScheme = PositiveWENO | ||||||
|
|
||||||
| @inline function div_Uc(i, j, k, grid, advection::BoundPreservingScheme, U, c) | ||||||
|
|
||||||
| div_x = bounded_tracer_flux_divergence_x(i, j, k, grid, advection, U.u, c) | ||||||
| div_y = bounded_tracer_flux_divergence_y(i, j, k, grid, advection, U.v, c) | ||||||
| div_z = bounded_tracer_flux_divergence_z(i, j, k, grid, advection, U.w, c) | ||||||
|
|
||||||
| return 1/Vᶜᶜᶜ(i, j, k, grid) * (div_x + div_y + div_z) | ||||||
| end | ||||||
|
|
||||||
| # Support for Flat directions | ||||||
| @inline bounded_tracer_flux_divergence_x(i, j, k, ::AbstractGrid{FT, Flat, TY, TZ}, args...) where {FT, TY, TZ} = zero(FT) | ||||||
| @inline bounded_tracer_flux_divergence_y(i, j, k, ::AbstractGrid{FT, TX, Flat, TZ}, args...) where {FT, TX, TZ} = zero(FT) | ||||||
| @inline bounded_tracer_flux_divergence_z(i, j, k, ::AbstractGrid{FT, TX, TY, Flat}, args...) where {FT, TX, TY} = zero(FT) | ||||||
|
|
||||||
| @inline function bounded_tracer_flux_divergence_x(i, j, k, grid, advection::BoundPreservingScheme, u, c) | ||||||
|
|
||||||
| lower_limit = advection.bounds[1] | ||||||
| upper_limit = advection.bounds[2] | ||||||
|
|
||||||
| cᵢⱼ = c[i, j, k] | ||||||
|
|
||||||
| c₊ᴸ = _left_biased_interpolate_xᶠᵃᵃ(i+1, j, k, grid, advection, c) | ||||||
| c₊ᴿ = _right_biased_interpolate_xᶠᵃᵃ(i+1, j, k, grid, advection, c) | ||||||
| c₋ᴸ = _left_biased_interpolate_xᶠᵃᵃ(i, j, k, grid, advection, c) | ||||||
| c₋ᴿ = _right_biased_interpolate_xᶠᵃᵃ(i, j, k, grid, advection, c) | ||||||
|
|
||||||
| p̃ = (cᵢⱼ - ω̂₁ * c₋ᴿ - ω̂ₙ * c₊ᴸ) / (1 - 2ω̂₁) | ||||||
| M = max(p̃, c₊ᴸ, c₋ᴿ) | ||||||
| m = min(p̃, c₊ᴸ, c₋ᴿ) | ||||||
| θ = min(abs((upper_limit - cᵢⱼ)/(M - cᵢⱼ + ε₂)), abs((lower_limit - cᵢⱼ)/(m - cᵢⱼ + ε₂)), one(grid)) | ||||||
|
|
||||||
| c₊ᴸ = θ * (c₊ᴸ - cᵢⱼ) + cᵢⱼ | ||||||
| c₋ᴿ = θ * (c₋ᴿ - cᵢⱼ) + cᵢⱼ | ||||||
|
|
||||||
| return Axᶠᶜᶜ(i+1, j, k, grid) * upwind_biased_product(u[i+1, j, k], c₊ᴸ, c₊ᴿ) - | ||||||
| Axᶠᶜᶜ(i, j, k, grid) * upwind_biased_product(u[i, j, k], c₋ᴸ, c₋ᴿ) | ||||||
| end | ||||||
|
|
||||||
| @inline function bounded_tracer_flux_divergence_y(i, j, k, grid, advection::BoundPreservingScheme, v, c) | ||||||
|
|
||||||
| lower_limit = advection.bounds[1] | ||||||
| upper_limit = advection.bounds[2] | ||||||
|
|
||||||
| cᵢⱼ = c[i, j, k] | ||||||
|
|
||||||
| c₊ᴸ = _left_biased_interpolate_yᵃᶠᵃ(i, j+1, k, grid, advection, c) | ||||||
| c₊ᴿ = _right_biased_interpolate_yᵃᶠᵃ(i, j+1, k, grid, advection, c) | ||||||
| c₋ᴸ = _left_biased_interpolate_yᵃᶠᵃ(i, j, k, grid, advection, c) | ||||||
| c₋ᴿ = _right_biased_interpolate_yᵃᶠᵃ(i, j, k, grid, advection, c) | ||||||
|
|
||||||
| p̃ = (cᵢⱼ - ω̂₁ * c₋ᴿ - ω̂ₙ * c₊ᴸ) / (1 - 2ω̂₁) | ||||||
| M = max(p̃, c₊ᴸ, c₋ᴿ) | ||||||
| m = min(p̃, c₊ᴸ, c₋ᴿ) | ||||||
| θ = min(abs((upper_limit - cᵢⱼ)/(M - cᵢⱼ + ε₂)), abs((lower_limit - cᵢⱼ)/(m - cᵢⱼ + ε₂)), one(grid)) | ||||||
|
|
||||||
| c₊ᴸ = θ * (c₊ᴸ - cᵢⱼ) + cᵢⱼ | ||||||
| c₋ᴿ = θ * (c₋ᴿ - cᵢⱼ) + cᵢⱼ | ||||||
|
|
||||||
| return Ayᶜᶠᶜ(i, j+1, k, grid) * upwind_biased_product(v[i, j+1, k], c₊ᴸ, c₊ᴿ) - | ||||||
| Ayᶜᶠᶜ(i, j, k, grid) * upwind_biased_product(v[i, j, k], c₋ᴸ, c₋ᴿ) | ||||||
| end | ||||||
|
|
||||||
| @inline function bounded_tracer_flux_divergence_z(i, j, k, grid, advection::BoundPreservingScheme, w, c) | ||||||
|
|
||||||
| lower_limit = advection.bounds[1] | ||||||
| upper_limit = advection.bounds[2] | ||||||
|
|
||||||
| cᵢⱼ = c[i, j, k] | ||||||
|
|
||||||
| c₊ᴸ = _left_biased_interpolate_zᵃᵃᶠ(i, j, k+1, grid, advection, c) | ||||||
| c₊ᴿ = _right_biased_interpolate_zᵃᵃᶠ(i, j, k+1, grid, advection, c) | ||||||
| c₋ᴸ = _left_biased_interpolate_zᵃᵃᶠ(i, j, k, grid, advection, c) | ||||||
| c₋ᴿ = _right_biased_interpolate_zᵃᵃᶠ(i, j, k, grid, advection, c) | ||||||
|
|
||||||
| p̃ = (cᵢⱼ - ω̂₁ * c₋ᴿ - ω̂ₙ * c₊ᴸ) / (1 - 2ω̂₁) | ||||||
| M = max(p̃, c₊ᴸ, c₋ᴿ) | ||||||
| m = min(p̃, c₊ᴸ, c₋ᴿ) | ||||||
| θ = min(abs((upper_limit - cᵢⱼ)/(M - cᵢⱼ + ε₂)), abs((lower_limit - cᵢⱼ)/(m - cᵢⱼ + ε₂)), 1.0) | ||||||
|
There was a problem hiding this comment.
Suggested change
|
||||||
|
|
||||||
| c₊ᴸ = θ * (c₊ᴸ - cᵢⱼ) + cᵢⱼ | ||||||
| c₋ᴿ = θ * (c₋ᴿ - cᵢⱼ) + cᵢⱼ | ||||||
|
|
||||||
| return Azᶜᶜᶠ(i, j, k+1, grid) * upwind_biased_product(w[i, j, k+1], c₊ᴸ, c₊ᴿ) - | ||||||
| Azᶜᶜᶠ(i, j, k, grid) * upwind_biased_product(w[i, j, k], c₋ᴸ, c₋ᴿ) | ||||||
| end | ||||||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -9,6 +9,10 @@ using KernelAbstractions.Extras.LoopInfo: @unroll | |
| using Adapt | ||
| import Base: show | ||
|
|
||
| const C3₀ = 3/10 | ||
| const C3₁ = 3/5 | ||
| const C3₂ = 1/10 | ||
|
|
||
| const two_32 = Int32(2) | ||
|
|
||
| const ƞ = Int32(2) # WENO exponent | ||
|
|
@@ -26,7 +30,7 @@ Weighted Essentially Non-Oscillatory (WENO) fifth-order advection scheme. | |
|
|
||
| $(TYPEDFIELDS) | ||
| """ | ||
| struct WENO5{FT, XT, YT, ZT, XS, YS, ZS, VI, WF, PP} <: AbstractUpwindBiasedAdvectionScheme{2} | ||
|
There was a problem hiding this comment. Is this the design we will want when we do #2454 ? If we can avoid digging ourselves into a deeper hole we should, but maybe it's not possible without a lot of effort... There was a problem hiding this comment. I am not still completely sure how #2454 will go along, this change only modifies the There was a problem hiding this comment. In #2454 we will disentangle vector invariant from WENO5, so There was a problem hiding this comment. I think the limiter will be independent on the order of the original scheme. You always want to blend your scheme into a first order upwind (when crossing bounds) |
||
|
|
||
| "coefficient for ENO reconstruction on x-faces" | ||
| coeff_xᶠᵃᵃ::XT | ||
|
|
@@ -54,22 +58,20 @@ struct WENO5{FT, XT, YT, ZT, XS, YS, ZS, VI, WF} <: AbstractUpwindBiasedAdvectio | |
| "coefficient for WENO smoothness indicators on z-centers" | ||
| smooth_zᵃᵃᶜ::ZS | ||
|
|
||
| "bounds for maximum-principle-satisfying WENO scheme" | ||
| bounds :: PP | ||
|
|
||
| function WENO5{FT, VI, WF}(coeff_xᶠᵃᵃ::XT, coeff_xᶜᵃᵃ::XT, | ||
| coeff_yᵃᶠᵃ::YT, coeff_yᵃᶜᵃ::YT, | ||
| coeff_zᵃᵃᶠ::ZT, coeff_zᵃᵃᶜ::ZT, | ||
| smooth_xᶠᵃᵃ::XS, smooth_xᶜᵃᵃ::XS, | ||
| smooth_yᵃᶠᵃ::YS, smooth_yᵃᶜᵃ::YS, | ||
| smooth_zᵃᵃᶠ::ZS, smooth_zᵃᵃᶜ::ZS, | ||
| bounds::PP) where {FT, XT, YT, ZT, XS, YS, ZS, VI, WF, PP} | ||
|
|
||
| return new{FT, XT, YT, ZT, XS, YS, ZS, VI, WF, PP}(coeff_xᶠᵃᵃ, coeff_xᶜᵃᵃ, coeff_yᵃᶠᵃ, coeff_yᵃᶜᵃ, coeff_zᵃᵃᶠ, coeff_zᵃᵃᶜ, | ||
| smooth_xᶠᵃᵃ, smooth_xᶜᵃᵃ, smooth_yᵃᶠᵃ, smooth_yᵃᶜᵃ, smooth_zᵃᵃᶠ, smooth_zᵃᵃᶜ, | ||
| bounds) | ||
| end | ||
| end | ||
|
|
||
|
|
@@ -161,26 +163,20 @@ WENO5(grid, FT::DataType=Float64; kwargs...) = WENO5(FT; grid = grid, kwargs...) | |
|
|
||
| function WENO5(FT::DataType = Float64; | ||
| grid = nothing, | ||
| stretched_smoothness = false, | ||
| zweno = true, | ||
| vector_invariant = nothing, | ||
| bounds = nothing) | ||
|
|
||
| if !(grid isa Nothing) | ||
| FT = eltype(grid) | ||
| end | ||
|
|
||
| weno_coefficients = compute_stretched_weno_coefficients(grid, stretched_smoothness, FT) | ||
|
|
||
| VI = typeof(vector_invariant) | ||
|
|
||
| return WENO5{FT, VI, zweno}(weno_coefficients..., bounds) | ||
| end | ||
|
|
||
| function compute_stretched_weno_coefficients(grid, stretched_smoothness, FT) | ||
|
|
@@ -216,14 +212,15 @@ return_metrics(::LatitudeLongitudeGrid) = (:λᶠᵃᵃ, :λᶜᵃᵃ, :φᵃᶠ | |
| return_metrics(::RectilinearGrid) = (:xᶠᵃᵃ, :xᶜᵃᵃ, :yᵃᶠᵃ, :yᵃᶜᵃ, :zᵃᵃᶠ, :zᵃᵃᶜ) | ||
|
|
||
| # Flavours of WENO | ||
| const ZWENO = WENO5{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, true} | ||
| const PositiveWENO = WENO5{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Tuple} | ||
|
|
||
| const WENOVectorInvariantVel{FT, XT, YT, ZT, XS, YS, ZS, VI, WF, PP} = | ||
| WENO5{FT, XT, YT, ZT, XS, YS, ZS, VI, WF, PP} where {FT, XT, YT, ZT, XS, YS, ZS, VI<:VelocityStencil, WF, PP} | ||
| const WENOVectorInvariantVort{FT, XT, YT, ZT, XS, YS, ZS, VI, WF, PP} = | ||
| WENO5{FT, XT, YT, ZT, XS, YS, ZS, VI, WF, PP} where {FT, XT, YT, ZT, XS, YS, ZS, VI<:VorticityStencil, WF, PP} | ||
|
|
||
| const WENOVectorInvariant = WENO5{FT, XT, YT, ZT, XS, YS, ZS, VI, WF, PP} where {FT, XT, YT, ZT, XS, YS, ZS, VI<:SmoothnessStencil, WF, PP} | ||
|
|
||
| function Base.show(io::IO, a::WENO5{FT, RX, RY, RZ}) where {FT, RX, RY, RZ} | ||
| print(io, "WENO5 advection scheme with: \n", | ||
|
|
@@ -232,14 +229,14 @@ function Base.show(io::IO, a::WENO5{FT, RX, RY, RZ}) where {FT, RX, RY, RZ} | |
| " └── Z $(RZ == Nothing ? "regular" : "stretched")" ) | ||
| end | ||
|
|
||
| Adapt.adapt_structure(to, scheme::WENO5{FT, XT, YT, ZT, XS, YS, ZS, VI, WF, PP}) where {FT, XT, YT, ZT, XS, YS, ZS, VI, WF, PP} = | ||
| WENO5{FT, VI, WF}(Adapt.adapt(to, scheme.coeff_xᶠᵃᵃ), Adapt.adapt(to, scheme.coeff_xᶜᵃᵃ), | ||
| Adapt.adapt(to, scheme.coeff_yᵃᶠᵃ), Adapt.adapt(to, scheme.coeff_yᵃᶜᵃ), | ||
| Adapt.adapt(to, scheme.coeff_zᵃᵃᶠ), Adapt.adapt(to, scheme.coeff_zᵃᵃᶜ), | ||
| Adapt.adapt(to, scheme.smooth_xᶠᵃᵃ), Adapt.adapt(to, scheme.smooth_xᶜᵃᵃ), | ||
| Adapt.adapt(to, scheme.smooth_yᵃᶠᵃ), Adapt.adapt(to, scheme.smooth_yᵃᶜᵃ), | ||
| Adapt.adapt(to, scheme.smooth_zᵃᵃᶠ), Adapt.adapt(to, scheme.smooth_zᵃᵃᶜ), | ||
| Adapt.adapt(to, scheme.bounds)) | ||
|
|
||
| @inline boundary_buffer(::WENO5) = 2 | ||
|
|
||
|
|
@@ -391,13 +388,13 @@ for (side, coeffs) in zip([:left, :right], ([:C3₀, :C3₁, :C3₂], [:C3₂, : | |
|
|
||
| if scheme isa ZWENO | ||
| τ₅ = abs(β₂ - β₀) | ||
| α₀ = FT($(coeffs[1])) * (1 + (τ₅ / (β₀ + FT(ε)))^ƞ) | ||
| α₁ = FT($(coeffs[2])) * (1 + (τ₅ / (β₁ + FT(ε)))^ƞ) | ||
| α₂ = FT($(coeffs[3])) * (1 + (τ₅ / (β₂ + FT(ε)))^ƞ) | ||
| else | ||
| α₀ = FT($(coeffs[1])) / (β₀ + FT(ε))^ƞ | ||
| α₁ = FT($(coeffs[2])) / (β₁ + FT(ε))^ƞ | ||
| α₂ = FT($(coeffs[3])) / (β₂ + FT(ε))^ƞ | ||
| end | ||
|
|
||
| Σα = α₀ + α₁ + α₂ | ||
|
|
@@ -428,13 +425,13 @@ for (side, coeffs) in zip([:left, :right], ([:C3₀, :C3₁, :C3₂], [:C3₂, : | |
|
|
||
| if scheme isa ZWENO | ||
| τ₅ = abs(β₂ - β₀) | ||
| α₀ = FT($(coeffs[1])) * (1 + (τ₅ / (β₀ + FT(ε)))^ƞ) | ||
| α₁ = FT($(coeffs[2])) * (1 + (τ₅ / (β₁ + FT(ε)))^ƞ) | ||
| α₂ = FT($(coeffs[3])) * (1 + (τ₅ / (β₂ + FT(ε)))^ƞ) | ||
| else | ||
| α₀ = FT($(coeffs[1])) / (β₀ + FT(ε))^ƞ | ||
| α₁ = FT($(coeffs[2])) / (β₁ + FT(ε))^ƞ | ||
| α₂ = FT($(coeffs[3])) / (β₂ + FT(ε))^ƞ | ||
| end | ||
|
|
||
| Σα = α₀ + α₁ + α₂ | ||
|
|
@@ -456,13 +453,13 @@ for (side, coeffs) in zip([:left, :right], ([:C3₀, :C3₁, :C3₂], [:C3₂, : | |
|
|
||
| if scheme isa ZWENO | ||
| τ₅ = abs(β₂ - β₀) | ||
| α₀ = FT($(coeffs[1])) * (1 + (τ₅ / (β₀ + FT(ε)))^ƞ) | ||
| α₁ = FT($(coeffs[2])) * (1 + (τ₅ / (β₁ + FT(ε)))^ƞ) | ||
| α₂ = FT($(coeffs[3])) * (1 + (τ₅ / (β₂ + FT(ε)))^ƞ) | ||
| else | ||
| α₀ = FT($(coeffs[1])) / (β₀ + FT(ε))^ƞ | ||
| α₁ = FT($(coeffs[2])) / (β₁ + FT(ε))^ƞ | ||
| α₂ = FT($(coeffs[3])) / (β₂ + FT(ε))^ƞ | ||
| end | ||
|
|
||
| Σα = α₀ + α₁ + α₂ | ||
|
|
||
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This looks great! It seems like these three funcionts are almost identical except for u -> v -> w. Is that true?
If yes, I wonder if there is away to merge these to one, rather than having three functions that do the same thing?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Indices, areas and interpolations differ. I can try merging it with some metaprogramming
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Not necessary. This is maybe cleaner but thought I would make the observation.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I prefer having separate functions over metaprogramming in this case.
@francispoulin the difference between x, y, z is that they depend on indices i, j, k respectively. So we also have advective_tracer_flux_x, etc.