feat(WENO): discretizer integration for non-uniform grids

src/MethodOfLines.jl

@@ -99,8 +99,8 @@ include("discretization/schemes/upwind_difference/upwind_difference.jl")
 include("discretization/schemes/half_offset_centred_difference.jl")
 include("discretization/schemes/nonlinear_laplacian/nonlinear_laplacian.jl")
 include("discretization/schemes/spherical_laplacian/spherical_laplacian.jl")
-include("discretization/schemes/WENO/WENO.jl")
 include("discretization/schemes/WENO/nonuniform_weno.jl")
+include("discretization/schemes/WENO/WENO.jl")
 include("discretization/schemes/integral_expansion/integral_expansion.jl")
 
 # System Discretization

src/discretization/schemes/WENO/WENO.jl

@@ -73,13 +73,23 @@ end
 end
 
 """
-`WENOScheme` of Jiang and Shu
+    WENOScheme(; epsilon = 1.0e-6)
+
+Jiang-Shu WENO-5 advection scheme for uniform and non-uniform grids.
+
 ## Keyword Arguments
 - `epsilon`: A quantity used to prevent vanishing denominators in the scheme, defaults to `1e-6`. More sensitive problems will benefit from a smaller value. It is defined as a functional scheme.
 """
 function WENOScheme(; epsilon = 1.0e-6)
