System Transformation - Automatically make PDESystems compatible

.github/workflows/CI.yml

@@ -26,6 +26,7 @@ jobs:
           - DAE
           - Burgers
           - Brusselator
+          - Mixed_Derivatives
         version:
           - '1'
           - '1.6'

docs/src/MOLFiniteDifference.md

@@ -4,20 +4,11 @@ struct MOLFiniteDifference{G} <: DiffEqBase.AbstractDiscretization
     dxs
     time
     approx_order::Int
-    upwind_order::Int
+    advection_scheme
     grid_align::G
-end
-
-# Constructors. If no order is specified, both upwind and centered differences will be 2nd order
-function MOLFiniteDifference(dxs, time=nothing; approx_order = 2, upwind_order = 1, grid_align=CenterAlignedGrid())
-
-    if approx_order % 2 != 0
-        @warn "Discretization approx_order must be even, rounding up to $(approx_order+1)"
-    end
-    @assert approx_order >= 1 "approx_order must be at least 1"
-    @assert upwind_order >= 1 "upwind_order must be at least 1"
-
-    return MOLFiniteDifference{typeof(grid_align)}(dxs, time, approx_order, upwind_order, grid_align)
+    should_transform::Bool
+    use_ODAE::Bool
+    kwargs
 end
 ```
 
@@ -33,20 +24,27 @@ discretization = MOLFiniteDifference(dxs,
                                       <your choice of continuous variable, usually time>; 
                                       advection_scheme = <UpwindScheme() or WENOScheme()>, 
                                       approx_order = <Order of derivative approximation, starting from 2> 
-                                      grid_align = <your grid type choice>)
+                                      grid_align = <your grid type choice>,
+                                      should_transform = <Whether to automatically transform the PDESystem (see below)>
+                                      use_ODAE = <Whether to use ODAEProblem>)
 prob = discretize(pdesys, discretization)
 ```
 Where `dxs` is a vector of pairs of parameters to the grid step in this dimension, i.e. `[x=>0.2, y=>0.1]`.
 For a non uniform rectilinear grid, replace any or all of the step sizes with the grid you'd like to use with that variable, must be an `AbstractVector` but not a `StepRangeLen`.
 
 Note that the second argument to `MOLFiniteDifference` is optional, all parameters can be discretized if all required boundary conditions are specified.
 
-Currently implemented advection schemes are `UpwindScheme()` and `WENOScheme()`, defaults to upwind. See [advection schemes](@ref adschemes) for more information.
+Currently implemented options for `advection_scheme` are `UpwindScheme()` and `WENOScheme()`, defaults to upwind. See [advection schemes](@ref adschemes) for more information.
 
-Currently supported grid types: `center_align` and `edge_align`. Edge align will give better accuracy with Neumann boundary conditions.
+Currently supported `grid_align`: `center_align` and `edge_align`. Edge align will give better accuracy with Neumann boundary conditions. Defaults tp `center_align`.
 
 `center_align`: naive grid, starting from lower boundary, ending on upper boundary with step of `dx`
 
 `edge_align`: offset grid, set halfway between the points that would be generated with center_align, with extra points at either end that are above and below the supremum and infimum by `dx/2`. This improves accuracy for Neumann BCs.
 
