Add ArrayDiscretization strategy for vectorized interior equations

Project.toml

@@ -31,7 +31,7 @@ Interpolations = "0.14, 0.15, 0.16"
 Latexify = "0.15, 0.16"
 ModelingToolkit = "11"
 OrdinaryDiffEq = "6"
-PDEBase = "0.1.20"
+PDEBase = "0.1.22"
 PrecompileTools = "1"
 RuntimeGeneratedFunctions = "0.5.12"
 SafeTestsets = "0.0.1, 0.1"

docs/src/advection_schemes.md

@@ -16,10 +16,19 @@ Changes the direction of the stencil based on the sign of the coefficient of the
 WENOScheme(epsilon = 1e-6)
 ```
 
-A more stable scheme, 5th order accurate, which is a weighted sum of several different schemes, weighted based on the curvature of the solution at the point in question. More stable and tolerant of discontinuities, at the cost of solve complexity.
+A more stable scheme, 5th order accurate in the interior, which is a weighted sum of several different schemes, weighted based on the curvature of the solution at the point in question. More stable and tolerant of discontinuities, at the cost of solve complexity.
 
 `epsilon`is a quantity used to prevent vanishing denominators in the scheme, defaults to `1e-6`. Problems with a lower magnitude solution will benefit from a smaller value.
 
+!!! note "Boundary behaviour"
+    The two grid points closest to each non-periodic boundary use a
+    3rd-order one-sided finite-difference stencil (Fornberg weights) in
+    place of the full 5-point WENO stencil.  Interior accuracy is still
+    5th order; overall accuracy on bounded domains is dominated by the
+    3rd-order boundary stencils.  Problems that need 5th-order accuracy
+    everywhere should be posed with periodic boundaries so only the
+    interior stencil is exercised.
+
 Problems which require this scheme may also benefit from a [Strong-Stability-Preserving (SSP) solver](https://docs.sciml.ai/DiffEqDocs/stable/solvers/ode_solve/#Explicit-Strong-Stability-Preserving-Runge-Kutta-Methods-for-Hyperbolic-PDEs-(Conservation-Laws)).
 
 Problems with first order derivatives which multiply one another will need to use this scheme over the upwind scheme.

src/MOL_discretization.jl

@@ -11,8 +11,8 @@ function PDEBase.interface_errors(
     if !any(s -> discretization.advection_scheme isa s, [UpwindScheme, FunctionalScheme])
         throw(ArgumentError("Only `UpwindScheme()` and `FunctionalScheme()` are supported advection schemes. Got $(typeof(discretization.advection_scheme))."))
     end
-    return if !(typeof(discretization.disc_strategy) ∈ [ScalarizedDiscretization])
-        throw(ArgumentError("Only `ScalarizedDiscretization()` are supported discretization strategies."))
+    return if !(typeof(discretization.disc_strategy) ∈ [ScalarizedDiscretization, ArrayDiscretization])
+        throw(ArgumentError("Only `ScalarizedDiscretization()` and `ArrayDiscretization()` are supported discretization strategies."))
     end
 end
 

src/MethodOfLines.jl

@@ -106,12 +106,14 @@ include("discretization/schemes/integral_expansion/integral_expansion.jl")
 
 # System Discretization
 include("discretization/generate_finite_difference_rules.jl")
+include("discretization/generate_array_fd_rules.jl")
 include("discretization/generate_bc_eqs.jl")
 include("discretization/generate_ic_defaults.jl")
 include("discretization/staggered_discretize.jl")
 
 # Main
 include("scalar_discretization.jl")
+include("array_discretization.jl")
 include("MOL_discretization.jl")
 
 ## PrecompileTools
@@ -121,5 +123,6 @@ include("precompile.jl")
 export MOLFiniteDifference, discretize, symbolic_discretize, ODEFunctionExpr, generate_code,
     grid_align, edge_align, center_align, get_discrete, chebyspace
 export UpwindScheme, WENOScheme, FunctionalScheme, MOLDiscCallback
+export ScalarizedDiscretization, ArrayDiscretization
 
 end

src/array_discretization.jl

@@ -0,0 +1,101 @@
+"""
+    _to_scalar_interiormap(interiormap)
+
+Convert an `InteriorMap` whose `.I` dict stores `[(lo, hi), …]` tuples
+(as produced by `generate_interior` for `ArrayDiscretization`) into one
+that stores `CartesianIndices` so that the existing boundary-equation code
+can consume it unchanged.
+
+The output dict is typed `Dict{Any, CartesianIndices}` so subsequent
+iteration in `generate_bc_eqs!` / `generate_extrap_eqs!` /
+`generate_corner_eqs!` sees concrete value types instead of `Any`.
+"""
+function _to_scalar_interiormap(interiormap)
+    # Use an `Any` key type because the PDE-equation key is typed
+    # heterogeneously upstream in `InteriorMap`; narrowing the value type
+    # is enough to lift the value-side type instability JET flagged.
+    scalar_I = Dict{Any, CartesianIndices}()
+    for (pde, ranges) in interiormap.I
+        cart_ranges = if ranges isa AbstractVector && !isempty(ranges) && first(ranges) isa Tuple
+            # ArrayDiscretization path: `ranges` is `Vector{Tuple{Int,Int}}`.
+            Tuple(r[1]:r[2] for r in ranges)
+        elseif ranges isa CartesianIndices
+            ranges.indices
+        else
+            # Defensive: any other shape we don't handle gets re-wrapped
+            # so downstream `CartesianIndices` iteration keeps working.
+            (ranges,)
+        end
+        scalar_I[pde] = CartesianIndices(cart_ranges)
+    end
+    return InteriorMap(
+        interiormap.var, interiormap.pde, scalar_I,
+        interiormap.lower, interiormap.upper, interiormap.stencil_extents
+    )
+end
+
+"""
+    discretize_equation!(disc_state, pde, interiormap, eqvar, bcmap, depvars,
+                         s, derivweights, indexmap,
+                         discretization::MOLFiniteDifference{G, D <: ArrayDiscretization})
+
+Array-based discretisation of a single PDE.
+
+Boundary equations (BCs, extrapolation, corners) are handled identically to the
+scalar path (using a CartesianIndices-based interior map).  Interior equations
+are generated via a *template-instantiation* strategy: substitution rules are
+built once with a symbolic index variable and the PDE is symbolically
+transformed once, then the resulting template is instantiated at each interior
+grid point.
+"""
+function PDEBase.discretize_equation!(
+        disc_state::PDEBase.EquationState, pde::Equation, interiormap,
+        eqvar, bcmap, depvars, s::DiscreteSpace, derivweights, indexmap,
+        discretization::MOLFiniteDifference{G, D}
+    ) where {G, D <: ArrayDiscretization}
+
+    # Convert tuple-range interior map to CartesianIndices for boundary code
+    scalar_interiormap = _to_scalar_interiormap(interiormap)
+    should_validate = discretization.disc_strategy.validate
+
+    # ── boundary handling (uses scalar-compatible interior map) ───────────────
+    boundaryvalfuncs = generate_boundary_val_funcs(
+        s, depvars, bcmap, indexmap, derivweights
+    )
+    eqvarbcs = mapreduce(x -> bcmap[operation(eqvar)][x], vcat, s.x̄)
+    for boundary in eqvarbcs
+        if boundary isa InterfaceBoundary
+            generate_bc_eqs_arrayop!(disc_state, s, boundaryvalfuncs, scalar_interiormap, boundary)
+        elseif boundary isa AbstractTruncatingBoundary
+            generate_bc_eqs_arrayop!(disc_state, s, boundaryvalfuncs, scalar_interiormap, boundary, derivweights; validate=should_validate)
+        else
+            generate_bc_eqs!(disc_state, s, boundaryvalfuncs, scalar_interiormap, boundary)
+        end
+    end
+    generate_extrap_eqs!(
+        disc_state, pde, eqvar, s, derivweights, scalar_interiormap, bcmap
+    )
+    generate_corner_eqs!(
+        disc_state, s, scalar_interiormap, ndims(s.discvars[eqvar]), eqvar
+    )
+
+    # ── interior equations ───────────────────────────────────────────────────
+    interior_ranges = interiormap.I[pde]   # [(lo, hi), …] for ArrayDiscretization
+
+    eqs = if length(interior_ranges) == 0
+        # No spatial dimensions — fall back to scalar point discretisation
+        II = CartesianIndex()
+        [discretize_equation_at_point(
+            II, s, depvars, pde, derivweights, bcmap,
+            eqvar, indexmap, boundaryvalfuncs
+        )]
+    else
+        generate_array_interior_eqs(
+            s, depvars, pde, derivweights, bcmap, eqvar,
+            indexmap, boundaryvalfuncs, interior_ranges;
+            validate=should_validate
+        )
+    end
+
+    return vcat!(disc_state.eqs, eqs)
+end