-    boundary_f = [nothing, nothing]
-    return FunctionalScheme{5, 0}(
-        weno_f, boundary_f, boundary_f, true, [epsilon], name = "WENO"
+    lower = WENONonUniformBoundary[WENONonUniformBoundary{1}(), WENONonUniformBoundary{2}()]
+    upper = WENONonUniformBoundary[WENONonUniformBoundary{5}(), WENONonUniformBoundary{4}()]
+    return FunctionalScheme{5, 5}(
+        weno_f, lower, upper, true, [epsilon], name = "WENO"
     )
 end
+
+# extent dispatch keys on typeof(lower): must stay Vector{<:WENONonUniformBoundary}.
+# Replacing lower/upper with nothing restores 2-arg extent=0 and breaks uniform routing.
+const _WENOBoundaryVec = AbstractVector{<:WENONonUniformBoundary}
+extent(::FunctionalScheme{<:Any, <:_WENOBoundaryVec}, dorder, dx::Number) = 2
+extent(::FunctionalScheme{<:Any, <:_WENOBoundaryVec}, dorder, dx::AbstractVector) = 0

src/discretization/schemes/WENO/nonuniform_weno.jl

@@ -183,16 +183,30 @@ Base.@propagate_inbounds @inline function weno_f_nonuniform(u, p, t, x, dx::Abst
     return _weno_f_nonuniform_core(u, p[1], x, Val(3))
 end
 
-# Scalar-dx method required by the FunctionalScheme{5,0} contract.
+# Scalar-dx method required by the FunctionalScheme{5,5} contract.
 Base.@propagate_inbounds @inline function weno_f_nonuniform(u, p, t, x, dx::Number)
     return _weno_f_nonuniform_core(u, p[1], x, Val(3))
 end
 
-# 6-arg: explicit Val{Target}; dx unused (FunctionalScheme{5,0} contract).
+# 6-arg: explicit Val{Target}; dx unused (FunctionalScheme{5,5} contract).
 Base.@propagate_inbounds @inline function weno_f_nonuniform(u, p, t, x, dx::AbstractVector, ::Val{T}) where {T}
     return _weno_f_nonuniform_core(u, p[1], x, Val(T))
 end
 
 Base.@propagate_inbounds @inline function weno_f_nonuniform(u, p, t, x, dx::Number, ::Val{T}) where {T}
     return _weno_f_nonuniform_core(u, p[1], x, Val(T))
 end
+
+struct WENONonUniformBoundary{T} <: Function end
+
+Base.@propagate_inbounds @inline function (::WENONonUniformBoundary{T})(
+        u, p, t, x, dx::AbstractVector
+    ) where {T}
+    return weno_f_nonuniform(u, p, t, x, dx, Val(T))
+end
+
+Base.@propagate_inbounds @inline function (::WENONonUniformBoundary{T})(
+        u, p, t, x, dx::Number
+    ) where {T}
+    return weno_f_nonuniform(u, p, t, x, dx, Val(T))
+end

src/interface/scheme_types.jl

@@ -9,6 +9,8 @@ end
 
 extent(scheme::UpwindScheme, dorder) = 0 # dorder + scheme.order - 1
 
+extent(scheme::AbstractScheme, dorder, dx) = extent(scheme, dorder)
+
 # Functional Schemes
 
 """

src/system_parsing/interior_map.jl

@@ -118,10 +118,10 @@ function calculate_stencil_extents(s, u, discretization, orders, bcmap)
         for dorder in filter(isodd, orders[x])
             ascheme = dorder == 1 ? advection_scheme : UpwindScheme()
             if !haslower
-                lowerextents[j] = max(lowerextents[j], extent(ascheme, dorder))
+                lowerextents[j] = max(lowerextents[j], extent(ascheme, dorder, s.dxs[x]))
             end
             if !hasupper
-                upperextents[j] = max(upperextents[j], extent(ascheme, dorder))
+                upperextents[j] = max(upperextents[j], extent(ascheme, dorder, s.dxs[x]))
             end
         end
     end

test/Components/weno_boundary_integration.jl

@@ -0,0 +1,137 @@
+# WENO discretizer integration: extent dispatch, interior map, boundary callable wiring.
+
+using Test
+using ModelingToolkit, MethodOfLines, DomainSets, OrdinaryDiffEq, SciMLBase
+
+const M = MethodOfLines
+
+@parameters t x
+@variables u(..)
+Dt = Differential(t)
+Dx = Differential(x)
+
+function build_discrete_system(pdesys, disc)
+    v = M.VariableMap(pdesys, disc)
+    bcorders = Dict(
+        map(xx -> xx => M.d_orders(xx, M.get_bcs(pdesys)), M.PDEBase.all_ivs(v))
+    )
+    bmap = M.PDEBase.parse_bcs(M.get_bcs(pdesys), v, bcorders)
+    s = M.construct_discrete_space(v, disc)
+    eqs = M.get_eqs(pdesys)
+    eqs = eqs isa AbstractVector ? Vector{Equation}(eqs) : Equation[eqs]
+    im = M.PDEBase.construct_var_equation_mapping(eqs, bmap, s, disc)
+    return im, s, bmap
+end
+
+function perturbed_nu_grid()
+    g = collect(range(0.0, 1.0, length = 21))
+    g[2:(end - 1)] .+= 0.004 .* sin.(1:19)
+    return g
+end
+
+@testset "WENONonUniformBoundary scalar dx" begin
+    xs = collect(range(0.0, 1.0, length = 5))
+    u = sin.(2π .* xs)
+    result = @inferred M.WENONonUniformBoundary{1}()(u, (1.0e-6,), 0.0, xs, 0.1)
+    @test result isa Float64
+    @test isfinite(result)
+end
+
+@testset "Grid-dispatched extent" begin
+    s = WENOScheme()
+    @test M.extent(s, 1, 1.0) == 2
+    @test M.extent(s, 1, [0.1, 0.2, 0.3]) == 0
+    @test M.extent(WENOScheme(), 1) == 0
+    @test M.extent(WENOScheme(), 1, 1.0) == 2
+
+    m_nonuniform = which(M.extent, (typeof(s), Int, Vector{Float64}))
+    m_fallback = which(M.extent, (M.AbstractScheme, Int, Any))
+    @test m_nonuniform !== m_fallback
+
+    @test M.extent(UpwindScheme(), 1, [0.1, 0.2]) == M.extent(UpwindScheme(), 1)
+end
+
+@testset "Interior map extent dispatch (non-uniform vs uniform)" begin
+    eq = Dt(u(t, x)) ~ -Dx(u(t, x))
+    bcs = [u(0, x) ~ sinpi(x), u(t, 0.0) ~ 0.0, u(t, 1.0) ~ 0.0]
+    domains = [t ∈ Interval(0.0, 1.0), x ∈ Interval(0.0, 1.0)]
+    @named pdesys = PDESystem(eq, bcs, domains, [t, x], [u(t, x)])
+
+    nonuniform_grid = perturbed_nu_grid()
+
+    im_nu, _, _ = build_discrete_system(
+        pdesys, MOLFiniteDifference([x => nonuniform_grid], t; advection_scheme = WENOScheme())
+    )
+    im_u, _, _ = build_discrete_system(
+        pdesys, MOLFiniteDifference([x => 1 / 20], t; advection_scheme = WENOScheme())
+    )
+
+    pde_nu = first(keys(im_nu.I))
+    pde_u = first(keys(im_u.I))
+
+    @test im_nu.stencil_extents[pde_nu] == ([0], [0])
+    @test im_u.stencil_extents[pde_u] == ([2], [2])
+
+    interior_nu = vec(collect(im_nu.I[pde_nu]))
+    @test !(CartesianIndex(1) in interior_nu)
+    @test CartesianIndex(2) in interior_nu
+    @test CartesianIndex(20) in interior_nu
+    @test !(CartesianIndex(21) in interior_nu)
+
+    interior_u = vec(collect(im_u.I[pde_u]))
+    @test !(CartesianIndex(1) in interior_u)
+    @test !(CartesianIndex(2) in interior_u)
+    @test !(CartesianIndex(20) in interior_u)
+    @test !(CartesianIndex(21) in interior_u)
+end
+
+@testset "Discretizer boundary callable routing and execution" begin
+    T_END = 0.1
+    u_exact(x, t) = sin(2π * (x - t))
+
+    nonuniform_grid = perturbed_nu_grid()
+    x0, xL = nonuniform_grid[1], nonuniform_grid[end]
+
+    eq = Dt(u(t, x)) ~ -Dx(u(t, x))
+    bcs = [
+        u(0.0, x) ~ sin(2π * x),
+        u(t, x0) ~ u_exact(x0, t),
+        Dx(u(t, xL)) ~ 0.0,
+    ]
+    domains = [t ∈ Interval(0.0, T_END), x ∈ Interval(x0, xL)]
+    @named pdesys = PDESystem(eq, bcs, domains, [t, x], [u(t, x)])
+
+    disc = MOLFiniteDifference([x => nonuniform_grid], t; advection_scheme = WENOScheme())
+    im, s, bmap = build_discrete_system(pdesys, disc)
+
+    pde = first(keys(im.I))
+    F = WENOScheme()
+    u_field = im.var[pde]
+    j = M.x2i(s, u_field, x)
+    bs = M.PDEBase.filter_interfaces(bmap[only(keys(bmap))][x])
+
+    f_lo, Itap_lo = M.get_f_and_taps(F, CartesianIndex(2), s, bs, (j, x), u_field)
+    f_hi, Itap_hi = M.get_f_and_taps(F, CartesianIndex(20), s, bs, (j, x), u_field)
+    @test f_lo === M.WENONonUniformBoundary{2}()
+    @test f_hi === M.WENONonUniformBoundary{4}()
+
+    xs = collect(s.grid[x])
+    uvals = sin.(2π .* xs)
+    dx_vec = collect(s.dxs[x])
+    ps = (F.ps[1],)
+    itap_lo = map(I -> I[j], Itap_lo)
+    itap_hi = map(I -> I[j], Itap_hi)
+    val_lo = @inferred f_lo(
+        uvals[itap_lo], ps, 0.0, @view(xs[itap_lo]), @view(dx_vec[itap_lo[1:(end - 1)]])
+    )
+    val_hi = @inferred f_hi(
+        uvals[itap_hi], ps, 0.0, @view(xs[itap_hi]), @view(dx_vec[itap_hi[1:(end - 1)]])
+    )
+    @test val_lo isa Float64 && isfinite(val_lo)
+    @test val_hi isa Float64 && isfinite(val_hi)
+
+    prob = discretize(pdesys, disc)
+    sol = solve(prob, Tsit5(); abstol = 1.0e-8, reltol = 1.0e-8, saveat = [T_END])
+    @test SciMLBase.successful_retcode(sol)
+    @test all(isfinite, sol[u(t, x)][end, :])
+end

test/runtests.jl

@@ -51,6 +51,9 @@ run_tests(;
             @safetestset "WENO Non-Uniform Boundary" begin
                 include(joinpath(@__DIR__, "Components", "weno_nonuniform_boundary.jl"))
             end
+            @safetestset "WENO Boundary Integration" begin
+                include(joinpath(@__DIR__, "Components", "weno_boundary_integration.jl"))
+            end
             @safetestset "ODEFunction" begin
                 include(joinpath(@__DIR__, "Components", "ODEFunction_test.jl"))
             end