-Any unrecognized keyword arguments will be passed to the `ODEProblem` constructor, see [its documentation](https://docs.sciml.ai/ModelingToolkit/stable/systems/ODESystem/#Standard-Problem-Constructors) for available options.
+`should_transform`: Whether to automatically transform the system to make it compatible with MethodOfLines where possible, defaults to true. If your system has no mixed derivatives, all derivatives are purely of a dependent variable i.e. `Dx(u_aux(t,x))` not `Dx(v(t,x)*u(t,x))`, excepting nonlinear and spherical laplacians for which this holds for the innermost derivative argument, and no expandable derivatives, this can be set to false for better discretization performance at the cost of generality, if you perform these transformations yourself.
+
+`use_ODAE`: MethodOfLines will automatically make use of `ODAEProblem` where relevant, which improves performance for DAEs (as discretized PDEs are in general), if this is set to true. Defaults to false.
+
+Any unrecognized keyword arguments will be passed to the `ODEProblem` constructor, see [its documentation](https://docs.sciml.ai/ModelingToolkit/stable/systems/ODESystem/#Standard-Problem-Constructors) for available options.
+

docs/src/index.md

@@ -62,13 +62,10 @@ packages.
 At the moment the package is able to discretize almost any system, with some assumptions listed below
 
 - That the grid is cartesian.
-- That the equation is first order in time.
 - Boundary conditions in time are supplied as initial conditions, not at the end of the simulation interal. If your system requires a final condition, please use a change of variables to rectify this. This is unlikely to change due to upstream constraints.
 - Intergral equations are not supported.
 - That dependant variables always have the same argument signature, except in BCs.
 - That periodic boundary conditions are of the simple form `u(t, x_min) ~ u(t, x_max)`, or the same with lhs and rhs reversed. Note that this generalises to higher dimensions. Please note that if you want to use a periodic condition on a dimension with WENO schemes, please use a periodic condition on all variables in that dimension.
 - That boundary conditions do not contain references to derivatives which are not in the direction of the boundary, except in time.
-- That initial conditions are of the form `u(...) ~ ...`, and don't reference the initial time derivative.
-- That simple derivative terms are purely of a dependant variable, for example `Dx(u(t,x,y))` is allowed but `Dx(u(t,x,y)*v(t,x,y))`, `Dx(u(t,x)+1)` or `Dx(f(u(t,x)))` are not. As a workaround please expand such terms with the product rule and use the linearity of the derivative operator, or define a new auxiliary dependant variable by adding an equation for it like `eqs = [Differential(x)(w(t,x))~ ... , w(t,x) ~ v(t,x)*u(t,x)]`, along with appropriate BCs/ICs. An exception to this is if the differential is a nonlinear or spherical laplacian, in which case only the innermost argument should be wrapped.
-- Note that the above also applies to mixed derivatives, please wrap the inner derivative.
-- That odd order derivatives do not multiply or divide each other. A workaround is to wrap all but one derivative per term in an auxiliary variable, such as `dxu(x, t) ~ Differential(x)(u(x, t))`. The performance hit from auxiliary variables should be negligable due to a structural simplification step.
+- That odd order derivatives do not multiply or divide each other. A workaround is to wrap all but one derivative per term in an auxiliary variable, such as `dxu(x, t) ~ Differential(x)(u(x, t))`. The performance hit from auxiliary variables should be negligable due to a structural simplification step.
+- That the WENO scheme must be used when there are mixed derivatives.

src/MOL_discretization.jl

@@ -2,67 +2,91 @@
 
 function interface_errors(depvars, indvars, discretization)
     for x in indvars
-        @assert findfirst(dxs -> isequal(x, dxs[1].val), discretization.dxs) !== nothing "Variable $x has no step size"
+        @assert haskey(discretization.dxs, Num(x)) || haskey(discretization.dxs, x) "Variable $x has no step size"
+    end
+    if !(typeof(discretization.advection_scheme) ∈ [UpwindScheme, WENOScheme])
+        throw(ArgumentError("Only `UpwindScheme()` and `WENOScheme()` are supported advection schemes."))
     end
 end
 
 function SciMLBase.symbolic_discretize(pdesys::PDESystem, discretization::MethodOfLines.MOLFiniteDifference{G}) where {G}
-    pdeeqs = [eq.lhs - eq.rhs ~ 0 for eq in pdesys.eqs]
-    bcs = pdesys.bcs
-    domain = pdesys.domain
-
     t = discretization.time
+    cardinalize_eqs!(pdesys)
+
+    ############################
+    # System Parsing and Transformation
+    ############################
+    # Parse the variables in to the right form and store useful information about the system
+    v = VariableMap(pdesys, t)
+    # Check for basic interface errors
+    interface_errors(v.ū, v.x̄, discretization)
+    # Extract tspan
+    tspan = t !== nothing ? v.intervals[t] : nothing
+    # Find the derivative orders in the bcs
+    bcorders = Dict(map(x -> x => d_orders(x, pdesys.bcs), all_ivs(v)))
+    # Create a map of each variable to their boundary conditions including initial conditions
+    boundarymap = parse_bcs(pdesys.bcs, v, bcorders)
+    # Generate a map of each variable to whether it is periodic in a given direction
+    pmap = PeriodicMap(boundarymap, v)
+    # Transform system so that it is compatible with the discretization
+    if discretization.should_transform
+        pdesys, pmap = transform_pde_system!(v, boundarymap, pmap, pdesys)
+    end
 
-    depvar_ops = map(x -> operation(x.val), pdesys.depvars)
-    # Get all dependent variables in the correct type
-    alldepvars = get_all_depvars(pdesys, depvar_ops)
-    alldepvars = filter(u -> !any(map(x -> x isa Number, arguments(u))), alldepvars)
-    # Get all independent variables in the correct type, removing time from the list
-    allindvars = remove(collect(filter(x -> !(x isa Number), reduce(union, filter(xs -> (!isequal(xs, [t])), map(arguments, alldepvars))))), t)
-    #@show allindvars, typeof.(allindvars)
+    # Check if the boundaries warrant using ODAEProblem, as long as this is allowed in the interface
+    use_ODAE = discretization.use_ODAE
+    if use_ODAE
+        bcivmap = reduce((d1, d2) -> mergewith(vcat, d1, d2), collect(values(boundarymap)))
+        allbcs = mapreduce(x -> bcivmap[x], vcat, v.x̄)
+        if all(bc -> bc.order > 0, allbcs)
+            use_ODAE = false
+        end
+    end
+
+    pdeeqs = pdesys.eqs
+    bcs = pdesys.bcs
 
-    interface_errors(alldepvars, allindvars, discretization)
     # @show alldepvars
     # @show allindvars
 
-    # Get tspan
-    tspan = nothing
-    # Check that inputs make sense
-    if t !== nothing
-        tdomain = pdesys.domain[findfirst(d -> isequal(t.val, d.variables), pdesys.domain)]
-        @assert tdomain.domain isa DomainSets.Interval
-        tspan = (DomainSets.infimum(tdomain.domain), DomainSets.supremum(tdomain.domain))
-    end
+    ############################
+    # Discretization of system
+    ############################
     alleqs = []
     bceqs = []
     # * We wamt to do this in 2 passes
     # * First parse the system and BCs, replacing with DiscreteVariables and DiscreteDerivatives
     # * periodic parameters get type info on whether they are periodic or not, and if they join up to any other parameters
     # * Then we can do the actual discretization by recursively indexing in to the DiscreteVariables
 
-    pdeorders = Dict(map(x -> x => d_orders(x, pdeeqs), allindvars))
-    bcorders = Dict(map(x -> x => d_orders(x, bcs), allindvars))
-
-    orders = Dict(map(x -> x => collect(union(pdeorders[x], bcorders[x])), allindvars))
-
     # Create discretized space and variables, this is called `s` throughout
-    s = DiscreteSpace(domain, alldepvars, allindvars, discretization)
-    # Create a map of each variable to their boundary conditions and get the initial condition
-    boundarymap, u0 = parse_bcs(pdesys.bcs, s, depvar_ops, tspan, bcorders)
-    # Generate a map of each variable to whether it is periodic in a given direction
-    pmap = PeriodicMap(boundarymap, s)
+    s = DiscreteSpace(v, discretization)
     # Get the interior and variable to solve for each equation
+    #TODO: do the interiormap before and independent of the discretization i.e. `s`
     interiormap = InteriorMap(pdeeqs, boundarymap, s, discretization, pmap)
+    # Get the derivative orders appearing in each equation
+    pdeorders = Dict(map(x -> x => d_orders(x, pdeeqs), v.x̄))
+    bcorders = Dict(map(x -> x => d_orders(x, bcs), v.x̄))
+    orders = Dict(map(x -> x => collect(union(pdeorders[x], bcorders[x])), v.x̄))
+
     # Generate finite difference weights
     derivweights = DifferentialDiscretizer(pdesys, s, discretization, orders)
 
+    ics = t === nothing ? [] : mapreduce(u -> boundarymap[u][t], vcat, operation.(s.ū))
+
+    bcmap = Dict(map(collect(keys(boundarymap))) do u
+        u => Dict(map(s.x̄) do x
+            x => boundarymap[u][x]
+        end)
+    end)
+
     ####
     # Loop over equations, Discretizing them and their dependent variables' boundary conditions
     ####
     for pde in pdeeqs
         # Read the dependent variables on both sides of the equation
-        depvars_lhs = get_depvars(pde.lhs, depvar_ops)
-        depvars_rhs = get_depvars(pde.rhs, depvar_ops)
+        depvars_lhs = get_depvars(pde.lhs, v.depvar_ops)
+        depvars_rhs = get_depvars(pde.rhs, v.depvar_ops)
         depvars = collect(depvars_lhs ∪ depvars_rhs)
         depvars = filter(u -> !any(map(x -> x isa Number, arguments(u))), depvars)
 
@@ -83,16 +107,18 @@ function SciMLBase.symbolic_discretize(pdesys::PDESystem, discretization::Method
 
             eqvar = interiormap.var[pde]
 
-            # * Assumes that all variables have same dimensionality
+            eqvarbcs = mapreduce(x -> bcmap[operation(eqvar)][x], vcat, s.x̄)
+
+            # * Assumes that all variables in the equation have same dimensionality except edgevals
             args = params(eqvar, s)
             indexmap = Dict([args[i] => i for i in 1:length(args)])
 
             # Handle boundary values appearing in the equation by creating functions that map each point on the interior to the correct replacement rule
             # Generate replacement rule gen closures for the boundary values like u(t, 1)
-            boundaryvalfuncs = generate_boundary_val_funcs(s, depvars, boundarymap, indexmap, derivweights)
+            boundaryvalfuncs = generate_boundary_val_funcs(s, depvars, bcmap, indexmap, derivweights)
 
             # Generate the boundary conditions for the correct variable
-            for boundary in reduce(vcat, collect(values(boundarymap[operation(eqvar)])))
+            for boundary in eqvarbcs
                 generate_bc_eqs!(bceqs, s, boundaryvalfuncs, interiormap, boundary)
             end
             # Generate extrapolation eqs
@@ -108,7 +134,7 @@ function SciMLBase.symbolic_discretize(pdesys::PDESystem, discretization::Method
         end
     end
 
-    u0 = !isempty(u0) ? reduce(vcat, u0) : u0
+    u0 = generate_ic_defaults(ics, s)
     bceqs = reduce(vcat, bceqs)
     alleqs = reduce(vcat, alleqs)
     alldepvarsdisc = unique(reduce(vcat, vec.(values(s.discvars))))
@@ -117,7 +143,7 @@ function SciMLBase.symbolic_discretize(pdesys::PDESystem, discretization::Method
     defaults = pdesys.ps === nothing || pdesys.ps === SciMLBase.NullParameters() ? u0 : vcat(u0, pdesys.ps)
     ps = pdesys.ps === nothing || pdesys.ps === SciMLBase.NullParameters() ? Num[] : first.(pdesys.ps)
     # Combine PDE equations and BC equations
-    metadata = MOLMetadata(s, discretization, pdesys)
+    metadata = MOLMetadata(s, discretization, pdesys, use_ODAE)
     try
         if t === nothing
             # At the time of writing, NonlinearProblems require that the system of equations be in this form:
@@ -163,49 +189,50 @@ function SciMLBase.discretize(pdesys::PDESystem,discretization::MethodOfLines.MO
     try
         simpsys = structural_simplify(sys)
         if tspan === nothing
+            add_metadata!(simpsys.metadata, sys)
             return prob = NonlinearProblem(simpsys, ones(length(simpsys.states)); discretization.kwargs...)
         else
-            return prob = ODEProblem(simpsys, Pair[], tspan; discretization.kwargs...)
+            # Use ODAE if nessesary
+            if getfield(sys, :metadata) isa MOLMetadata && getfield(sys, :metadata).use_ODAE
+                add_metadata!(simpsys.metadata, DAEProblem(simpsys; discretization.kwargs...))
+                return prob = ODAEProblem(simpsys, Pair[], tspan; discretization.kwargs...)
+            else
+                add_metadata!(simpsys.metadata, sys)
+                return prob = ODEProblem(simpsys, Pair[], tspan; discretization.kwargs...)
+            end
         end
     catch e
         error_analysis(sys, e)
     end
 end
 
 function get_discrete(pdesys, discretization)
-    domain = pdesys.domain
-
     t = discretization.time
-
-    depvar_ops = map(x->operation(x.val),pdesys.depvars)
-    # Get all dependent variables in the correct type
-    alldepvars = get_all_depvars(pdesys, depvar_ops)
-    alldepvars = filter(u -> !any(map(x-> x isa Number, arguments(u))), alldepvars)
-    # Get all independent variables in the correct type, removing time from the list
-    allindvars = remove(collect(filter(x->!(x isa Number), reduce(union, filter(xs->(!isequal(xs, [t])), map(arguments, alldepvars))))), t)
-    #@show allindvars, typeof.(allindvars)
-
-    @warn "`get_discrete` is deprecated, The solution is now automatically wrapped in a PDESolution object, which retrieves the shaped solution much faster than the previously recommended method. See the documentation for more information."
-
-    interface_errors(alldepvars, allindvars, discretization)
-    # @show alldepvars
-    # @show allindvars
-
-    # Get tspan
-    tspan = nothing
-    # Check that inputs make sense
-    if t !== nothing
-        tdomain = pdesys.domain[findfirst(d->isequal(t.val, d.variables), pdesys.domain)]
-        @assert tdomain.domain isa DomainSets.Interval
-        tspan = (DomainSets.infimum(tdomain.domain), DomainSets.supremum(tdomain.domain))
+    cardinalize_eqs!(pdesys)
+
+    ############################
+    # System Parsing and Transformation
+    ############################
+    # Parse the variables in to the right form and store useful information about the system
+    v = VariableMap(pdesys, t)
+    # Check for basic interface errors
+    interface_errors(v.ū, v.x̄, discretization)
+    # Extract tspan
+    tspan = t !== nothing ? v.intervals[t] : nothing
+    # Find the derivative orders in the bcs
+    bcorders = Dict(map(x -> x => d_orders(x, pdesys.bcs), all_ivs(v)))
+    # Create a map of each variable to their boundary conditions including initial conditions
+    boundarymap = parse_bcs(pdesys.bcs, v, bcorders)
+    # Generate a map of each variable to whether it is periodic in a given direction
+    pmap = PeriodicMap(boundarymap, v)
+    # Transform system so that it is compatible with the discretization
+    if discretization.should_transform
+        pdesys, pmap = transform_pde_system!(v, boundarymap, pmap, pdesys)
     end
-    alleqs = []
-    bceqs = []
 
-    # Create discretized space and variables
-    s = DiscreteSpace(domain, alldepvars, allindvars, discretization)
+    s = DiscreteSpace(v, discretization)
 
-    return Dict(vcat([Num(x) => s.grid[x] for x in s.x̄], [Num(u) => s.discvars[u] for u in s.ū]) )
+    return Dict(vcat([Num(x) => s.grid[x] for x in s.x̄], [Num(u) => s.discvars[u] for u in s.ū]))
 end
 
 function ModelingToolkit.ODEFunctionExpr(pdesys::PDESystem,discretization::MethodOfLines.MOLFiniteDifference)