diff --git a/Manifest.toml b/Manifest.toml index 961740127e..93e23ccfeb 100644 --- a/Manifest.toml +++ b/Manifest.toml @@ -2,7 +2,7 @@ julia_version = "1.10.3" manifest_format = "2.0" -project_hash = "04d395caf937b0921325a77873167e8baa293a99" +project_hash = "ebd246712231f728236bc79e507cd0d30883b432" [[deps.AbstractFFTs]] deps = ["LinearAlgebra"] @@ -403,15 +403,9 @@ version = "1.5.0" [[deps.JSON3]] deps = ["Dates", "Mmap", "Parsers", "PrecompileTools", "StructTypes", "UUIDs"] -git-tree-sha1 = "eb3edce0ed4fa32f75a0a11217433c31d56bd48b" +git-tree-sha1 = "95220473901735a0f4df9d1ca5b171b568b2daa3" uuid = "0f8b85d8-7281-11e9-16c2-39a750bddbf1" -version = "1.14.0" - - [deps.JSON3.extensions] - JSON3ArrowExt = ["ArrowTypes"] - - [deps.JSON3.weakdeps] - ArrowTypes = "31f734f8-188a-4ce0-8406-c8a06bd891cd" +version = "1.13.2" [[deps.JuliaNVTXCallbacks_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] diff --git a/Project.toml b/Project.toml index 5d2eb1405e..f43d41cb53 100644 --- a/Project.toml +++ b/Project.toml @@ -17,6 +17,7 @@ IncompleteLU = "40713840-3770-5561-ab4c-a76e7d0d7895" InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153" JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" +JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Logging = "56ddb016-857b-54e1-b83d-db4d58db5568" @@ -32,6 +33,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Rotations = "6038ab10-8711-5258-84ad-4b1120ba62dc" SeawaterPolynomials = "d496a93d-167e-4197-9f49-d3af4ff8fe40" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a" diff --git a/src/Advection/Advection.jl b/src/Advection/Advection.jl index b8d7d96469..e85f51009c 100644 --- a/src/Advection/Advection.jl +++ b/src/Advection/Advection.jl @@ -40,6 +40,8 @@ import Base: show, summary import Oceananigans.Grids: required_halo_size import Oceananigans.Architectures: on_architecture +using KernelAbstractions.Extras.LoopInfo: @unroll + abstract type AbstractAdvectionScheme{B, FT} end abstract type AbstractCenteredAdvectionScheme{B, FT} <: AbstractAdvectionScheme{B, FT} end abstract type AbstractUpwindBiasedAdvectionScheme{B, FT} <: AbstractAdvectionScheme{B, FT} end @@ -65,6 +67,9 @@ include("centered_reconstruction.jl") include("upwind_biased_reconstruction.jl") include("weno_reconstruction.jl") include("weno_interpolants.jl") +include("weno_default_stencils.jl") +include("weno_function_stencils.jl") +include("weno_velocity_stencils.jl") include("stretched_weno_smoothness.jl") include("multi_dimensional_reconstruction.jl") include("vector_invariant_upwinding.jl") diff --git a/src/Advection/tracer_advection_operators.jl b/src/Advection/tracer_advection_operators.jl index 86cc43229b..deb6501857 100644 --- a/src/Advection/tracer_advection_operators.jl +++ b/src/Advection/tracer_advection_operators.jl @@ -41,6 +41,7 @@ end ##### Tracer advection operator ##### + """ div_uc(i, j, k, grid, advection, U, c) diff --git a/src/Advection/upwind_biased_advective_fluxes.jl b/src/Advection/upwind_biased_advective_fluxes.jl index 701c266519..d6ff6687db 100644 --- a/src/Advection/upwind_biased_advective_fluxes.jl +++ b/src/Advection/upwind_biased_advective_fluxes.jl @@ -100,7 +100,7 @@ end ##### Tracer advection operators ##### -@inline function advective_tracer_flux_x(i, j, k, grid, scheme::UpwindScheme, U, c) +@inline function advective_tracer_flux_x(i, j, k, grid, scheme::UpwindScheme, U, c) @inbounds ũ = U[i, j, k] cᴸ = _left_biased_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, c) diff --git a/src/Advection/vector_invariant_advection.jl b/src/Advection/vector_invariant_advection.jl index 537dcf6eb3..27403d1324 100644 --- a/src/Advection/vector_invariant_advection.jl +++ b/src/Advection/vector_invariant_advection.jl @@ -108,9 +108,9 @@ Vector Invariant, Dimension-by-dimension reconstruction └── smoothness δv²: FunctionStencil f = v_smoothness ``` """ -function VectorInvariant(; vorticity_scheme = EnstrophyConserving(), +function VectorInvariant(; vorticity_scheme = EnstrophyConserving(), vorticity_stencil = VelocityStencil(), - vertical_scheme = EnergyConserving(), + vertical_scheme = EnergyConserving(), divergence_scheme = vertical_scheme, kinetic_energy_gradient_scheme = divergence_scheme, upwinding = OnlySelfUpwinding(; cross_scheme = divergence_scheme), @@ -152,15 +152,33 @@ const VectorInvariantVelocityVerticalUpwinding = VectorInvariant{<:Any, <:Any, Base.summary(a::VectorInvariant) = string("Vector Invariant, Dimension-by-dimension reconstruction") Base.summary(a::MultiDimensionalVectorInvariant) = string("Vector Invariant, Multidimensional reconstruction") -Base.show(io::IO, a::VectorInvariant{N, FT}) where {N, FT} = - print(io, summary(a), " \n", - " Vorticity flux scheme: ", "\n", - " $(a.vorticity_scheme isa WENO ? "├" : "└")── $(summary(a.vorticity_scheme))", - " $(a.vorticity_scheme isa WENO ? "\n └── smoothness ζ: $(a.vorticity_stencil)\n" : "\n")", - " Vertical advection / Divergence flux scheme: ", "\n", - " $(a.vertical_scheme isa WENO ? "├" : "└")── $(summary(a.vertical_scheme))", - "$(a.vertical_scheme isa AbstractUpwindBiasedAdvectionScheme ? - "\n └── upwinding treatment: $(a.upwinding)" : "")") +function Base.show(io::IO, a::VectorInvariant{N, FT}) where {N, FT} + + δscheme = a.divergence_scheme + vscheme = a.vertical_scheme + ζscheme = a.vorticity_scheme + kscheme = a.kinetic_energy_gradient_scheme + + msg1 = " Vorticity flux scheme: \n" + msg2 = "└── $(summary(ζscheme)) \n" + msg3 = " Kinetic energy gradient flux scheme: \n" + msg4 = "└── $(summary(kscheme)) \n" + msg5 = " Vertical advection scheme: \n" + msg6 = "└── $(summary(vscheme)) \n" + msg7 = (a.vertical_scheme isa EnergyConserving) ? "" : " Divergence flux scheme: \n" + msg8 = isempty(msg7) ? "" : "└── $(summary(a.divergence_scheme)) \n" + + upwinding = (δscheme isa WENO) || (kscheme isa WENO) || (ζscheme isa WENO) + + msg9 = upwinding ? " WENO smoothness stencils: \n" : "" + msg10 = !(ζscheme isa WENO) ? "" : "└── smoothness ζ: $(a.vorticity_stencil)\n" + msg11 = !(δscheme isa WENO) ? "" : "└── smoothness δx_U : $(a.upwinding.δU_stencil)\n" + msg12 = !(δscheme isa WENO) ? "" : "└── smoothness δy_V : $(a.upwinding.δV_stencil)\n" + msg11 = !(kscheme isa WENO) ? "" : "└── smoothness δx_u² : $(a.upwinding.δu²_stencil)\n" + msg12 = !(kscheme isa WENO) ? "" : "└── smoothness δy_v² : $(a.upwinding.δv²_stencil)\n" + + return print(io, summary(a), "\n", msg1, msg2, msg3, msg4, msg5, msg6, msg7, msg8, msg9, msg10, msg11, msg12) +end ##### ##### Convenience for WENO Vector Invariant @@ -309,6 +327,10 @@ end return 1/Vᶜᶠᶜ(i, j, k, grid) * (Φᵟ + 𝒜ᶻ) end +# Fallback for centered advection schemes +@inline upwinded_divergence_flux_Uᶠᶜᶜ(i, j, k, grid, scheme, u, v) = @inbounds u[i, j, k] * _symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme.divergence_scheme, flux_div_xyᶜᶜᶜ, u, v) +@inline upwinded_divergence_flux_Vᶜᶠᶜ(i, j, k, grid, scheme, u, v) = @inbounds v[i, j, k] * _symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme.divergence_scheme, flux_div_xyᶜᶜᶜ, u, v) + ##### ##### Horizontal advection 4 formulations: ##### 1. Energy conservative @@ -346,7 +368,7 @@ end return - upwind_biased_product(v̂, ζᴸ, ζᴿ) end -@inline function horizontal_advection_V(i, j, k, grid, scheme::VectorInvariantUpwindVorticity, u, v) +@inline function horizontal_advection_V(i, j, k, grid, scheme::VectorInvariantUpwindVorticity, u, v) Sζ = scheme.vorticity_stencil diff --git a/src/Advection/vector_invariant_self_upwinding.jl b/src/Advection/vector_invariant_self_upwinding.jl index d8f798006e..daaa9f6534 100644 --- a/src/Advection/vector_invariant_self_upwinding.jl +++ b/src/Advection/vector_invariant_self_upwinding.jl @@ -26,7 +26,7 @@ end @inline function upwinded_divergence_flux_Vᶜᶠᶜ(i, j, k, grid, scheme::VectorInvariantSelfVerticalUpwinding, u, v) - + δV_stencil = scheme.upwinding.δV_stencil cross_scheme = scheme.upwinding.cross_scheme diff --git a/src/Advection/weno_default_stencils.jl b/src/Advection/weno_default_stencils.jl new file mode 100644 index 0000000000..18d9bbccfa --- /dev/null +++ b/src/Advection/weno_default_stencils.jl @@ -0,0 +1,455 @@ +@inline get_shifted_value_x(i, j, k, grid, shift, ψ, args...) = @inbounds ψ[i + shift, j, k] +@inline get_shifted_value_x(i, j, k, grid, shift, ψ::Function, args...) = ψ(i + shift, j, k, grid, args...) + +@inline get_shifted_value_y(i, j, k, grid, shift, ψ, args...) = @inbounds ψ[i, j + shift, k] +@inline get_shifted_value_y(i, j, k, grid, shift, ψ::Function, args...) = ψ(i, j + shift, k, grid, args...) + +@inline get_shifted_value_z(i, j, k, grid, shift, ψ, args...) = @inbounds ψ[i, j, k + shift] +@inline get_shifted_value_z(i, j, k, grid, shift, ψ::Function, args...) = ψ(i, j, k + shift, grid, args...) + +##### +##### STENCILS IN X +##### + +for (side, add) in zip([:left, :right], (1, 0)), (dir, loc, val) in zip((:x, :y, :z), (:ᶠᵃᵃ, :ᵃᶠᵃ, :ᵃᵃᶠ), (1, 2, 3)) + biased_interpolate = Symbol(:inner_, side, :_biased_interpolate_, dir, loc) + coeff = Symbol(:coeff_, side) + weno_interpolant = Symbol(side, :_weno_interpolant_, dir, loc) + get_shifted_value = Symbol(:get_shifted_value_, dir) + + @eval begin + # Fallback for DefaultStencil formulations and disambiguation + @inline $biased_interpolate(i, j, k, grid, scheme::WENO{2}, ψ, idx, loc, ::DefaultStencil, args...) = + $biased_interpolate(i, j, k, grid, scheme, ψ, idx, loc, args...) + @inline $biased_interpolate(i, j, k, grid, scheme::WENO{3}, ψ, idx, loc, ::DefaultStencil, args...) = + $biased_interpolate(i, j, k, grid, scheme, ψ, idx, loc, args...) + @inline $biased_interpolate(i, j, k, grid, scheme::WENO{4}, ψ, idx, loc, ::DefaultStencil, args...) = + $biased_interpolate(i, j, k, grid, scheme, ψ, idx, loc, args...) + @inline $biased_interpolate(i, j, k, grid, scheme::WENO{5}, ψ, idx, loc, ::DefaultStencil, args...) = + $biased_interpolate(i, j, k, grid, scheme, ψ, idx, loc, args...) + @inline $biased_interpolate(i, j, k, grid, scheme::WENO{6}, ψ, idx, loc, ::DefaultStencil, args...) = + $biased_interpolate(i, j, k, grid, scheme, ψ, idx, loc, args...) + + @inline function $biased_interpolate(i, j, k, grid, + scheme::WENO{2}, + ψ, idx, loc, args...) + + # All stencils + ψ₀ = $get_shifted_value(i, j, k, grid, - $add, ψ, args...) + ψ₁ = $get_shifted_value(i, j, k, grid, 1 - $add, ψ, args...) + + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁), 1, scheme, $val, idx, loc) + τ = add_to_global_smoothness(β, Val(2), Val(1)) + ψ̂₁ = ψ̅ * α + ψ̂₂ = ψ̅ * C + w₁ = α + + ψ₁ = ψ₀ + ψ₀ = $get_shifted_value(i, j, k, grid, - 1 - $add, ψ, args...) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁), 2, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(2), Val(2)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + τ = τ * τ + + return (ψ̂₁ * τ + ψ̂₂) / (w₁ * τ + 1) + end + + @inline function $biased_interpolate(i, j, k, grid, + scheme::WENO{3, FT, XT, YT, ZT}, + ψ, idx, loc, args...) where {FT, XT, YT, ZT} + + # All stencils + ψ₀ = $get_shifted_value(i, j, k, grid, - $add, ψ, args...) + ψ₁ = $get_shifted_value(i, j, k, grid, 1 - $add, ψ, args...) + ψ₂ = $get_shifted_value(i, j, k, grid, 2 - $add, ψ, args...) + + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂), 1, scheme, $val, idx, loc) + τ = add_to_global_smoothness(β, Val(3), Val(1)) + ψ̂₁ = ψ̅ * α + ψ̂₂ = ψ̅ * C + w₁ = α + + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = $get_shifted_value(i, j, k, grid, - 1 - $add, ψ, args...) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂), 2, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(3), Val(2)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = $get_shifted_value(i, j, k, grid, - 2 - $add, ψ, args...) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂), 3, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(3), Val(3)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + τ = τ * τ + + return (ψ̂₁ * τ + ψ̂₂) / (w₁ * τ + 1) + end + + @inline function $biased_interpolate(i, j, k, grid, + scheme::WENO{4, FT, XT, YT, ZT}, + ψ, idx, loc, args...) where {FT, XT, YT, ZT} + + # All stencils + ψ₀ = $get_shifted_value(i, j, k, grid, - $add, ψ, args...) + ψ₁ = $get_shifted_value(i, j, k, grid, 1 - $add, ψ, args...) + ψ₂ = $get_shifted_value(i, j, k, grid, 2 - $add, ψ, args...) + ψ₃ = $get_shifted_value(i, j, k, grid, 3 - $add, ψ, args...) + + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃), 1, scheme, $val, idx, loc) + τ = add_to_global_smoothness(β, Val(4), Val(1)) + ψ̂₁ = ψ̅ * α + ψ̂₂ = ψ̅ * C + w₁ = α + + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = $get_shifted_value(i, j, k, grid, - 1 - $add, ψ, args...) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃), 2, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(4), Val(2)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = $get_shifted_value(i, j, k, grid, - 2 - $add, ψ, args...) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃), 3, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(4), Val(3)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = $get_shifted_value(i, j, k, grid, - 3 - $add, ψ, args...) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃), 4, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(4), Val(4)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + τ = τ * τ + + return (ψ̂₁ * τ + ψ̂₂) / (w₁ * τ + 1) + end + + @inline function $biased_interpolate(i, j, k, grid, + scheme::WENO{5, FT, XT, YT, ZT}, + ψ, idx, loc, args...) where {FT, XT, YT, ZT} + + # All stencils + ψ₀ = $get_shifted_value(i, j, k, grid, - $add, ψ, args...) + ψ₁ = $get_shifted_value(i, j, k, grid, + 1 - $add, ψ, args...) + ψ₂ = $get_shifted_value(i, j, k, grid, + 2 - $add, ψ, args...) + ψ₃ = $get_shifted_value(i, j, k, grid, + 3 - $add, ψ, args...) + ψ₄ = $get_shifted_value(i, j, k, grid, + 4 - $add, ψ, args...) + + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃, ψ₄), 1, scheme, $val, idx, loc) + τ = β + ψ̂₁ = ψ̅ * α + ψ̂₂ = ψ̅ * C + w₁ = α + + ψ₄ = ψ₃ + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = $get_shifted_value(i, j, k, grid, - 1 - $add, ψ, args...) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃, ψ₄), 2, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(5), Val(2)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₄ = ψ₃ + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = $get_shifted_value(i, j, k, grid, - 2 - $add, ψ, args...) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃, ψ₄), 3, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(5), Val(3)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₄ = ψ₃ + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = $get_shifted_value(i, j, k, grid, - 3 - $add, ψ, args...) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃, ψ₄), 4, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(5), Val(4)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₄ = ψ₃ + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = $get_shifted_value(i, j, k, grid, - 4 - $add, ψ, args...) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃, ψ₄), 5, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(5), Val(5)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + τ = τ * τ + + return (ψ̂₁ * τ + ψ̂₂) / (w₁ * τ + 1) + end + end +end + +# # Metaprogrammed implementation of WENO reconstruction of order N + +# # Example: +# # julia> stencil(3, :ψ) +# # :((ψ0, ψ1, ψ2, ψ3)) +# # +# function stencil(N, v) +# expr = [] +# for i in 0:N-1 +# vs = Symbol(v, i) +# push!(expr, vs) +# end +# return Expr(:tuple, expr...) +# end + +# # julia> retrieve_initial_stencil(3, :x, 1, :ψ) +# # quote +# # ψ0 = get_shifted_value_x(i, j, k, grid, 0 - 1, ψ, args...) +# # ψ1 = get_shifted_value_x(i, j, k, grid, 1 - 1, ψ, args...) +# # ψ2 = get_shifted_value_x(i, j, k, grid, 2 - 1, ψ, args...) +# # end +# # +# # julia> +# function retrieve_initial_stencil(N, dir, add, variable) +# expr = [] +# get_shifted_value = Symbol(:get_shifted_value_, dir) +# for i in 0:N-1 +# ψ = Symbol(variable, i) +# push!(expr, :($ψ = $get_shifted_value(i, j, k, grid, $i - $add, $variable, args...))) +# end +# eblk = Expr(:block, expr...) + +# return eblk +# end + +# # julia> update_stencil(5, :y, 2, :𝓋) +# # quote +# # 𝓋4 = 𝓋3 +# # 𝓋3 = 𝓋2 +# # 𝓋2 = 𝓋1 +# # 𝓋1 = 𝓋0 +# # 𝓋0 = get_shifted_value_y(i, j, k, grid, (-s + 1) - 2, 𝓋, args...) +# # end +# # +# # julia> +# function update_stencil(N, dir, add, variable) +# expr = [] +# get_shifted_value = Symbol(:get_shifted_value_, dir) +# for i in N-1:-1:1 +# ψ = Symbol(variable, i) +# ψp = Symbol(variable, i-1) +# push!(expr, :($ψ = $ψp)) +# end +# ψ = Symbol(variable, 0) +# push!(expr, :($ψ = $get_shifted_value(i, j, k, grid, - s + 1 - $add, $variable, args...))) +# eblk = Expr(:block, expr...) + +# return eblk +# end + +# # julia> retrieve_velocity_stencil(3, :x, 1, :ψ) +# # quote +# # ψ0 = get_shifted_value_x(i, j, k, grid, 0 - 1, ψ, u, v, args...) +# # ψ1 = get_shifted_value_x(i, j, k, grid, 1 - 1, ψ, u, v, args...) +# # ψ2 = get_shifted_value_x(i, j, k, grid, 2 - 1, ψ, u, v, args...) +# # end +# # +# # julia> +# function retrieve_velocity_stencil(N, dir, add, variable) +# expr = [] +# get_shifted_value = Symbol(:get_shifted_value_, dir) +# for i in 0:N-1 +# ψ = Symbol(variable, i) +# push!(expr, :($ψ = $get_shifted_value(i, j, k, grid, $i - $add, $variable, u, v, args...))) +# end +# eblk = Expr(:block, expr...) + +# return eblk +# end + +# # julia> update_velocity_stencil(5, :y, 2, :𝓋) +# # quote +# # 𝓋4 = 𝓋3 +# # 𝓋3 = 𝓋2 +# # 𝓋2 = 𝓋1 +# # 𝓋1 = 𝓋0 +# # 𝓋0 = get_shifted_value_y(i, j, k, grid, (-s + 1) - 2, 𝓋, u, v, args...) +# # end +# # +# # julia> +# function update_velocity_stencil(N, dir, add, variable) +# expr = [] +# get_shifted_value = Symbol(:get_shifted_value_, dir) +# for i in N-1:-1:1 +# ψ = Symbol(variable, i) +# ψp = Symbol(variable, i-1) +# push!(expr, :($ψ = $ψp)) +# end +# ψ = Symbol(variable, 0) +# push!(expr, :($ψ = $get_shifted_value(i, j, k, grid, - s + 1 - $add, $variable, u, v, args...))) +# eblk = Expr(:block, expr...) + +# return eblk +# end + +# @inline ℑyᵃᶠᵃ_u(i, j, k, grid, u, v, args...) = ℑyᵃᶠᵃ(i, j, k, grid, u) +# @inline ℑxᶠᵃᵃ_v(i, j, k, grid, u, v, args...) = ℑxᶠᵃᵃ(i, j, k, grid, v) + +# for (side, add) in zip([:left, :right], (1, 0)), (dir, loc, val) in zip((:x, :y, :z), (:ᶠᵃᵃ, :ᵃᶠᵃ, :ᵃᵃᶠ), (1, 2, 3)) +# biased_interpolate = Symbol(:inner_, side, :_biased_interpolate_, dir, loc) +# weno_interpolant = Symbol(side, :_weno_interpolant_, dir, loc) + +# for N in [2, 3, 4, 5, 6] +# @eval begin +# @inline function $biased_interpolate(i, j, k, grid, +# scheme::WENO{$N}, +# ψ, idx, loc, args...) + +# # All stencils +# $(retrieve_initial_stencil(N, dir, add, :ψ)) + +# β, ψ̅, C, α = $weno_interpolant($(stencil(N, :ψ)), 1, scheme, $val, idx, loc) +# τ = β +# ψ̂₁ = ψ̅ * α +# ψ̂₂ = ψ̅ * C +# w₁ = α + +# @unroll for s in 2:$N +# $(update_stencil(N, dir, add, :ψ)) +# β, ψ̅, C, α = $weno_interpolant($(stencil(N, :ψ)), s, scheme, $val, idx, loc) +# τ += add_to_global_smoothness(β, Val($N), Val(s)) +# ψ̂₁ += ψ̅ * α +# ψ̂₂ += ψ̅ * C +# w₁ += α +# end + +# τ = τ * τ + +# return (ψ̂₁ * τ + ψ̂₂) / (w₁ * τ + 1) +# end + +# @inline function $biased_interpolate(i, j, k, grid, +# scheme::WENO{$N}, +# ψ, idx, loc, VI::FunctionStencil, args...) + +# ϕ = VI.func + +# # All stencils +# $(retrieve_initial_stencil(N, dir, add, :ψ)) +# $(retrieve_initial_stencil(N, dir, add, :ϕ)) + +# β, ψ̅, C, α = $weno_interpolant($(stencil(N, :ψ)), +# $(stencil(N, :ϕ)), 1, scheme, $val, idx, loc) +# τ = β +# ψ̂₁ = ψ̅ * α +# ψ̂₂ = ψ̅ * C +# w₁ = α + +# @unroll for s in 2:$N +# $(update_stencil(N, dir, add, :ψ)) +# $(update_stencil(N, dir, add, :ϕ)) + +# β, ψ̅, C, α = $weno_interpolant($(stencil(N, :ψ)), +# $(stencil(N, :ϕ)), +# s, scheme, $val, idx, loc) + +# τ += add_to_global_smoothness(β, Val($N), Val(s)) +# ψ̂₁ += ψ̅ * α +# ψ̂₂ += ψ̅ * C +# w₁ += α +# end + +# τ = τ * τ + +# return (ψ̂₁ * τ + ψ̂₂) / (w₁ * τ + 1) +# end + +# @inline function $biased_interpolate(i, j, k, grid, +# scheme::WENO{$N}, +# ψ, idx, loc, ::VelocityStencil, u, v, args...) + +# 𝓊 = ℑyᵃᶠᵃ_u +# 𝓋 = ℑxᶠᵃᵃ_v + +# # All stencils +# $(retrieve_velocity_stencil(N, dir, add, :ψ)) +# $(retrieve_velocity_stencil(N, dir, add, :𝓊)) +# $(retrieve_velocity_stencil(N, dir, add, :𝓋)) + +# β, ψ̅, C, α = $weno_interpolant($(stencil(N, :ψ)), +# $(stencil(N, :𝓊)), +# $(stencil(N, :𝓋)), 1, scheme, $val, idx, loc) +# τ = β +# ψ̂₁ = ψ̅ * α +# ψ̂₂ = ψ̅ * C +# w₁ = α + +# @unroll for s in 2:$N +# $(update_velocity_stencil(N, dir, add, :ψ)) +# $(update_velocity_stencil(N, dir, add, :𝓊)) +# $(update_velocity_stencil(N, dir, add, :𝓋)) + +# β, ψ̅, C, α = $weno_interpolant($(stencil(N, :ψ)), +# $(stencil(N, :𝓊)), +# $(stencil(N, :𝓋)), s, scheme, $val, idx, loc) + +# τ += add_to_global_smoothness(β, Val($N), Val(s)) +# ψ̂₁ += ψ̅ * α +# ψ̂₂ += ψ̅ * C +# w₁ += α +# end + +# τ = τ^2 + +# return (ψ̂₁ * τ + ψ̂₂) / (w₁ * τ + 1) +# end +# end +# end +# end \ No newline at end of file diff --git a/src/Advection/weno_function_stencils.jl b/src/Advection/weno_function_stencils.jl new file mode 100644 index 0000000000..eb0981e7ae --- /dev/null +++ b/src/Advection/weno_function_stencils.jl @@ -0,0 +1,283 @@ +##### +##### STENCILS IN X +##### + + +for (side, add) in zip([:left, :right], (1, 0)), (dir, loc, val) in zip((:x, :y, :z), (:ᶠᵃᵃ, :ᵃᶠᵃ, :ᵃᵃᶠ), (1, 2, 3)) + biased_interpolate = Symbol(:inner_, side, :_biased_interpolate_, dir, loc) + coeff = Symbol(:coeff_, side) + weno_interpolant = Symbol(side, :_weno_interpolant_, dir, loc) + get_shifted_value = Symbol(:get_shifted_value_, dir) + + @eval begin + @inline function $biased_interpolate(i, j, k, grid, + scheme::WENO{2, FT, XT, YT, ZT}, + ψ, idx, loc, VI::FunctionStencil, args...) where {FT, XT, YT, ZT} + + # All stencils + ψ₀ = $get_shifted_value(i, j, k, grid, - $add, ψ, args...) + ψ₁ = $get_shifted_value(i, j, k, grid, + 1 - $add, ψ, args...) + + ϕ₀ = $get_shifted_value(i, j, k, grid, - $add, VI.func, args...) + ϕ₁ = $get_shifted_value(i, j, k, grid, + 1 - $add, VI.func, args...) + + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁), (ϕ₀, ϕ₁), 1, scheme, $val, idx, loc) + τ = β + ψ̂₁ = ψ̅ * α + ψ̂₂ = ψ̅ * C + w₁ = α + + ψ₁ = ψ₀ + ψ₀ = $get_shifted_value(i, j, k, grid, - 1 - $add, ψ, args...) + + ϕ₁ = ϕ₀ + ϕ₀ = $get_shifted_value(i, j, k, grid, - 1 - $add, VI.func, args...) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁), (ϕ₀, ϕ₁), 2, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(2), Val(2)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + τ = τ^2 + + return (ψ̂₁ * τ + ψ̂₂) / (w₁ * τ + 1) + end + + @inline function $biased_interpolate(i, j, k, grid, + scheme::WENO{3, FT, XT, YT, ZT}, + ψ, idx, loc, VI::FunctionStencil, args...) where {FT, XT, YT, ZT} + + # All stencils + ψ₀ = $get_shifted_value(i, j, k, grid, - $add, ψ, args...) + ψ₁ = $get_shifted_value(i, j, k, grid, + 1 - $add, ψ, args...) + ψ₂ = $get_shifted_value(i, j, k, grid, + 2 - $add, ψ, args...) + + ϕ₀ = $get_shifted_value(i, j, k, grid, - $add, VI.func, args...) + ϕ₁ = $get_shifted_value(i, j, k, grid, + 1 - $add, VI.func, args...) + ϕ₂ = $get_shifted_value(i, j, k, grid, + 2 - $add, VI.func, args...) + + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂), (ϕ₀, ϕ₁, ϕ₂), 1, scheme, $val, idx, loc) + τ = β + ψ̂₁ = ψ̅ * α + ψ̂₂ = ψ̅ * C + w₁ = α + + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = $get_shifted_value(i, j, k, grid, - 1 - $add, ψ, args...) + + ϕ₂ = ϕ₁ + ϕ₁ = ϕ₀ + ϕ₀ = $get_shifted_value(i, j, k, grid, - 1 - $add, VI.func, args...) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂), (ϕ₀, ϕ₁, ϕ₂), 2, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(3), Val(2)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = $get_shifted_value(i, j, k, grid, - 2 - $add, ψ, args...) + + ϕ₂ = ϕ₁ + ϕ₁ = ϕ₀ + ϕ₀ = $get_shifted_value(i, j, k, grid, - 2 - $add, VI.func, args...) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂), (ϕ₀, ϕ₁, ϕ₂), 3, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(3), Val(3)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + τ = τ^2 + + return (ψ̂₁ * τ + ψ̂₂) / (w₁ * τ + 1) + end + + @inline function $biased_interpolate(i, j, k, grid, + scheme::WENO{4, FT, XT, YT, ZT}, + ψ, idx, loc, VI::FunctionStencil, args...) where {FT, XT, YT, ZT} + + # All stencils + ψ₀ = $get_shifted_value(i, j, k, grid, - $add, ψ, args...) + ψ₁ = $get_shifted_value(i, j, k, grid, + 1 - $add, ψ, args...) + ψ₂ = $get_shifted_value(i, j, k, grid, + 2 - $add, ψ, args...) + ψ₃ = $get_shifted_value(i, j, k, grid, + 3 - $add, ψ, args...) + + ϕ₀ = $get_shifted_value(i, j, k, grid, - $add, VI.func, args...) + ϕ₁ = $get_shifted_value(i, j, k, grid, + 1 - $add, VI.func, args...) + ϕ₂ = $get_shifted_value(i, j, k, grid, + 2 - $add, VI.func, args...) + ϕ₃ = $get_shifted_value(i, j, k, grid, + 3 - $add, VI.func, args...) + + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃), (ϕ₀, ϕ₁, ϕ₂, ϕ₃), 1, scheme, $val, idx, loc) + τ = β + ψ̂₁ = ψ̅ * α + ψ̂₂ = ψ̅ * C + w₁ = α + + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = $get_shifted_value(i, j, k, grid, - 1 - $add, ψ, args...) + + ϕ₃ = ϕ₂ + ϕ₂ = ϕ₁ + ϕ₁ = ϕ₀ + ϕ₀ = $get_shifted_value(i, j, k, grid, - 1 - $add, VI.func, args...) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃), (ϕ₀, ϕ₁, ϕ₂, ϕ₃), 2, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(4), Val(1)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = $get_shifted_value(i, j, k, grid, - 2 - $add, ψ, args...) + + ϕ₃ = ϕ₂ + ϕ₂ = ϕ₁ + ϕ₁ = ϕ₀ + ϕ₀ = $get_shifted_value(i, j, k, grid, - 2 - $add, VI.func, args...) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃), (ϕ₀, ϕ₁, ϕ₂, ϕ₃), 3, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(4), Val(3)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = $get_shifted_value(i, j, k, grid, - 3 - $add, ψ, args...) + + ϕ₃ = ϕ₂ + ϕ₂ = ϕ₁ + ϕ₁ = ϕ₀ + ϕ₀ = $get_shifted_value(i, j, k, grid, - 3 - $add, VI.func, args...) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃), (ϕ₀, ϕ₁, ϕ₂, ϕ₃), 4, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(4), Val(4)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + τ = τ^2 + + return (ψ̂₁ * τ + ψ̂₂) / (w₁ * τ + 1) + end + + @inline function $biased_interpolate(i, j, k, grid, + scheme::WENO{5, FT, XT, YT, ZT}, + ψ, idx, loc, VI::FunctionStencil, args...) where {FT, XT, YT, ZT} + + # All stencils + ψ₀ = $get_shifted_value(i, j, k, grid, - $add, ψ, args...) + ψ₁ = $get_shifted_value(i, j, k, grid, + 1 - $add, ψ, args...) + ψ₂ = $get_shifted_value(i, j, k, grid, + 2 - $add, ψ, args...) + ψ₃ = $get_shifted_value(i, j, k, grid, + 3 - $add, ψ, args...) + ψ₄ = $get_shifted_value(i, j, k, grid, + 4 - $add, ψ, args...) + + ϕ₀ = $get_shifted_value(i, j, k, grid, - $add, VI.func, args...) + ϕ₁ = $get_shifted_value(i, j, k, grid, + 1 - $add, VI.func, args...) + ϕ₂ = $get_shifted_value(i, j, k, grid, + 2 - $add, VI.func, args...) + ϕ₃ = $get_shifted_value(i, j, k, grid, + 3 - $add, VI.func, args...) + ϕ₄ = $get_shifted_value(i, j, k, grid, + 4 - $add, VI.func, args...) + + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃, ψ₄), (ϕ₀, ϕ₁, ϕ₂, ϕ₃, ϕ₄), 1, scheme, $val, idx, loc) + τ = β + ψ̂₁ = ψ̅ * α + ψ̂₂ = ψ̅ * C + w₁ = α + + ψ₄ = ψ₃ + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = $get_shifted_value(i, j, k, grid, - 1 - $add, ψ, args...) + + ϕ₄ = ϕ₃ + ϕ₃ = ϕ₂ + ϕ₂ = ϕ₁ + ϕ₁ = ϕ₀ + ϕ₀ = $get_shifted_value(i, j, k, grid, - 1 - $add, VI.func, args...) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃, ψ₄), (ϕ₀, ϕ₁, ϕ₂, ϕ₃, ϕ₄), 2, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(5), Val(2)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₄ = ψ₃ + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = $get_shifted_value(i, j, k, grid, - 2 - $add, ψ, args...) + + ϕ₄ = ϕ₃ + ϕ₃ = ϕ₂ + ϕ₂ = ϕ₁ + ϕ₁ = ϕ₀ + ϕ₀ = $get_shifted_value(i, j, k, grid, - 2 - $add, VI.func, args...) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃, ψ₄), (ϕ₀, ϕ₁, ϕ₂, ϕ₃, ϕ₄), 3, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(5), Val(3)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₄ = ψ₃ + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = $get_shifted_value(i, j, k, grid, - 3 - $add, ψ, args...) + + ϕ₄ = ϕ₃ + ϕ₃ = ϕ₂ + ϕ₂ = ϕ₁ + ϕ₁ = ϕ₀ + ϕ₀ = $get_shifted_value(i, j, k, grid, - 3 - $add, VI.func, args...) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃, ψ₄), (ϕ₀, ϕ₁, ϕ₂, ϕ₃, ϕ₄), 4, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(5), Val(4)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₄ = ψ₃ + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = $get_shifted_value(i, j, k, grid, - 4 - $add, ψ, args...) + + ϕ₄ = ϕ₃ + ϕ₃ = ϕ₂ + ϕ₂ = ϕ₁ + ϕ₁ = ϕ₀ + ϕ₀ = $get_shifted_value(i, j, k, grid, - 4 - $add, VI.func, args...) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃, ψ₄), (ϕ₀, ϕ₁, ϕ₂, ϕ₃, ϕ₄), 5, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(5), Val(5)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + τ = τ^2 + + return (ψ̂₁ * τ + ψ̂₂) / (w₁ * τ + 1) + end + end +end diff --git a/src/Advection/weno_interpolants.jl b/src/Advection/weno_interpolants.jl index 7c96b00adb..6a9e1c7561 100644 --- a/src/Advection/weno_interpolants.jl +++ b/src/Advection/weno_interpolants.jl @@ -1,3 +1,7 @@ +using Oceananigans.Fields: Field +using Oceananigans.Operators: ℑyᵃᶠᵃ, ℑxᶠᵃᵃ +using KernelAbstractions: @localmem, @index, @uniform, @groupsize + # WENO reconstruction of order `M` entails reconstructions of order `N` # on `N` different stencils, where `N = (M + 1) / 2`. # @@ -65,86 +69,66 @@ end Base.show(io::IO, a::FunctionStencil) = print(io, "FunctionStencil f = $(a.func)") -const ƞ = Int32(2) # WENO exponent const ε = 1e-8 # Optimal values taken from # Balsara & Shu, "Monotonicity Preserving Weighted Essentially Non-oscillatory Schemes with Inceasingly High Order of Accuracy" -@inline Cl(::WENO{2}, ::Val{0}) = 2/3 -@inline Cl(::WENO{2}, ::Val{1}) = 1/3 - -@inline Cl(::WENO{3}, ::Val{0}) = 3/10 -@inline Cl(::WENO{3}, ::Val{1}) = 3/5 -@inline Cl(::WENO{3}, ::Val{2}) = 1/10 - -@inline Cl(::WENO{4}, ::Val{0}) = 4/35 -@inline Cl(::WENO{4}, ::Val{1}) = 18/35 -@inline Cl(::WENO{4}, ::Val{2}) = 12/35 -@inline Cl(::WENO{4}, ::Val{3}) = 1/35 - -@inline Cl(::WENO{5}, ::Val{0}) = 5/126 -@inline Cl(::WENO{5}, ::Val{1}) = 20/63 -@inline Cl(::WENO{5}, ::Val{2}) = 10/21 -@inline Cl(::WENO{5}, ::Val{3}) = 10/63 -@inline Cl(::WENO{5}, ::Val{4}) = 1/126 - -@inline Cl(::WENO{6}, ::Val{0}) = 1/77 -@inline Cl(::WENO{6}, ::Val{1}) = 25/154 -@inline Cl(::WENO{6}, ::Val{2}) = 100/231 -@inline Cl(::WENO{6}, ::Val{3}) = 25/77 -@inline Cl(::WENO{6}, ::Val{4}) = 5/77 -@inline Cl(::WENO{6}, ::Val{5}) = 1/462 - -# ENO reconstruction procedure per stencil -for buffer in [2, 3, 4, 5, 6] - for stencil in collect(0:1:buffer-1) - - # ENO coefficients for uniform direction (when T<:Nothing) and stretched directions (when T<:Any) - @eval begin - @inline Cr(scheme::WENO{$buffer}, ::Val{$stencil}) = @inbounds Cl(scheme, Val($(buffer-stencil-1))) +@inline coeff_left(::WENO{2}, ::Val{0}) = 2/3 +@inline coeff_left(::WENO{2}, ::Val{1}) = 1/3 + +@inline coeff_left(::WENO{3}, ::Val{0}) = 3/10 +@inline coeff_left(::WENO{3}, ::Val{1}) = 3/5 +@inline coeff_left(::WENO{3}, ::Val{2}) = 1/10 + +@inline coeff_left(::WENO{4}, ::Val{0}) = 4/35 +@inline coeff_left(::WENO{4}, ::Val{1}) = 18/35 +@inline coeff_left(::WENO{4}, ::Val{2}) = 12/35 +@inline coeff_left(::WENO{4}, ::Val{3}) = 1/35 + +@inline coeff_left(::WENO{5}, ::Val{0}) = 5/126 +@inline coeff_left(::WENO{5}, ::Val{1}) = 20/63 +@inline coeff_left(::WENO{5}, ::Val{2}) = 10/21 +@inline coeff_left(::WENO{5}, ::Val{3}) = 10/63 +@inline coeff_left(::WENO{5}, ::Val{4}) = 1/126 + +@inline coeff_left(::WENO{6}, ::Val{0}) = 1/77 +@inline coeff_left(::WENO{6}, ::Val{1}) = 25/154 +@inline coeff_left(::WENO{6}, ::Val{2}) = 100/231 +@inline coeff_left(::WENO{6}, ::Val{3}) = 25/77 +@inline coeff_left(::WENO{6}, ::Val{4}) = 5/77 +@inline coeff_left(::WENO{6}, ::Val{5}) = 1/462 - # uniform coefficients are independent on direction and location - @inline coeff_left_p(scheme::WENO{$buffer, FT}, ::Val{$stencil}, ::Type{Nothing}, args...) where FT = @inbounds FT.($(stencil_coefficients(50, stencil , collect(1:100), collect(1:100); order = buffer))) - @inline coeff_right_p(scheme::WENO{$buffer, FT}, ::Val{$stencil}, ::Type{Nothing}, args...) where FT = @inbounds FT.($(stencil_coefficients(50, stencil-1, collect(1:100), collect(1:100); order = buffer))) - - # stretched coefficients are retrieved from precalculated coefficients - @inline coeff_left_p(scheme::WENO{$buffer}, ::Val{$stencil}, T, dir, i, loc) = @inbounds retrieve_coeff(scheme, $stencil, dir, i, loc) - @inline coeff_right_p(scheme::WENO{$buffer}, ::Val{$stencil}, T, dir, i, loc) = @inbounds retrieve_coeff(scheme, $(stencil-1), dir, i, loc) - end - - # left biased and right biased reconstruction value for each stencil - @eval begin - @inline left_biased_p(scheme::WENO{$buffer}, ::Val{$stencil}, ψ, T, dir, i, loc) = @inbounds sum(coeff_left_p(scheme, Val($stencil), T, dir, i, loc) .* ψ) - @inline right_biased_p(scheme::WENO{$buffer}, ::Val{$stencil}, ψ, T, dir, i, loc) = @inbounds sum(coeff_right_p(scheme, Val($stencil), T, dir, i, loc) .* ψ) - end +for buffer in [2, 3, 4, 5, 6] + for stencil in [0, 1, 2, 3, 4, 5] + @eval @inline coeff_right(scheme::WENO{$buffer}, ::Val{$stencil}) = @inbounds coeff_left(scheme, Val($(buffer-stencil-1))) end end # _UNIFORM_ smoothness coefficients (stretched smoothness coefficients are to be fixed!) -@inline smoothness_coefficients(scheme::WENO{2, FT}, ::Val{0}) where FT = @inbounds FT.((1, -2, 1)) -@inline smoothness_coefficients(scheme::WENO{2, FT}, ::Val{1}) where FT = @inbounds FT.((1, -2, 1)) - -@inline smoothness_coefficients(scheme::WENO{3, FT}, ::Val{0}) where FT = @inbounds FT.((10, -31, 11, 25, -19, 4)) -@inline smoothness_coefficients(scheme::WENO{3, FT}, ::Val{1}) where FT = @inbounds FT.((4, -13, 5, 13, -13, 4)) -@inline smoothness_coefficients(scheme::WENO{3, FT}, ::Val{2}) where FT = @inbounds FT.((4, -19, 11, 25, -31, 10)) - -@inline smoothness_coefficients(scheme::WENO{4, FT}, ::Val{0}) where FT = @inbounds FT.((2.107, -9.402, 7.042, -1.854, 11.003, -17.246, 4.642, 7.043, -3.882, 0.547)) -@inline smoothness_coefficients(scheme::WENO{4, FT}, ::Val{1}) where FT = @inbounds FT.((0.547, -2.522, 1.922, -0.494, 3.443, - 5.966, 1.602, 2.843, -1.642, 0.267)) -@inline smoothness_coefficients(scheme::WENO{4, FT}, ::Val{2}) where FT = @inbounds FT.((0.267, -1.642, 1.602, -0.494, 2.843, - 5.966, 1.922, 3.443, -2.522, 0.547)) -@inline smoothness_coefficients(scheme::WENO{4, FT}, ::Val{3}) where FT = @inbounds FT.((0.547, -3.882, 4.642, -1.854, 7.043, -17.246, 7.042, 11.003, -9.402, 2.107)) - -@inline smoothness_coefficients(scheme::WENO{5, FT}, ::Val{0}) where FT = @inbounds FT.((1.07918, -6.49501, 7.58823, -4.11487, 0.86329, 10.20563, -24.62076, 13.58458, -2.88007, 15.21393, -17.04396, 3.64863, 4.82963, -2.08501, 0.22658)) -@inline smoothness_coefficients(scheme::WENO{5, FT}, ::Val{1}) where FT = @inbounds FT.((0.22658, -1.40251, 1.65153, -0.88297, 0.18079, 2.42723, -6.11976, 3.37018, -0.70237, 4.06293, -4.64976, 0.99213, 1.38563, -0.60871, 0.06908)) -@inline smoothness_coefficients(scheme::WENO{5, FT}, ::Val{2}) where FT = @inbounds FT.((0.06908, -0.51001, 0.67923, -0.38947, 0.08209, 1.04963, -2.99076, 1.79098, -0.38947, 2.31153, -2.99076, 0.67923, 1.04963, -0.51001, 0.06908)) -@inline smoothness_coefficients(scheme::WENO{5, FT}, ::Val{3}) where FT = @inbounds FT.((0.06908, -0.60871, 0.99213, -0.70237, 0.18079, 1.38563, -4.64976, 3.37018, -0.88297, 4.06293, -6.11976, 1.65153, 2.42723, -1.40251, 0.22658)) -@inline smoothness_coefficients(scheme::WENO{5, FT}, ::Val{4}) where FT = @inbounds FT.((0.22658, -2.08501, 3.64863, -2.88007, 0.86329, 4.82963, -17.04396, 13.58458, -4.11487, 15.21393, -24.62076, 7.58823, 10.20563, -6.49501, 1.07918)) - -@inline smoothness_coefficients(scheme::WENO{6, FT}, ::Val{0}) where FT = @inbounds FT.((0.6150211, -4.7460464, 7.6206736, -6.3394124, 2.7060170, -0.4712740, 9.4851237, -31.1771244, 26.2901672, -11.3206788, 1.9834350, 26.0445372, -44.4003904, 19.2596472, -3.3918804, 19.0757572, -16.6461044, 2.9442256, 3.6480687, -1.2950184, 0.1152561)) -@inline smoothness_coefficients(scheme::WENO{6, FT}, ::Val{1}) where FT = @inbounds FT.((0.1152561, -0.9117992, 1.4742480, -1.2183636, 0.5134574, -0.0880548, 1.9365967, -6.5224244, 5.5053752, -2.3510468, 0.4067018, 5.6662212, -9.7838784, 4.2405032, -0.7408908, 4.3093692, -3.7913324, 0.6694608, 0.8449957, -0.3015728, 0.0271779)) -@inline smoothness_coefficients(scheme::WENO{6, FT}, ::Val{2}) where FT = @inbounds FT.((0.0271779, -0.2380800, 0.4086352, -0.3462252, 0.1458762, -0.0245620, 0.5653317, -2.0427884, 1.7905032, -0.7727988, 0.1325006, 1.9510972, -3.5817664, 1.5929912, -0.2792660, 1.7195652, -1.5880404, 0.2863984, 0.3824847, -0.1429976, 0.0139633)) -@inline smoothness_coefficients(scheme::WENO{6, FT}, ::Val{3}) where FT = @inbounds FT.((0.0139633, -0.1429976, 0.2863984, -0.2792660, 0.1325006, -0.0245620, 0.3824847, -1.5880404, 1.5929912, -0.7727988, 0.1458762, 1.7195652, -3.5817664, 1.7905032, -0.3462252, 1.9510972, -2.0427884, 0.4086352, 0.5653317, -0.2380800, 0.0271779)) -@inline smoothness_coefficients(scheme::WENO{6, FT}, ::Val{4}) where FT = @inbounds FT.((0.0271779, -0.3015728, 0.6694608, -0.7408908, 0.4067018, -0.0880548, 0.8449957, -3.7913324, 4.2405032, -2.3510468, 0.5134574, 4.3093692, -9.7838784, 5.5053752, -1.2183636, 5.6662212, -6.5224244, 1.4742480, 1.9365967, -0.9117992, 0.1152561)) -@inline smoothness_coefficients(scheme::WENO{6, FT}, ::Val{5}) where FT = @inbounds FT.((0.1152561, -1.2950184, 2.9442256, -3.3918804, 1.9834350, -0.4712740, 3.6480687, -16.6461044, 19.2596472, -11.3206788, 2.7060170, 19.0757572, -44.4003904, 26.2901672, -6.3394124, 26.0445372, -31.1771244, 7.6206736, 9.4851237, -4.7460464, 0.6150211)) +@inline smoothness_coefficients(::WENO{2, FT}, ::Val{0}) where FT = @inbounds FT.((1, -2, 1)) +@inline smoothness_coefficients(::WENO{2, FT}, ::Val{1}) where FT = @inbounds FT.((1, -2, 1)) + +@inline smoothness_coefficients(::WENO{3, FT}, ::Val{0}) where FT = @inbounds FT.((10, -31, 11, 25, -19, 4)) +@inline smoothness_coefficients(::WENO{3, FT}, ::Val{1}) where FT = @inbounds FT.((4, -13, 5, 13, -13, 4)) +@inline smoothness_coefficients(::WENO{3, FT}, ::Val{2}) where FT = @inbounds FT.((4, -19, 11, 25, -31, 10)) + +@inline smoothness_coefficients(::WENO{4, FT}, ::Val{0}) where FT = @inbounds FT.((2.107, -9.402, 7.042, -1.854, 11.003, -17.246, 4.642, 7.043, -3.882, 0.547)) +@inline smoothness_coefficients(::WENO{4, FT}, ::Val{1}) where FT = @inbounds FT.((0.547, -2.522, 1.922, -0.494, 3.443, - 5.966, 1.602, 2.843, -1.642, 0.267)) +@inline smoothness_coefficients(::WENO{4, FT}, ::Val{2}) where FT = @inbounds FT.((0.267, -1.642, 1.602, -0.494, 2.843, - 5.966, 1.922, 3.443, -2.522, 0.547)) +@inline smoothness_coefficients(::WENO{4, FT}, ::Val{3}) where FT = @inbounds FT.((0.547, -3.882, 4.642, -1.854, 7.043, -17.246, 7.042, 11.003, -9.402, 2.107)) + +@inline smoothness_coefficients(::WENO{5, FT}, ::Val{0}) where FT = @inbounds FT.((1.07918, -6.49501, 7.58823, -4.11487, 0.86329, 10.20563, -24.62076, 13.58458, -2.88007, 15.21393, -17.04396, 3.64863, 4.82963, -2.08501, 0.22658)) +@inline smoothness_coefficients(::WENO{5, FT}, ::Val{1}) where FT = @inbounds FT.((0.22658, -1.40251, 1.65153, -0.88297, 0.18079, 2.42723, -6.11976, 3.37018, -0.70237, 4.06293, -4.64976, 0.99213, 1.38563, -0.60871, 0.06908)) +@inline smoothness_coefficients(::WENO{5, FT}, ::Val{2}) where FT = @inbounds FT.((0.06908, -0.51001, 0.67923, -0.38947, 0.08209, 1.04963, -2.99076, 1.79098, -0.38947, 2.31153, -2.99076, 0.67923, 1.04963, -0.51001, 0.06908)) +@inline smoothness_coefficients(::WENO{5, FT}, ::Val{3}) where FT = @inbounds FT.((0.06908, -0.60871, 0.99213, -0.70237, 0.18079, 1.38563, -4.64976, 3.37018, -0.88297, 4.06293, -6.11976, 1.65153, 2.42723, -1.40251, 0.22658)) +@inline smoothness_coefficients(::WENO{5, FT}, ::Val{4}) where FT = @inbounds FT.((0.22658, -2.08501, 3.64863, -2.88007, 0.86329, 4.82963, -17.04396, 13.58458, -4.11487, 15.21393, -24.62076, 7.58823, 10.20563, -6.49501, 1.07918)) + +@inline smoothness_coefficients(::WENO{6, FT}, ::Val{0}) where FT = @inbounds FT.((0.6150211, -4.7460464, 7.6206736, -6.3394124, 2.7060170, -0.4712740, 9.4851237, -31.1771244, 26.2901672, -11.3206788, 1.9834350, 26.0445372, -44.4003904, 19.2596472, -3.3918804, 19.0757572, -16.6461044, 2.9442256, 3.6480687, -1.2950184, 0.1152561)) +@inline smoothness_coefficients(::WENO{6, FT}, ::Val{1}) where FT = @inbounds FT.((0.1152561, -0.9117992, 1.4742480, -1.2183636, 0.5134574, -0.0880548, 1.9365967, -6.5224244, 5.5053752, -2.3510468, 0.4067018, 5.6662212, -9.7838784, 4.2405032, -0.7408908, 4.3093692, -3.7913324, 0.6694608, 0.8449957, -0.3015728, 0.0271779)) +@inline smoothness_coefficients(::WENO{6, FT}, ::Val{2}) where FT = @inbounds FT.((0.0271779, -0.2380800, 0.4086352, -0.3462252, 0.1458762, -0.0245620, 0.5653317, -2.0427884, 1.7905032, -0.7727988, 0.1325006, 1.9510972, -3.5817664, 1.5929912, -0.2792660, 1.7195652, -1.5880404, 0.2863984, 0.3824847, -0.1429976, 0.0139633)) +@inline smoothness_coefficients(::WENO{6, FT}, ::Val{3}) where FT = @inbounds FT.((0.0139633, -0.1429976, 0.2863984, -0.2792660, 0.1325006, -0.0245620, 0.3824847, -1.5880404, 1.5929912, -0.7727988, 0.1458762, 1.7195652, -3.5817664, 1.7905032, -0.3462252, 1.9510972, -2.0427884, 0.4086352, 0.5653317, -0.2380800, 0.0271779)) +@inline smoothness_coefficients(::WENO{6, FT}, ::Val{4}) where FT = @inbounds FT.((0.0271779, -0.3015728, 0.6694608, -0.7408908, 0.4067018, -0.0880548, 0.8449957, -3.7913324, 4.2405032, -2.3510468, 0.5134574, 4.3093692, -9.7838784, 5.5053752, -1.2183636, 5.6662212, -6.5224244, 1.4742480, 1.9365967, -0.9117992, 0.1152561)) +@inline smoothness_coefficients(::WENO{6, FT}, ::Val{5}) where FT = @inbounds FT.((0.1152561, -1.2950184, 2.9442256, -3.3918804, 1.9834350, -0.4712740, 3.6480687, -16.6461044, 19.2596472, -11.3206788, 2.7060170, 19.0757572, -44.4003904, 26.2901672, -6.3394124, 26.0445372, -31.1771244, 7.6206736, 9.4851237, -4.7460464, 0.6150211)) # The rule for calculating smoothness indicators is the following (example WENO{4} which is seventh order) # ψ[1] (C[1] * ψ[1] + C[2] * ψ[2] + C[3] * ψ[3] + C[4] * ψ[4]) + @@ -158,12 +142,12 @@ end elem = Vector(undef, buffer) c_idx = 1 for stencil = 1:buffer - 1 - stencil_sum = Expr(:call, :+, (:(@inbounds C[$(c_idx + i - stencil)] * ψ[$i]) for i in stencil:buffer)...) - elem[stencil] = :(@inbounds ψ[$stencil] * $stencil_sum) + stencil_sum = Expr(:call, :+, (:(C[$(c_idx + i - stencil)] * ψ[$i]) for i in stencil:buffer)...) + elem[stencil] = :(ψ[$stencil] * $stencil_sum) c_idx += buffer - stencil + 1 end - elem[buffer] = :(@inbounds ψ[$buffer] * ψ[$buffer] * C[$c_idx]) + elem[buffer] = :(ψ[$buffer] * ψ[$buffer] * C[$c_idx]) return Expr(:call, :+, elem...) end @@ -171,272 +155,145 @@ end # Smoothness indicators for stencil `stencil` for left and right biased reconstruction for buffer in [2, 3, 4, 5, 6] @eval begin - @inline smoothness_sum(scheme::WENO{$buffer}, ψ, C) = @inbounds $(metaprogrammed_smoothness_sum(buffer)) + @inline smoothness_sum(scheme::WENO{$buffer}, ψ, C) = @inbounds @fastmath $(metaprogrammed_smoothness_sum(buffer)) end for stencil in [0, 1, 2, 3, 4, 5] @eval begin - @inline left_biased_β(ψ, scheme::WENO{$buffer}, ::Val{$stencil}) = @inbounds smoothness_sum(scheme, ψ, smoothness_coefficients(scheme, Val($stencil))) - @inline right_biased_β(ψ, scheme::WENO{$buffer}, ::Val{$stencil}) = @inbounds smoothness_sum(scheme, ψ, smoothness_coefficients(scheme, Val($stencil))) + @inline left_biased_β(ψ, scheme::WENO{$buffer}, ::Val{$stencil}) = smoothness_sum(scheme, ψ, smoothness_coefficients(scheme, Val($stencil))) + @inline right_biased_β(ψ, scheme::WENO{$buffer}, ::Val{$stencil}) = smoothness_sum(scheme, ψ, smoothness_coefficients(scheme, Val($stencil))) end - end -end - -# Shenanigans for WENO weights calculation for vector invariant formulation -> [β[i] = 0.5*(βᵤ[i] + βᵥ[i]) for i in 1:buffer] -@inline function metaprogrammed_beta_sum(buffer) - elem = Vector(undef, buffer) - for stencil = 1:buffer - elem[stencil] = :(@inbounds (β₁[$stencil] + β₂[$stencil])/2) - end - - return :($(elem...),) -end -# left and right biased_β calculation for scheme and stencil = 0:buffer - 1 -@inline function metaprogrammed_beta_loop(buffer) - elem = Vector(undef, buffer) - for stencil = 1:buffer - elem[stencil] = :(@inbounds func(ψ[$stencil], scheme, Val($(stencil-1)))) - end - - return :($(elem...),) -end - -# ZWENO α weights dᵣ * (1 + (τ₂ᵣ₋₁ / (βᵣ + ε))ᵖ) -@inline function metaprogrammed_zweno_alpha_loop(buffer) - elem = Vector(undef, buffer) - for stencil = 1:buffer - elem[stencil] = :(@inbounds FT(coeff(scheme, Val($(stencil-1)))) * (1 + (τ / (β[$stencil] + FT(ε)))^ƞ)) - end - - return :($(elem...),) -end - -# JSWENO α weights dᵣ / (βᵣ + ε)² -@inline function metaprogrammed_js_alpha_loop(buffer) - elem = Vector(undef, buffer) - for stencil = 1:buffer - elem[stencil] = :(@inbounds FT(coeff(scheme, Val($(stencil-1)))) / (β[$stencil] + FT(ε))^ƞ) - end - - return :($(elem...),) -end - -for buffer in [2, 3, 4, 5, 6] - @eval begin - @inline beta_sum(scheme::WENO{$buffer}, β₁, β₂) = @inbounds $(metaprogrammed_beta_sum(buffer)) - @inline beta_loop(scheme::WENO{$buffer}, ψ, func) = @inbounds $(metaprogrammed_beta_loop(buffer)) - @inline zweno_alpha_loop(scheme::WENO{$buffer}, β, τ, coeff, FT) = @inbounds $(metaprogrammed_zweno_alpha_loop(buffer)) - @inline js_alpha_loop(scheme::WENO{$buffer}, β, coeff, FT) = @inbounds $(metaprogrammed_js_alpha_loop(buffer)) - end -end - -# Global smoothness indicator τ₂ᵣ₋₁ taken from "Accuracy of the weighted essentially non-oscillatory conservative finite difference schemes", Don & Borges, 2013 -@inline global_smoothness_indicator(::Val{2}, β) = @inbounds abs(β[1] - β[2]) -@inline global_smoothness_indicator(::Val{3}, β) = @inbounds abs(β[1] - β[3]) -@inline global_smoothness_indicator(::Val{4}, β) = @inbounds abs(β[1] + 3β[2] - 3β[3] - β[4]) -@inline global_smoothness_indicator(::Val{5}, β) = @inbounds abs(β[1] + 2β[2] - 6β[3] + 2β[4] + β[5]) -@inline global_smoothness_indicator(::Val{6}, β) = @inbounds abs(β[1] + 36β[2] + 135β[3] - 135β[4] - 36β[5] - β[6]) - -# Calculating Dynamic WENO Weights (wᵣ), either with JS weno, Z weno or VectorInvariant WENO -for (side, coeff) in zip([:left, :right], (:Cl, :Cr)) - biased_weno_weights = Symbol(side, :_biased_weno_weights) - biased_β = Symbol(side, :_biased_β) - - tangential_stencil_u = Symbol(:tangential_, side, :_stencil_u) - tangential_stencil_v = Symbol(:tangential_, side, :_stencil_v) - - biased_stencil_z = Symbol(side, :_stencil_z) - - @eval begin - @inline function $biased_weno_weights(ψ, scheme::WENO{N, FT}, args...) where {N, FT} - @inbounds begin - β = beta_loop(scheme, ψ, $biased_β) - - if scheme isa ZWENO - τ = global_smoothness_indicator(Val(N), β) - α = zweno_alpha_loop(scheme, β, τ, $coeff, FT) - else - α = js_alpha_loop(scheme, β, $coeff, FT) - end - return α ./ sum(α) - end - end + # ENO coefficients for uniform direction (when T<:Nothing) and stretched directions (when T<:Any) + @eval begin + # uniform coefficients are independent on direction and location + @inline coeff_left_p(scheme::WENO{$buffer, FT}, ::Val{$stencil}, ::Type{Nothing}, args...) where FT = @inbounds FT.($(stencil_coefficients(50, stencil , collect(1:100), collect(1:100); order = buffer))) + @inline coeff_right_p(scheme::WENO{$buffer, FT}, ::Val{$stencil}, ::Type{Nothing}, args...) where FT = @inbounds FT.($(stencil_coefficients(50, stencil-1, collect(1:100), collect(1:100); order = buffer))) - @inline function $biased_weno_weights(ijk, scheme::WENO{N, FT}, dir, ::VelocityStencil, u, v) where {N, FT} - @inbounds begin - i, j, k = ijk - - uₛ = $tangential_stencil_u(i, j, k, scheme, dir, u) - vₛ = $tangential_stencil_v(i, j, k, scheme, dir, v) - βᵤ = beta_loop(scheme, uₛ, $biased_β) - βᵥ = beta_loop(scheme, vₛ, $biased_β) - - β = beta_sum(scheme, βᵤ, βᵥ) - - if scheme isa ZWENO - τ = global_smoothness_indicator(Val(N), β) - α = zweno_alpha_loop(scheme, β, τ, $coeff, FT) - else - α = js_alpha_loop(scheme, β, $coeff, FT) - end - return α ./ sum(α) - end + # stretched coefficients are retrieved from precalculated coefficients + @inline coeff_left_p(scheme::WENO{$buffer}, ::Val{$stencil}, T, dir, i, loc) = @inbounds retrieve_coeff(scheme, $stencil, dir, i, loc) + @inline coeff_right_p(scheme::WENO{$buffer}, ::Val{$stencil}, T, dir, i, loc) = @inbounds retrieve_coeff(scheme, $(stencil-1), dir, i, loc) end - end -end - -""" - calc_weno_stencil(buffer, shift, dir, func::Bool = false) - -Stencils for WENO reconstruction calculations - -The first argument is the `buffer`, not the `order`! -- `order = 2 * buffer - 1` for WENO reconstruction - -Examples -======== - -```jldoctest -julia> using Oceananigans.Advection: calc_weno_stencil - -julia> calc_weno_stencil(3, :left, :x) -:(((ψ[i + -1, j, k], ψ[i + 0, j, k], ψ[i + 1, j, k]), (ψ[i + -2, j, k], ψ[i + -1, j, k], ψ[i + 0, j, k]), (ψ[i + -3, j, k], ψ[i + -2, j, k], ψ[i + -1, j, k]))) - -julia> calc_weno_stencil(2, :right, :x) -:(((ψ[i + 0, j, k], ψ[i + 1, j, k]), (ψ[i + -1, j, k], ψ[i + 0, j, k]))) - -""" -@inline function calc_weno_stencil(buffer, shift, dir, func::Bool = false) - N = buffer * 2 - if shift != :none - N -=1 - end - stencil_full = Vector(undef, buffer) - rng = 1:N - if shift == :right - rng = rng .+ 1 - end - for stencil in 1:buffer - stencil_point = Vector(undef, buffer) - rngstencil = rng[stencil:stencil+buffer-1] - for (idx, n) in enumerate(rngstencil) - c = n - buffer - 1 - if func - stencil_point[idx] = dir == :x ? - :(@inbounds ψ(i + $c, j, k, args...)) : - dir == :y ? - :(@inbounds ψ(i, j + $c, k, args...)) : - :(@inbounds ψ(i, j, k + $c, args...)) - else - stencil_point[idx] = dir == :x ? - :(@inbounds ψ[i + $c, j, k]) : - dir == :y ? - :(@inbounds ψ[i, j + $c, k]) : - :(@inbounds ψ[i, j, k + $c]) - end + + # left biased and right biased reconstruction value for each stencil + @eval begin + @inline left_biased_p(scheme::WENO{$buffer}, ::Val{$stencil}, ψ, T, dir, i, loc) = @inbounds sum(coeff_left_p(scheme, Val($stencil), T, dir, i, loc) .* ψ) + @inline right_biased_p(scheme::WENO{$buffer}, ::Val{$stencil}, ψ, T, dir, i, loc) = @inbounds sum(coeff_right_p(scheme, Val($stencil), T, dir, i, loc) .* ψ) end - stencil_full[buffer - stencil + 1] = :($(stencil_point...), ) end - return :($(stencil_full...),) end -# Stencils for left and right biased reconstruction ((ψ̅ᵢ₋ᵣ₊ⱼ for j in 0:k) for r in 0:k) to calculate v̂ᵣ = ∑ⱼ(cᵣⱼψ̅ᵢ₋ᵣ₊ⱼ) -# where `k = N - 1`. Coefficients (cᵣⱼ for j in 0:N) for stencil r are given by `coeff_side_p(scheme, Val(r), ...)` -for side in (:left, :right), dir in (:x, :y, :z) - stencil = Symbol(side, :_stencil_, dir) - - for buffer in [2, 3, 4, 5, 6] +# Global smoothness indicator τ₂ᵣ₋₁ taken from "Accuracy of the weighted essentially non-oscillatory conservative finite difference schemes", Don & Borges, 2013 +@inline add_to_global_smoothness(β, ::Val{2}, ::Val{1}) = + β +@inline add_to_global_smoothness(β, ::Val{2}, ::Val{2}) = - β + +@inline add_to_global_smoothness(β, ::Val{3}, ::Val{1}) = + β +@inline add_to_global_smoothness(β, ::Val{3}, ::Val{2}) = 0 +@inline add_to_global_smoothness(β, ::Val{3}, ::Val{3}) = - β + +@inline add_to_global_smoothness(β, ::Val{4}, ::Val{1}) = + β +@inline add_to_global_smoothness(β, ::Val{4}, ::Val{2}) = + 3β +@inline add_to_global_smoothness(β, ::Val{4}, ::Val{3}) = - 3β +@inline add_to_global_smoothness(β, ::Val{4}, ::Val{4}) = - β + +@inline add_to_global_smoothness(β, ::Val{5}, ::Val{1}) = + β +@inline add_to_global_smoothness(β, ::Val{5}, ::Val{2}) = + 2β +@inline add_to_global_smoothness(β, ::Val{5}, ::Val{3}) = - 6β +@inline add_to_global_smoothness(β, ::Val{5}, ::Val{4}) = + 2β +@inline add_to_global_smoothness(β, ::Val{5}, ::Val{5}) = + β + +@inline add_to_global_smoothness(β, ::Val{6}, ::Val{1}) = + β +@inline add_to_global_smoothness(β, ::Val{6}, ::Val{2}) = + β +@inline add_to_global_smoothness(β, ::Val{6}, ::Val{3}) = - 8β +@inline add_to_global_smoothness(β, ::Val{6}, ::Val{4}) = + 8β +@inline add_to_global_smoothness(β, ::Val{6}, ::Val{5}) = - β +@inline add_to_global_smoothness(β, ::Val{6}, ::Val{6}) = - β + +# Reciprocal implementation from Valentin Churavy copied from +# https://github.com/JuliaGPU/CUDA.jl/blob/668b79c6378aa9443faf27d1efa3c117e0701077/perf/volumerhs.jl#L40 +let (jlf, f) = (:div_arcp, :div) + for (T, llvmT) in ((:Float32, "float"), (:Float64, "double")) + ir = """ + %x = f$f fast $llvmT %0, %1 + ret $llvmT %x + """ @eval begin - @inline $stencil(i, j, k, scheme::WENO{$buffer}, ψ, args...) = @inbounds $(calc_weno_stencil(buffer, side, dir, false)) - @inline $stencil(i, j, k, scheme::WENO{$buffer}, ψ::Function, args...) = @inbounds $(calc_weno_stencil(buffer, side, dir, true)) + # the @pure is necessary so that we can constant propagate. + @inline Base.@pure function $jlf(a::$T, b::$T) + Base.llvmcall($ir, $T, Tuple{$T, $T}, a, b) + end end end + @eval function $jlf(args...) + Base.$jlf(args...) + end end -# Stencil for vector invariant calculation of smoothness indicators in the horizontal direction -# Parallel to the interpolation direction! (same as left/right stencil) -using Oceananigans.Operators: ℑyᵃᶠᵃ, ℑxᶠᵃᵃ +rcp(x) = div_arcp(one(x), x) # still leads to rcp.rn which is also a function call -@inline tangential_left_stencil_u(i, j, k, scheme, ::Val{1}, u) = @inbounds left_stencil_x(i, j, k, scheme, ℑyᵃᶠᵃ, u) -@inline tangential_left_stencil_u(i, j, k, scheme, ::Val{2}, u) = @inbounds left_stencil_y(i, j, k, scheme, ℑyᵃᶠᵃ, u) -@inline tangential_left_stencil_v(i, j, k, scheme, ::Val{1}, v) = @inbounds left_stencil_x(i, j, k, scheme, ℑxᶠᵃᵃ, v) -@inline tangential_left_stencil_v(i, j, k, scheme, ::Val{2}, v) = @inbounds left_stencil_y(i, j, k, scheme, ℑxᶠᵃᵃ, v) +for side in [:left, :right], (dir, val, CT) in zip([:xᶠᵃᵃ, :yᵃᶠᵃ, :zᵃᵃᶠ], [1, 2, 3], [:XT, :YT, :ZT]) + weno_interpolant = Symbol(side, :_weno_interpolant_, dir) + biased_β = Symbol(side, :_biased_β) + biased_p = Symbol(side, :_biased_p) + coeff = Symbol(:coeff_, side) -@inline tangential_right_stencil_u(i, j, k, scheme, ::Val{1}, u) = @inbounds right_stencil_x(i, j, k, scheme, ℑyᵃᶠᵃ, u) -@inline tangential_right_stencil_u(i, j, k, scheme, ::Val{2}, u) = @inbounds right_stencil_y(i, j, k, scheme, ℑyᵃᶠᵃ, u) -@inline tangential_right_stencil_v(i, j, k, scheme, ::Val{1}, v) = @inbounds right_stencil_x(i, j, k, scheme, ℑxᶠᵃᵃ, v) -@inline tangential_right_stencil_v(i, j, k, scheme, ::Val{2}, v) = @inbounds right_stencil_y(i, j, k, scheme, ℑxᶠᵃᵃ, v) + @eval begin + ##### + ##### Picking up the stencil prior + ##### -# Trick to force compilation of Val(stencil-1) and avoid loops on the GPU -@inline function metaprogrammed_stencil_sum(buffer) - elem = Vector(undef, buffer) - for stencil = 1:buffer - elem[stencil] = :(@inbounds w[$stencil] * func(scheme, Val($(stencil-1)), ψ[$stencil], cT, Val(val), idx, loc)) - end + # Standard weno reconstruction at `i, j, k` for stencil `s` where `1 ≤ s ≤ N` + @inline function $weno_interpolant(ψs::Tuple, s, scheme::WENO{N, FT, XT, YT, ZT}, val, idx, loc) where {N, FT, XT, YT, ZT} + + # Calculate smoothness of stencil `s` + β = $biased_β(ψs, scheme, Val(s-1)) - return Expr(:call, :+, elem...) -end + # Calculate the `α` coefficient of stencil `s` following a WENO-JS formulation + C = FT($coeff(scheme, Val(s-1))) + α = C * rcp(β + FT(ε))^2 -# Calculation of WENO reconstructed value v⋆ = ∑ᵣ(wᵣv̂ᵣ) -for buffer in [2, 3, 4, 5, 6] - @eval begin - @inline stencil_sum(scheme::WENO{$buffer}, ψ, w, func, cT, val, idx, loc) = @inbounds $(metaprogrammed_stencil_sum(buffer)) - end -end + # Reconstruction of `ψ` from stencil `s` + ψ̅ = $biased_p(scheme, Val(s-1), ψs, $CT, Val(val), idx, loc) + + return β, ψ̅, C, α + end -# Interpolation functions -for (interp, dir, val, cT) in zip([:xᶠᵃᵃ, :yᵃᶠᵃ, :zᵃᵃᶠ], [:x, :y, :z], [1, 2, 3], [:XT, :YT, :ZT]) - for side in (:left, :right) - interpolate_func = Symbol(:inner_, side, :_biased_interpolate_, interp) - stencil = Symbol(side, :_stencil_, dir) - weno_weights = Symbol(side, :_biased_weno_weights) - biased_p = Symbol(side, :_biased_p) - - @eval begin - @inline function $interpolate_func(i, j, k, grid, - scheme::WENO{N, FT, XT, YT, ZT}, - ψ, idx, loc, args...) where {N, FT, XT, YT, ZT} - @inbounds begin - ψₜ = $stencil(i, j, k, scheme, ψ, grid, args...) - w = $weno_weights(ψₜ, scheme, Val($val), Nothing, args...) - return stencil_sum(scheme, ψₜ, w, $biased_p, $cT, $val, idx, loc) - end - end + # Using velocity interpolated at `(Face, Face, Center)` to assess smoothness. + # Can be used only for `(Face, Face, Center)` variables like vorticity + @inline function $weno_interpolant(ψs::Tuple, us::Tuple, vs::Tuple, s, scheme::WENO{N, FT, XT, YT, ZT}, val, idx, loc) where {N, FT, XT, YT, ZT} + + # Calculate x-velocity smoothness at stencil `s` + βu = $biased_β(us, scheme, Val(s-1)) + # Calculate y-velocity smoothness at stencil `s` + βv = $biased_β(vs, scheme, Val(s-1)) + + # total smoothness + βᵁ = (βu + βv) / 2 + + # Calculate the `α` coefficient of stencil `s` following a WENO-JS formulation + C = FT($coeff(scheme, Val(s-1))) + α = C * rcp(βᵁ + FT(ε))^2 - @inline function $interpolate_func(i, j, k, grid, - scheme::WENO{N, FT, XT, YT, ZT}, - ψ, idx, loc, VI::AbstractSmoothnessStencil, args...) where {N, FT, XT, YT, ZT} + # Retrieve stencil `s` and reconstruct `ψ` from stencil `s` + ψ̅ = $biased_p(scheme, Val(s-1), ψs, $CT, Val(val), idx, loc) - @inbounds begin - ψₜ = $stencil(i, j, k, scheme, ψ, grid, args...) - w = $weno_weights(ψₜ, scheme, Val($val), VI, args...) - return stencil_sum(scheme, ψₜ, w, $biased_p, $cT, $val, idx, loc) - end - end + return βᵁ, ψ̅, C, α + end - @inline function $interpolate_func(i, j, k, grid, - scheme::WENO{N, FT, XT, YT, ZT}, - ψ, idx, loc, VI::VelocityStencil, u, v, args...) where {N, FT, XT, YT, ZT} + # The smoothness is assessed using the stencil calculated from the function `VI.func(i, j, k, grid, args...)` + @inline function $weno_interpolant(ψs::Tuple, ϕs::Tuple, s, scheme::WENO{N, FT, XT, YT, ZT}, val, idx, loc) where {N, FT, XT, YT, ZT} + + # Calculate `ϕ` smoothness at `s` + βᵠ = $biased_β(ϕs, scheme, Val(s-1)) - @inbounds begin - ψₜ = $stencil(i, j, k, scheme, ψ, grid, u, v, args...) - w = $weno_weights((i, j, k), scheme, Val($val), VI, u, v) - return stencil_sum(scheme, ψₜ, w, $biased_p, $cT, $val, idx, loc) - end - end + # Calculate the `α` coefficient of stencil `s` following a WENO-JS formulation + C = FT($coeff(scheme, Val(s-1))) + α = C * rcp(βᵠ + FT(ε))^2 - @inline function $interpolate_func(i, j, k, grid, - scheme::WENO{N, FT, XT, YT, ZT}, - ψ, idx, loc, VI::FunctionStencil, args...) where {N, FT, XT, YT, ZT} + # Retrieve stencil `s` and reconstruct `ψ` from stencil `s` + ψ̅ = $biased_p(scheme, Val(s-1), ψs, $CT, Val(val), idx, loc) - @inbounds begin - ψₜ = $stencil(i, j, k, scheme, ψ, grid, args...) - ψₛ = $stencil(i, j, k, scheme, VI.func, grid, args...) - w = $weno_weights(ψₛ, scheme, Val($val), VI, args...) - return stencil_sum(scheme, ψₜ, w, $biased_p, $cT, $val, idx, loc) - end - end + return βᵠ, ψ̅, C, α end end end diff --git a/src/Advection/weno_reconstruction.jl b/src/Advection/weno_reconstruction.jl index 49e00f9a0b..b00f1a54ae 100644 --- a/src/Advection/weno_reconstruction.jl +++ b/src/Advection/weno_reconstruction.jl @@ -1,8 +1,9 @@ ##### ##### Weighted Essentially Non-Oscillatory (WENO) advection scheme ##### +using Oceananigans.Fields: CenterField -struct WENO{N, FT, XT, YT, ZT, WF, PP, CA, SI} <: AbstractUpwindBiasedAdvectionScheme{N, FT} +struct WENO{N, FT, XT, YT, ZT, PP, CA, SI} <: AbstractUpwindBiasedAdvectionScheme{N, FT} "Coefficient for ENO reconstruction on x-faces" coeff_xᶠᵃᵃ::XT @@ -25,16 +26,16 @@ struct WENO{N, FT, XT, YT, ZT, WF, PP, CA, SI} <: AbstractUpwindBiasedAdvectionS "Reconstruction scheme used for symmetric interpolation" advecting_velocity_scheme :: SI - function WENO{N, FT, WF}(coeff_xᶠᵃᵃ::XT, coeff_xᶜᵃᵃ::XT, - coeff_yᵃᶠᵃ::YT, coeff_yᵃᶜᵃ::YT, - coeff_zᵃᵃᶠ::ZT, coeff_zᵃᵃᶜ::ZT, - bounds::PP, buffer_scheme::CA, - advecting_velocity_scheme :: SI) where {N, FT, XT, YT, ZT, WF, PP, CA, SI} + function WENO{N, FT}(coeff_xᶠᵃᵃ::XT, coeff_xᶜᵃᵃ::XT, + coeff_yᵃᶠᵃ::YT, coeff_yᵃᶜᵃ::YT, + coeff_zᵃᵃᶠ::ZT, coeff_zᵃᵃᶜ::ZT, + bounds::PP, buffer_scheme::CA, + advecting_velocity_scheme :: SI) where {N, FT, XT, YT, ZT, PP, CA, SI} - return new{N, FT, XT, YT, ZT, WF, PP, CA, SI}(coeff_xᶠᵃᵃ, coeff_xᶜᵃᵃ, - coeff_yᵃᶠᵃ, coeff_yᵃᶜᵃ, - coeff_zᵃᵃᶠ, coeff_zᵃᵃᶜ, - bounds, buffer_scheme, advecting_velocity_scheme) + return new{N, FT, XT, YT, ZT, PP, CA, SI}(coeff_xᶠᵃᵃ, coeff_xᶜᵃᵃ, + coeff_yᵃᶠᵃ, coeff_yᵃᶜᵃ, + coeff_zᵃᵃᶠ, coeff_zᵃᵃᶜ, + bounds, buffer_scheme, advecting_velocity_scheme) end end @@ -103,7 +104,6 @@ WENO reconstruction order 7 function WENO(FT::DataType=Float64; order = 5, grid = nothing, - zweno = true, bounds = nothing) if !(grid isa Nothing) @@ -119,11 +119,12 @@ function WENO(FT::DataType=Float64; N = Int((order + 1) ÷ 2) weno_coefficients = compute_reconstruction_coefficients(grid, FT, :WENO; order = N) - buffer_scheme = WENO(FT; grid, order = order - 2, zweno, bounds) + buffer_scheme = WENO(FT; grid, order = order - 2, bounds) advecting_velocity_scheme = Centered(FT; grid, order = order - 1) end - return WENO{N, FT, zweno}(weno_coefficients..., bounds, buffer_scheme, advecting_velocity_scheme) + return WENO{N, FT}(weno_coefficients..., + bounds, buffer_scheme, advecting_velocity_scheme) end WENO(grid, FT::DataType=Float64; kwargs...) = WENO(FT; grid, kwargs...) @@ -133,15 +134,12 @@ WENOThirdOrder(grid=nothing, FT::DataType=Float64; kwargs...) = WENO(grid, FT; WENOFifthOrder(grid=nothing, FT::DataType=Float64; kwargs...) = WENO(grid, FT; order=5, kwargs...) # Flavours of WENO -const ZWENO = WENO{<:Any, <:Any, <:Any, <:Any, <:Any, true} -const PositiveWENO = WENO{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Tuple} +const PositiveWENO = WENO{<:Any, <:Any, <:Any, <:Any, <:Any, <:Tuple} Base.summary(a::WENO{N}) where N = string("WENO reconstruction order ", N*2-1) Base.show(io::IO, a::WENO{N, FT, RX, RY, RZ, WF, PP}) where {N, FT, RX, RY, RZ, WF, PP} = print(io, summary(a), " \n", - " Smoothness formulation: ", "\n", - " └── $(WF ? "Z-weno" : "JS-weno") \n", a.bounds isa Nothing ? "" : " Bounds : \n └── $(a.bounds) \n", " Boundary scheme: ", "\n", " └── ", summary(a.buffer_scheme) , "\n", @@ -152,21 +150,21 @@ Base.show(io::IO, a::WENO{N, FT, RX, RY, RZ, WF, PP}) where {N, FT, RX, RY, RZ, " ├── Y $(RY == Nothing ? "regular" : "stretched") \n", " └── Z $(RZ == Nothing ? "regular" : "stretched")" ) -Adapt.adapt_structure(to, scheme::WENO{N, FT, XT, YT, ZT, WF, PP}) where {N, FT, XT, YT, ZT, WF, PP} = - WENO{N, FT, 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.bounds), - Adapt.adapt(to, scheme.buffer_scheme), - Adapt.adapt(to, scheme.advecting_velocity_scheme)) - -on_architecture(to, scheme::WENO{N, FT, XT, YT, ZT, WF, PP}) where {N, FT, XT, YT, ZT, WF, PP} = - WENO{N, FT, WF}(on_architecture(to, scheme.coeff_xᶠᵃᵃ), on_architecture(to, scheme.coeff_xᶜᵃᵃ), - on_architecture(to, scheme.coeff_yᵃᶠᵃ), on_architecture(to, scheme.coeff_yᵃᶜᵃ), - on_architecture(to, scheme.coeff_zᵃᵃᶠ), on_architecture(to, scheme.coeff_zᵃᵃᶜ), - on_architecture(to, scheme.bounds), - on_architecture(to, scheme.buffer_scheme), - on_architecture(to, scheme.advecting_velocity_scheme)) +Adapt.adapt_structure(to, scheme::WENO{N, FT, XT, YT, ZT}) where {N, FT, XT, YT, ZT} = + WENO{N, FT}(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.bounds), + Adapt.adapt(to, scheme.buffer_scheme), + Adapt.adapt(to, scheme.advecting_velocity_scheme)) + +on_architecture(to, scheme::WENO{N, FT, XT, YT, ZT}) where {N, FT, XT, YT, ZT} = + WENO{N, FT}(on_architecture(to, scheme.coeff_xᶠᵃᵃ), on_architecture(to, scheme.coeff_xᶜᵃᵃ), + on_architecture(to, scheme.coeff_yᵃᶠᵃ), on_architecture(to, scheme.coeff_yᵃᶜᵃ), + on_architecture(to, scheme.coeff_zᵃᵃᶠ), on_architecture(to, scheme.coeff_zᵃᵃᶜ), + on_architecture(to, scheme.bounds), + on_architecture(to, scheme.buffer_scheme), + on_architecture(to, scheme.advecting_velocity_scheme)) # Retrieve precomputed coefficients (+2 for julia's 1 based indices) @inline retrieve_coeff(scheme::WENO, r, ::Val{1}, i, ::Type{Face}) = @inbounds scheme.coeff_xᶠᵃᵃ[r+2][i] diff --git a/src/Advection/weno_velocity_stencils.jl b/src/Advection/weno_velocity_stencils.jl new file mode 100644 index 0000000000..971f943732 --- /dev/null +++ b/src/Advection/weno_velocity_stencils.jl @@ -0,0 +1,878 @@ +##### +##### STENCILS IN X +##### + +for (side, add) in zip([:left, :right], (1, 0)) + biased_interpolate = Symbol(:inner_, side, :_biased_interpolate_xᶠᵃᵃ) + coeff = Symbol(:coeff_, side) + weno_interpolant = Symbol(side, :_weno_interpolant_xᶠᵃᵃ) + val = 1 + + @eval begin + @inline function $biased_interpolate(i, j, k, grid, + scheme::WENO{2, FT, XT, YT, ZT}, + ψ, idx, loc, ::VelocityStencil, u, v, args...) where {FT, XT, YT, ZT} + + # All stencils + ψ₀ = ψ(i - $add, j, k, grid, u, v) + ψ₁ = ψ(i + 1 - $add, j, k, grid, u, v) + + u₀ = ℑyᵃᶠᵃ(i - $add, j, k, grid, u) + u₁ = ℑyᵃᶠᵃ(i + 1 - $add, j, k, grid, u) + + v₀ = ℑxᶠᵃᵃ(i - $add, j, k, grid, v) + v₁ = ℑxᶠᵃᵃ(i + 1 - $add, j, k, grid, v) + + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁), (u₀, u₁), (v₀, v₁), 1, scheme, $val, idx, loc) + τ = β + ψ̂₁ = ψ̅ * α + ψ̂₂ = ψ̅ * C + w₁ = α + + ψ₁ = ψ₀ + ψ₀ = ψ(i - 1 - $add, j, k, grid, u, v) + + u₁ = u₀ + u₀ = ℑyᵃᶠᵃ(i - 1 - $add, j, k, grid, u) + + v₁ = v₀ + v₀ = ℑxᶠᵃᵃ(i - 1 - $add, j, k, grid, v) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁), (u₀, u₁), (v₀, v₁), 2, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(2), Val(2)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + τ = τ^2 + + return (ψ̂₁ * τ + ψ̂₂) / (w₁ * τ + 1) + end + + @inline function $biased_interpolate(i, j, k, grid, + scheme::WENO{3, FT, XT, YT, ZT}, + ψ, idx, loc, ::VelocityStencil, u, v, args...) where {FT, XT, YT, ZT} + + # All stencils + ψ₀ = ψ(i - $add, j, k, grid, u, v) + ψ₁ = ψ(i + 1 - $add, j, k, grid, u, v) + ψ₂ = ψ(i + 2 - $add, j, k, grid, u, v) + + u₀ = ℑyᵃᶠᵃ(i - $add, j, k, grid, u) + u₁ = ℑyᵃᶠᵃ(i + 1 - $add, j, k, grid, u) + u₂ = ℑyᵃᶠᵃ(i + 2 - $add, j, k, grid, u) + + v₀ = ℑxᶠᵃᵃ(i - $add, j, k, grid, v) + v₁ = ℑxᶠᵃᵃ(i + 1 - $add, j, k, grid, v) + v₂ = ℑxᶠᵃᵃ(i + 2 - $add, j, k, grid, v) + + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂), (u₀, u₁, u₂), (v₀, v₁, v₂), 1, scheme, $val, idx, loc) + τ = β + ψ̂₁ = ψ̅ * α + ψ̂₂ = ψ̅ * C + w₁ = α + + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = ψ(i - 1 - $add, j, k, grid, u, v) + + u₂ = u₁ + u₁ = u₀ + u₀ = ℑyᵃᶠᵃ(i - 1 - $add, j, k, grid, u) + + v₂ = v₁ + v₁ = v₀ + v₀ = ℑxᶠᵃᵃ(i - 1 - $add, j, k, grid, v) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂), (u₀, u₁, u₂), (v₀, v₁, v₂), 2, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(3), Val(2)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = ψ(i - 2 - $add, j, k, grid, u, v) + + u₂ = u₁ + u₁ = u₀ + u₀ = ℑyᵃᶠᵃ(i - 2 - $add, j, k, grid, u) + + v₂ = v₁ + v₁ = v₀ + v₀ = ℑxᶠᵃᵃ(i - 2 - $add, j, k, grid, v) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂), (u₀, u₁, u₂), (v₀, v₁, v₂), 3, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(3), Val(3)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + τ = τ^2 + + return (ψ̂₁ * τ + ψ̂₂) / (w₁ * τ + 1) + end + + @inline function $biased_interpolate(i, j, k, grid, + scheme::WENO{4, FT, XT, YT, ZT}, + ψ, idx, loc, ::VelocityStencil, u, v, args...) where {FT, XT, YT, ZT} + + # All stencils + ψ₀ = ψ(i - $add, j, k, grid, u, v) + ψ₁ = ψ(i + 1 - $add, j, k, grid, u, v) + ψ₂ = ψ(i + 2 - $add, j, k, grid, u, v) + ψ₃ = ψ(i + 3 - $add, j, k, grid, u, v) + + u₀ = ℑyᵃᶠᵃ(i - $add, j, k, grid, u) + u₁ = ℑyᵃᶠᵃ(i + 1 - $add, j, k, grid, u) + u₂ = ℑyᵃᶠᵃ(i + 2 - $add, j, k, grid, u) + u₃ = ℑyᵃᶠᵃ(i + 3 - $add, j, k, grid, u) + + v₀ = ℑxᶠᵃᵃ(i - $add, j, k, grid, v) + v₁ = ℑxᶠᵃᵃ(i + 1 - $add, j, k, grid, v) + v₂ = ℑxᶠᵃᵃ(i + 2 - $add, j, k, grid, v) + v₃ = ℑxᶠᵃᵃ(i + 3 - $add, j, k, grid, v) + + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃), (u₀, u₁, u₂, u₃), (v₀, v₁, v₂, v₃), 1, scheme, $val, idx, loc) + τ = β + ψ̂₁ = ψ̅ * α + ψ̂₂ = ψ̅ * C + w₁ = α + + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = ψ(i - 1 - $add, j, k, grid, u, v) + + u₃ = u₂ + u₂ = u₁ + u₁ = u₀ + u₀ = ℑyᵃᶠᵃ(i - 1 - $add, j, k, grid, u) + + v₃ = v₂ + v₂ = v₁ + v₁ = v₀ + v₀ = ℑxᶠᵃᵃ(i - 1 - $add, j, k, grid, v) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃), (u₀, u₁, u₂, u₃), (v₀, v₁, v₂, v₃), 2, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(4), Val(2)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = ψ(i - 2 - $add, j, k, grid, u, v) + + u₃ = u₂ + u₂ = u₁ + u₁ = u₀ + u₀ = ℑyᵃᶠᵃ(i - 2 - $add, j, k, grid, u) + + v₃ = v₂ + v₂ = v₁ + v₁ = v₀ + v₀ = ℑxᶠᵃᵃ(i - 2 - $add, j, k, grid, v) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃), (u₀, u₁, u₂, u₃), (v₀, v₁, v₂, v₃), 3, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(4), Val(3)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = ψ(i - 3 - $add, j, k, grid, u, v) + + u₃ = u₂ + u₂ = u₁ + u₁ = u₀ + u₀ = ℑyᵃᶠᵃ(i - 3 - $add, j, k, grid, u) + + v₃ = v₂ + v₂ = v₁ + v₁ = v₀ + v₀ = ℑxᶠᵃᵃ(i - 3 - $add, j, k, grid, v) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃), (u₀, u₁, u₂, u₃), (v₀, v₁, v₂, v₃), 4, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(4), Val(4)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + τ = τ^2 + + return (ψ̂₁ * τ + ψ̂₂) / (w₁ * τ + 1) + end + + @inline function $biased_interpolate(i, j, k, grid, + scheme::WENO{5, FT, XT, YT, ZT}, + ψ, idx, loc, ::VelocityStencil, u, v, args...) where {FT, XT, YT, ZT} + + # All stencils + ψ₀ = ψ(i - $add, j, k, grid, u, v) + ψ₁ = ψ(i + 1 - $add, j, k, grid, u, v) + ψ₂ = ψ(i + 2 - $add, j, k, grid, u, v) + ψ₃ = ψ(i + 3 - $add, j, k, grid, u, v) + ψ₄ = ψ(i + 4 - $add, j, k, grid, u, v) + + u₀ = ℑyᵃᶠᵃ(i - $add, j, k, grid, u) + u₁ = ℑyᵃᶠᵃ(i + 1 - $add, j, k, grid, u) + u₂ = ℑyᵃᶠᵃ(i + 2 - $add, j, k, grid, u) + u₃ = ℑyᵃᶠᵃ(i + 3 - $add, j, k, grid, u) + u₄ = ℑyᵃᶠᵃ(i + 4 - $add, j, k, grid, u) + + v₀ = ℑxᶠᵃᵃ(i - $add, j, k, grid, v) + v₁ = ℑxᶠᵃᵃ(i + 1 - $add, j, k, grid, v) + v₂ = ℑxᶠᵃᵃ(i + 2 - $add, j, k, grid, v) + v₃ = ℑxᶠᵃᵃ(i + 3 - $add, j, k, grid, v) + v₄ = ℑxᶠᵃᵃ(i + 4 - $add, j, k, grid, v) + + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃, ψ₄), (u₀, u₁, u₂, u₃, u₄), (v₀, v₁, v₂, v₃, v₄), 1, scheme, $val, idx, loc) + τ = β + ψ̂₁ = ψ̅ * α + ψ̂₂ = ψ̅ * C + w₁ = α + + ψ₄ = ψ₃ + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = ψ(i - 1 - $add, j, k, grid, u, v) + + u₄ = u₃ + u₃ = u₂ + u₂ = u₁ + u₁ = u₀ + u₀ = ℑyᵃᶠᵃ(i - 1 - $add, j, k, grid, u) + + v₄ = v₃ + v₃ = v₂ + v₂ = v₁ + v₁ = v₀ + v₀ = ℑxᶠᵃᵃ(i - 1 - $add, j, k, grid, v) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃, ψ₄), (u₀, u₁, u₂, u₃, u₄), (v₀, v₁, v₂, v₃, v₄), 2, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(5), Val(2)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₄ = ψ₃ + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = ψ(i - 2 - $add, j, k, grid, u, v) + + u₄ = u₃ + u₃ = u₂ + u₂ = u₁ + u₁ = u₀ + u₀ = ℑyᵃᶠᵃ(i - 2 - $add, j, k, grid, u) + + v₄ = v₃ + v₃ = v₂ + v₂ = v₁ + v₁ = v₀ + v₀ = ℑxᶠᵃᵃ(i - 2 - $add, j, k, grid, v) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃, ψ₄), (u₀, u₁, u₂, u₃, u₄), (v₀, v₁, v₂, v₃, v₄), 3, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(5), Val(3)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₄ = ψ₃ + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = ψ(i - 3 - $add, j, k, grid, u, v) + + u₄ = u₃ + u₃ = u₂ + u₂ = u₁ + u₁ = u₀ + u₀ = ℑyᵃᶠᵃ(i - 3 - $add, j, k, grid, u) + + v₄ = v₃ + v₃ = v₂ + v₂ = v₁ + v₁ = v₀ + v₀ = ℑxᶠᵃᵃ(i - 3 - $add, j, k, grid, v) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃, ψ₄), (u₀, u₁, u₂, u₃, u₄), (v₀, v₁, v₂, v₃, v₄), 4, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(5), Val(4)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₄ = ψ₃ + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = ψ(i - 4 - $add, j, k, grid, u, v) + + u₄ = u₃ + u₃ = u₂ + u₂ = u₁ + u₁ = u₀ + u₀ = ℑyᵃᶠᵃ(i - 4 - $add, j, k, grid, u) + + v₄ = v₃ + v₃ = v₂ + v₂ = v₁ + v₁ = v₀ + v₀ = ℑxᶠᵃᵃ(i - 4 - $add, j, k, grid, v) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃, ψ₄), (u₀, u₁, u₂, u₃, u₄), (v₀, v₁, v₂, v₃, v₄), 5, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(5), Val(5)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + τ = τ^2 + + return (ψ̂₁ * τ + ψ̂₂) / (w₁ * τ + 1) + end + end +end + +##### +##### STENCILS IN Y +##### + +for (side, add) in zip([:left, :right], (1, 0)) + biased_interpolate = Symbol(:inner_, side, :_biased_interpolate_yᵃᶠᵃ) + coeff = Symbol(:coeff_, side) + weno_interpolant = Symbol(side, :_weno_interpolant_yᵃᶠᵃ) + val = 2 + + @eval begin + @inline function $biased_interpolate(i, j, k, grid, + scheme::WENO{2, FT, XT, YT, ZT}, + ψ, idx, loc, ::VelocityStencil, u, v, args...) where {FT, XT, YT, ZT} + + # All stencils + ψ₀ = ψ(i, j - $add, k, grid, u, v) + ψ₁ = ψ(i, j + 1 - $add, k, grid, u, v) + + u₀ = ℑyᵃᶠᵃ(i, j - $add, k, grid, u) + u₁ = ℑyᵃᶠᵃ(i, j + 1 - $add, k, grid, u) + + v₀ = ℑxᶠᵃᵃ(i, j - $add, k, grid, v) + v₁ = ℑxᶠᵃᵃ(i, j + 1 - $add, k, grid, v) + + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁), (u₀, u₁), (v₀, v₁), 1, scheme, $val, idx, loc) + τ = β + ψ̂₁ = ψ̅ * α + ψ̂₂ = ψ̅ * C + w₁ = α + + ψ₁ = ψ₀ + ψ₀ = ψ(i, j - 1 - $add, k, grid, u, v) + + u₁ = u₀ + u₀ = ℑyᵃᶠᵃ(i, j - 1 - $add, k, grid, u) + + v₁ = v₀ + v₀ = ℑxᶠᵃᵃ(i, j - 1 - $add, k, grid, v) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁), (u₀, u₁), (v₀, v₁), 2, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(2), Val(2)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + τ = τ^2 + + return (ψ̂₁ * τ + ψ̂₂) / (w₁ * τ + 1) + end + + @inline function $biased_interpolate(i, j, k, grid, + scheme::WENO{3, FT, XT, YT, ZT}, + ψ, idx, loc, ::VelocityStencil, u, v, args...) where {FT, XT, YT, ZT} + + # All stencils + ψ₀ = ψ(i, j - $add, k, grid, u, v) + ψ₁ = ψ(i, j + 1 - $add, k, grid, u, v) + ψ₂ = ψ(i, j + 2 - $add, k, grid, u, v) + + u₀ = ℑyᵃᶠᵃ(i, j - $add, k, grid, u) + u₁ = ℑyᵃᶠᵃ(i, j + 1 - $add, k, grid, u) + u₂ = ℑyᵃᶠᵃ(i, j + 2 - $add, k, grid, u) + + v₀ = ℑxᶠᵃᵃ(i, j - $add, k, grid, v) + v₁ = ℑxᶠᵃᵃ(i, j + 1 - $add, k, grid, v) + v₂ = ℑxᶠᵃᵃ(i, j + 2 - $add, k, grid, v) + + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂), (u₀, u₁, u₂), (v₀, v₁, v₂), 1, scheme, $val, idx, loc) + τ = β + ψ̂₁ = ψ̅ * α + ψ̂₂ = ψ̅ * C + w₁ = α + + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = ψ(i, j - 1 - $add, k, grid, u, v) + + u₂ = u₁ + u₁ = u₀ + u₀ = ℑyᵃᶠᵃ(i, j - 1 - $add, k, grid, u) + + v₂ = v₁ + v₁ = v₀ + v₀ = ℑxᶠᵃᵃ(i, j - 1 - $add, k, grid, v) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂), (u₀, u₁, u₂), (v₀, v₁, v₂), 2, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(3), Val(2)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = ψ(i, j - 2 - $add, k, grid, u, v) + + u₂ = u₁ + u₁ = u₀ + u₀ = ℑyᵃᶠᵃ(i, j - 2 - $add, k, grid, u) + + v₂ = v₁ + v₁ = v₀ + v₀ = ℑxᶠᵃᵃ(i, j - 2 - $add, k, grid, v) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂), (u₀, u₁, u₂), (v₀, v₁, v₂), 3, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(3), Val(3)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + τ = τ^2 + + return (ψ̂₁ * τ + ψ̂₂) / (w₁ * τ + 1) + end + + @inline function $biased_interpolate(i, j, k, grid, + scheme::WENO{4, FT, XT, YT, ZT}, + ψ, idx, loc, ::VelocityStencil, u, v, args...) where {FT, XT, YT, ZT} + + # All stencils + ψ₀ = ψ(i, j - $add, k, grid, u, v) + ψ₁ = ψ(i, j + 1 - $add, k, grid, u, v) + ψ₂ = ψ(i, j + 2 - $add, k, grid, u, v) + ψ₃ = ψ(i, j + 3 - $add, k, grid, u, v) + + u₀ = ℑyᵃᶠᵃ(i, j - $add, k, grid, u) + u₁ = ℑyᵃᶠᵃ(i, j + 1 - $add, k, grid, u) + u₂ = ℑyᵃᶠᵃ(i, j + 2 - $add, k, grid, u) + u₃ = ℑyᵃᶠᵃ(i, j + 3 - $add, k, grid, u) + + v₀ = ℑxᶠᵃᵃ(i, j - $add, k, grid, v) + v₁ = ℑxᶠᵃᵃ(i, j + 1 - $add, k, grid, v) + v₂ = ℑxᶠᵃᵃ(i, j + 2 - $add, k, grid, v) + v₃ = ℑxᶠᵃᵃ(i, j + 3 - $add, k, grid, v) + + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃), (u₀, u₁, u₂, u₃), (v₀, v₁, v₂, v₃), 1, scheme, $val, idx, loc) + τ = β + ψ̂₁ = ψ̅ * α + ψ̂₂ = ψ̅ * C + w₁ = α + + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = ψ(i, j - 1 - $add, k, grid, u, v) + + u₃ = u₂ + u₂ = u₁ + u₁ = u₀ + u₀ = ℑyᵃᶠᵃ(i, j - 1 - $add, k, grid, u) + + v₃ = v₂ + v₂ = v₁ + v₁ = v₀ + v₀ = ℑxᶠᵃᵃ(i, j - 1 - $add, k, grid, v) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃), (u₀, u₁, u₂, u₃), (v₀, v₁, v₂, v₃), 2, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(4), Val(2)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = ψ(i, j - 2 - $add, k, grid, u, v) + + u₃ = u₂ + u₂ = u₁ + u₁ = u₀ + u₀ = ℑyᵃᶠᵃ(i, j - 2 - $add, k, grid, u) + + v₃ = v₂ + v₂ = v₁ + v₁ = v₀ + v₀ = ℑxᶠᵃᵃ(i, j - 2 - $add, k, grid, v) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃), (u₀, u₁, u₂, u₃), (v₀, v₁, v₂, v₃), 3, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(4), Val(3)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = ψ(i, j - 3 - $add, k, grid, u, v) + + u₃ = u₂ + u₂ = u₁ + u₁ = u₀ + u₀ = ℑyᵃᶠᵃ(i, j - 3 - $add, k, grid, u) + + v₃ = v₂ + v₂ = v₁ + v₁ = v₀ + v₀ = ℑxᶠᵃᵃ(i, j - 3 - $add, k, grid, v) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃), (u₀, u₁, u₂, u₃), (v₀, v₁, v₂, v₃), 4, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(4), Val(4)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + τ = τ^2 + + return (ψ̂₁ * τ + ψ̂₂) / (w₁ * τ + 1) + end + + @inline function $biased_interpolate(i, j, k, grid, + scheme::WENO{5, FT, XT, YT, ZT}, + ψ, idx, loc, ::VelocityStencil, u, v, args...) where {FT, XT, YT, ZT} + + # All stencils + ψ₀ = ψ(i, j - $add, k, grid, u, v) + ψ₁ = ψ(i, j + 1 - $add, k, grid, u, v) + ψ₂ = ψ(i, j + 2 - $add, k, grid, u, v) + ψ₃ = ψ(i, j + 3 - $add, k, grid, u, v) + ψ₄ = ψ(i, j + 4 - $add, k, grid, u, v) + + u₀ = ℑyᵃᶠᵃ(i, j - $add, k, grid, u) + u₁ = ℑyᵃᶠᵃ(i, j + 1 - $add, k, grid, u) + u₂ = ℑyᵃᶠᵃ(i, j + 2 - $add, k, grid, u) + u₃ = ℑyᵃᶠᵃ(i, j + 3 - $add, k, grid, u) + u₄ = ℑyᵃᶠᵃ(i, j + 4 - $add, k, grid, u) + + v₀ = ℑxᶠᵃᵃ(i, j - $add, k, grid, v) + v₁ = ℑxᶠᵃᵃ(i, j + 1 - $add, k, grid, v) + v₂ = ℑxᶠᵃᵃ(i, j + 2 - $add, k, grid, v) + v₃ = ℑxᶠᵃᵃ(i, j + 3 - $add, k, grid, v) + v₄ = ℑxᶠᵃᵃ(i, j + 4 - $add, k, grid, v) + + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃, ψ₄), (u₀, u₁, u₂, u₃, u₄), (v₀, v₁, v₂, v₃, v₄), 1, scheme, $val, idx, loc) + τ = β + ψ̂₁ = ψ̅ * α + ψ̂₂ = ψ̅ * C + w₁ = α + + ψ₄ = ψ₃ + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = ψ(i, j - 1 - $add, k, grid, u, v) + + u₄ = u₃ + u₃ = u₂ + u₂ = u₁ + u₁ = u₀ + u₀ = ℑyᵃᶠᵃ(i, j - 1 - $add, k, grid, u) + + v₄ = v₃ + v₃ = v₂ + v₂ = v₁ + v₁ = v₀ + v₀ = ℑxᶠᵃᵃ(i, j - 1 - $add, k, grid, v) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃, ψ₄), (u₀, u₁, u₂, u₃, u₄), (v₀, v₁, v₂, v₃, v₄), 2, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(5), Val(2)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₄ = ψ₃ + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = ψ(i, j - 2 - $add, k, grid, u, v) + + u₄ = u₃ + u₃ = u₂ + u₂ = u₁ + u₁ = u₀ + u₀ = ℑyᵃᶠᵃ(i, j - 2 - $add, k, grid, u) + + v₄ = v₃ + v₃ = v₂ + v₂ = v₁ + v₁ = v₀ + v₀ = ℑxᶠᵃᵃ(i, j - 2 - $add, k, grid, v) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃, ψ₄), (u₀, u₁, u₂, u₃, u₄), (v₀, v₁, v₂, v₃, v₄), 3, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(5), Val(3)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₄ = ψ₃ + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = ψ(i, j - 3 - $add, k, grid, u, v) + + u₄ = u₃ + u₃ = u₂ + u₂ = u₁ + u₁ = u₀ + u₀ = ℑyᵃᶠᵃ(i, j - 3 - $add, k, grid, u) + + v₄ = v₃ + v₃ = v₂ + v₂ = v₁ + v₁ = v₀ + v₀ = ℑxᶠᵃᵃ(i, j - 3 - $add, k, grid, v) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃, ψ₄), (u₀, u₁, u₂, u₃, u₄), (v₀, v₁, v₂, v₃, v₄), 4, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(5), Val(4)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₄ = ψ₃ + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = ψ(i, j - 4 - $add, k, grid, u, v) + + u₄ = u₃ + u₃ = u₂ + u₂ = u₁ + u₁ = u₀ + u₀ = ℑyᵃᶠᵃ(i, j - 4 - $add, k, grid, u) + + v₄ = v₃ + v₃ = v₂ + v₂ = v₁ + v₁ = v₀ + v₀ = ℑxᶠᵃᵃ(i, j - 4 - $add, k, grid, v) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃, ψ₄), (u₀, u₁, u₂, u₃, u₄), (v₀, v₁, v₂, v₃, v₄), 5, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(5), Val(5)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + τ = τ^2 + + return (ψ̂₁ * τ + ψ̂₂) / (w₁ * τ + 1) + end + + @inline function $biased_interpolate(i, j, k, grid, + scheme::WENO{6, FT, XT, YT, ZT}, + ψ, idx, loc, ::VelocityStencil, u, v, args...) where {FT, XT, YT, ZT} + + # All stencils + ψ₀ = ψ(i, j - $add, k, grid, u, v) + ψ₁ = ψ(i, j + 1 - $add, k, grid, u, v) + ψ₂ = ψ(i, j + 2 - $add, k, grid, u, v) + ψ₃ = ψ(i, j + 3 - $add, k, grid, u, v) + ψ₄ = ψ(i, j + 4 - $add, k, grid, u, v) + ψ₅ = ψ(i, j + 5 - $add, k, grid, u, v) + + u₀ = ℑyᵃᶠᵃ(i, j - $add, k, grid, u) + u₁ = ℑyᵃᶠᵃ(i, j + 1 - $add, k, grid, u) + u₂ = ℑyᵃᶠᵃ(i, j + 2 - $add, k, grid, u) + u₃ = ℑyᵃᶠᵃ(i, j + 3 - $add, k, grid, u) + u₄ = ℑyᵃᶠᵃ(i, j + 4 - $add, k, grid, u) + u₅ = ℑyᵃᶠᵃ(i, j + 5 - $add, k, grid, u) + + v₀ = ℑxᶠᵃᵃ(i, j - $add, k, grid, v) + v₁ = ℑxᶠᵃᵃ(i, j + 1 - $add, k, grid, v) + v₂ = ℑxᶠᵃᵃ(i, j + 2 - $add, k, grid, v) + v₃ = ℑxᶠᵃᵃ(i, j + 3 - $add, k, grid, v) + v₄ = ℑxᶠᵃᵃ(i, j + 4 - $add, k, grid, v) + v₅ = ℑxᶠᵃᵃ(i, j + 5 - $add, k, grid, v) + + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃, ψ₄, ψ₅), (u₀, u₁, u₂, u₃, u₄, u₅), (v₀, v₁, v₂, v₃, v₄, v₅), 1, scheme, $val, idx, loc) + τ = β + ψ̂₁ = ψ̅ * α + ψ̂₂ = ψ̅ * C + w₁ = α + + ψ₅ = ψ₄ + ψ₄ = ψ₃ + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = ψ(i, j - 1 - $add, k, grid, u, v) + + u₅ = u₄ + u₄ = u₃ + u₃ = u₂ + u₂ = u₁ + u₁ = u₀ + u₀ = ℑyᵃᶠᵃ(i, j - 1 - $add, k, grid, u) + + v₅ = v₄ + v₄ = v₃ + v₃ = v₂ + v₂ = v₁ + v₁ = v₀ + v₀ = ℑxᶠᵃᵃ(i, j - 1 - $add, k, grid, v) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃, ψ₄, ψ₅), (u₀, u₁, u₂, u₃, u₄, u₅), (v₀, v₁, v₂, v₃, v₄, v₅), 2, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(6), Val(2)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₅ = ψ₄ + ψ₄ = ψ₃ + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = ψ(i, j - 2 - $add, k, grid, u, v) + + u₅ = u₄ + u₄ = u₃ + u₃ = u₂ + u₂ = u₁ + u₁ = u₀ + u₀ = ℑyᵃᶠᵃ(i, j - 2 - $add, k, grid, u) + + v₅ = v₄ + v₄ = v₃ + v₃ = v₂ + v₂ = v₁ + v₁ = v₀ + v₀ = ℑxᶠᵃᵃ(i, j - 2 - $add, k, grid, v) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃, ψ₄, ψ₅), (u₀, u₁, u₂, u₃, u₄, u₅), (v₀, v₁, v₂, v₃, v₄, v₅), 3, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(6), Val(3)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₅ = ψ₄ + ψ₄ = ψ₃ + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = ψ(i, j - 3 - $add, k, grid, u, v) + + u₅ = u₄ + u₄ = u₃ + u₃ = u₂ + u₂ = u₁ + u₁ = u₀ + u₀ = ℑyᵃᶠᵃ(i, j - 3 - $add, k, grid, u) + + v₅ = v₄ + v₄ = v₃ + v₃ = v₂ + v₂ = v₁ + v₁ = v₀ + v₀ = ℑxᶠᵃᵃ(i, j - 3 - $add, k, grid, v) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃, ψ₄, ψ₅), (u₀, u₁, u₂, u₃, u₄, u₅), (v₀, v₁, v₂, v₃, v₄, v₅), 4, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(6), Val(4)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₅ = ψ₄ + ψ₄ = ψ₃ + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = ψ(i, j - 4 - $add, k, grid, u, v) + + u₅ = u₄ + u₄ = u₃ + u₃ = u₂ + u₂ = u₁ + u₁ = u₀ + u₀ = ℑyᵃᶠᵃ(i, j - 4 - $add, k, grid, u) + + v₅ = v₄ + v₄ = v₃ + v₃ = v₂ + v₂ = v₁ + v₁ = v₀ + v₀ = ℑxᶠᵃᵃ(i, j - 4 - $add, k, grid, v) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃, ψ₄, ψ₅), (u₀, u₁, u₂, u₃, u₄, u₅), (v₀, v₁, v₂, v₃, v₄, v₅), 5, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(6), Val(5)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + ψ₅ = ψ₄ + ψ₄ = ψ₃ + ψ₃ = ψ₂ + ψ₂ = ψ₁ + ψ₁ = ψ₀ + ψ₀ = ψ(i, j - 5 - $add, k, grid, u, v) + + u₅ = u₄ + u₄ = u₃ + u₃ = u₂ + u₂ = u₁ + u₁ = u₀ + u₀ = ℑyᵃᶠᵃ(i, j - 5 - $add, k, grid, u) + + v₅ = v₄ + v₄ = v₃ + v₃ = v₂ + v₂ = v₁ + v₁ = v₀ + v₀ = ℑxᶠᵃᵃ(i, j - 5 - $add, k, grid, v) + + # Stencil S₁ + β, ψ̅, C, α = $weno_interpolant((ψ₀, ψ₁, ψ₂, ψ₃, ψ₄, ψ₅), (u₀, u₁, u₂, u₃, u₄, u₅), (v₀, v₁, v₂, v₃, v₄, v₅), 6, scheme, $val, idx, loc) + τ += add_to_global_smoothness(β, Val(6), Val(6)) + ψ̂₁ += ψ̅ * α + ψ̂₂ += ψ̅ * C + w₁ += α + + τ = τ^2 + + return (ψ̂₁ * τ + ψ̂₂) / (w₁ * τ + 1) + end + end +end \ No newline at end of file diff --git a/src/Architectures.jl b/src/Architectures.jl index 9b03fbf7fa..01b5e706a7 100644 --- a/src/Architectures.jl +++ b/src/Architectures.jl @@ -4,7 +4,6 @@ export AbstractArchitecture, AbstractSerialArchitecture export CPU, GPU export device, architecture, unified_array, device_copy_to! export array_type, on_architecture, arch_array - using CUDA using KernelAbstractions using Adapt @@ -44,7 +43,7 @@ struct GPU <: AbstractSerialArchitecture end ##### device(::CPU) = KernelAbstractions.CPU() -device(::GPU) = CUDA.CUDABackend(; always_inline=true) +device(::GPU) = CUDABackend( ; always_inline=true) architecture() = nothing architecture(::Number) = nothing diff --git a/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl b/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl index 0c9f6b17b5..aaaf4edbcf 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl @@ -13,6 +13,9 @@ using Oceananigans.TurbulenceClosures.CATKEVerticalDiffusivities: FlavorOfCATKE using Oceananigans.ImmersedBoundaries: active_interior_map, ActiveCellsIBG, InteriorMap, active_linear_index_to_tuple +using KernelAbstractions: @private, @uniform, @groupsize, @index, @localmem + + """ compute_tendencies!(model::HydrostaticFreeSurfaceModel, callbacks) diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl index 6de795e2b6..b0ef4e4095 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl @@ -9,7 +9,6 @@ using Oceananigans.Advection: div_Uc, U_dot_∇u, U_dot_∇v using Oceananigans.Forcings: with_advective_forcing using Oceananigans.TurbulenceClosures: shear_production, buoyancy_flux, dissipation using Oceananigans.Utils: SumOfArrays -using KernelAbstractions: @private import Oceananigans.TurbulenceClosures: hydrostatic_turbulent_kinetic_energy_tendency @@ -26,7 +25,7 @@ The tendency for ``u`` is called ``G_u`` and defined via where `p_n` is the part of the barotropic kinematic pressure that's treated implicitly during time-stepping. """ -@inline function hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, +@inline function hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, advection, coriolis, closure, @@ -62,10 +61,10 @@ The tendency for ``v`` is called ``G_v`` and defined via ∂_t v = G_v - ∂_y p_n ``` -where `p_n` is the part of the barotropic kinematic pressure that's treated +where `p_n` is the part of the barotropic kinematic pressure that's treated implicitly during time-stepping. """ -@inline function hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, +@inline function hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, advection, coriolis, closure, @@ -103,7 +102,7 @@ The tendency is called ``G_c`` and defined via where `c = C[tracer_index]`. """ -@inline function hydrostatic_free_surface_tracer_tendency(i, j, k, grid, +@inline function hydrostatic_free_surface_tracer_tendency(i, j, k, grid, val_tracer_index::Val{tracer_index}, val_tracer_name, advection, diff --git a/src/Utils/kernel_launching.jl b/src/Utils/kernel_launching.jl index b0d3338669..3f4b5906c7 100644 --- a/src/Utils/kernel_launching.jl +++ b/src/Utils/kernel_launching.jl @@ -158,7 +158,6 @@ If the worksize is 0 after adjusting for active cells, the function returns `not - `active_cells_map`: A map indicating the active cells in the grid. If the map is not a nothing, the workspec will be disregarded and the kernel is configured as a linear kernel with a worksize equal to the length of the active cell map. Default is `nothing`. """ - function configured_kernel(arch, grid, workspec, kernel!; include_right_boundaries = false, reduced_dimensions = (), @@ -309,3 +308,19 @@ function partition(kernel::OffsetKernel, inrange, ingroupsize) return iterspace, dynamic end + +using KernelAbstractions: Kernel, backend, mkcontext, workgroupsize, launch_config, workgroupsize +using CUDA: launch_configuration, registers, @cuda, CUDABackend + +function retrieve_kernel(obj::Kernel{CUDABackend}, args...; ndrange=nothing, workgroupsize=nothing) + back = backend(obj) + + ndrange, workgroupsize, iterspace, dynamic = launch_config(obj, ndrange, workgroupsize) + # this might not be the final context, since we may tune the workgroupsize + ctx = mkcontext(obj, ndrange, iterspace) + maxthreads = nothing + + kernel = @cuda launch=false always_inline=back.always_inline maxthreads=maxthreads obj.f(ctx, args...) + + return kernel +end \ No newline at end of file diff --git a/test_gpu.jl b/test_gpu.jl new file mode 100644 index 0000000000..e561ad4071 --- /dev/null +++ b/test_gpu.jl @@ -0,0 +1,95 @@ +using Oceananigans +using Oceananigans.Advection: DefaultStencil, OnlySelfUpwinding +using Oceananigans.Architectures: architecture +using Oceananigans.Utils: configured_kernel +using Oceananigans.Models.HydrostaticFreeSurfaceModels: + immersed_boundary_condition, + top_tracer_boundary_conditions, + compute_hydrostatic_free_surface_Gu!, + compute_hydrostatic_free_surface_Gv!, + compute_hydrostatic_free_surface_Gc! + +using Oceananigans.ImmersedBoundaries: active_interior_map + +using CUDA +CUDA.device!(1) +a = CPU() + +grid = LatitudeLongitudeGrid(a, size = (20, 20, 20), latitude = (-1, 1), longitude = (-1, 1), z = (0, 1), halo = (7, 7, 7)) +# grid = ImmersedBoundaryGrid(grid, GridFittedBottom((x, y)->x); active_cells_map = true) + +momentum_advection = VectorInvariant(vorticity_scheme = WENO(; order = 9), + divergence_scheme = WENO(), + vertical_scheme = Centered()) #, vorticity_stencil = DefaultStencil()) + +tracer_advection = Oceananigans.Advection.TracerAdvection(WENO(; order = 7), WENO(; order = 7), Centered()) + +# momentum_advection = nothing + +model = HydrostaticFreeSurfaceModel(; grid, momentum_advection) #, tracer_advection) + +arch = architecture(grid) +velocities = model.velocities + +u_immersed_bc = immersed_boundary_condition(velocities.u) +v_immersed_bc = immersed_boundary_condition(velocities.v) + +start_momentum_kernel_args = (model.advection.momentum, + model.coriolis, + model.closure) + +end_momentum_kernel_args = (velocities, + model.free_surface, + model.tracers, + model.buoyancy, + model.diffusivity_fields, + model.pressure.pHY′, + model.auxiliary_fields, + model.forcing, + model.clock) + +u_kernel_args = tuple(start_momentum_kernel_args..., u_immersed_bc, end_momentum_kernel_args...) +v_kernel_args = tuple(start_momentum_kernel_args..., v_immersed_bc, end_momentum_kernel_args...) + +c_immersed_bc = immersed_boundary_condition(model.tracers.T) +top_tracer_bcs = top_tracer_boundary_conditions(grid, model.tracers) + +T_kernel_args = tuple(Val(1), + Val(:T), + model.advection.T, + nothing, + c_immersed_bc, + model.buoyancy, + model.biogeochemistry, + model.velocities, + model.free_surface, + model.tracers, + top_tracer_bcs, + model.diffusivity_fields, + model.auxiliary_fields, + model.forcing.T, + model.clock) + +active_cells_map = active_interior_map(model.grid) + +loop_u! = configured_kernel(arch, grid, :xyz, + compute_hydrostatic_free_surface_Gu!; + active_cells_map) + +loop_v! = configured_kernel(arch, grid, :xyz, + compute_hydrostatic_free_surface_Gv!; + active_cells_map) + +loop_T! = configured_kernel(arch, grid, :xyz, + compute_hydrostatic_free_surface_Gc!; + active_cells_map) + +@show "Kernels are configured" + +kernel_u = Oceananigans.Utils.retrieve_kernel(loop_u!, model.timestepper.Gⁿ.u, grid, active_cells_map, u_kernel_args) +kernel_v = Oceananigans.Utils.retrieve_kernel(loop_v!, model.timestepper.Gⁿ.v, grid, active_cells_map, v_kernel_args) +kernel_T = Oceananigans.Utils.retrieve_kernel(loop_T!, model.timestepper.Gⁿ.T, grid, active_cells_map, T_kernel_args) + +@show CUDA.registers(kernel_u) CUDA.registers(kernel_v) CUDA.registers(kernel_T) + +nothing diff --git a/validation/advection/validate_one_dimensional_advection.jl b/validation/advection/validate_one_dimensional_advection.jl index f04c3c3661..6fc48f52d5 100644 --- a/validation/advection/validate_one_dimensional_advection.jl +++ b/validation/advection/validate_one_dimensional_advection.jl @@ -1,5 +1,5 @@ using Oceananigans -using Oceananigans.Advection: AbstractCenteredAdvectionScheme, AbstractUpwindBiasedAdvectionScheme, VelocityStencil, VorticityStencil +using Oceananigans.Advection: AbstractCenteredAdvectionScheme, AbstractUpwindBiasedAdvectionScheme, VelocityStencil using Oceananigans.Models.ShallowWaterModels: VectorInvariantFormulation, ConservativeFormulation using JLD2 using OffsetArrays @@ -58,7 +58,7 @@ Z = -0.7 a = 0.5 α = 10 -@inline function c₀_1D(x, y, z) +@inline function c₀_1D(x, y) if x <= -0.6 && x >= -0.8 return 1/6*(G(x, β, Z-δ) + 4*G(x, β, Z) + G(x, β, Z+δ)) elseif x <= -0.2 && x >= -0.4 @@ -78,8 +78,8 @@ Schemes = [:Centered, :UpwindBiased, :WENO] @inline grid_or_not(::Nothing) = -1 # # Checking the accuracy of different schemes with different settings -buffers = [2, 3] -for (gr, grid) in enumerate([grid_str]) +buffers = [2, 3, 4, 5, 6] +for (gr, grid) in enumerate([grid_reg]) @info "testing grid number $gr" @@ -92,7 +92,7 @@ for (gr, grid) in enumerate([grid_str]) c_real = CenterField(grid) formulation = ConservativeFormulation() - for Scheme in [Schemes[2]] + for Scheme in [Schemes[3]] for buffer in buffers, gr in (nothing, grid) scheme = eval(Scheme)(gr, order = advection_order(buffer, eval(Scheme))) @@ -124,7 +124,7 @@ for (gr, grid) in enumerate([grid_str]) solution[(buffer, Int(i), grid_or_not(gr))] = csim end - c_sol(x, y, z) = @. c₀_1D(x - model.clock.time + 2, y, z) + c_sol(x, y) = @. c₀_1D(x - model.clock.time + 2, y) set!(c_real, c_sol) end