allow multiple depvars in MOLFiniteDifference

src/MOLFiniteDifference/MOL_discretization.jl

@@ -12,10 +12,12 @@ MOLFiniteDifference(dxs, time; upwind_order = 1, centered_order = 2) =
     MOLFiniteDifference(dxs, time, upwind_order, centered_order)
 
 function SciMLBase.symbolic_discretize(pdesys::ModelingToolkit.PDESystem,discretization::DiffEqOperators.MOLFiniteDifference)
+
     pdeeqs = pdesys.eqs isa Vector ? pdesys.eqs : [pdesys.eqs]
     t = discretization.time
     nottime = filter(x->~isequal(x.val, t.val),pdesys.indvars)
     nspace = length(nottime)
+    depvars = pdesys.depvars
 
     # Discretize space
     space = map(nottime) do x
@@ -31,7 +33,7 @@ function SciMLBase.symbolic_discretize(pdesys::ModelingToolkit.PDESystem,discret
 
     # Build symbolic variables
     indices = CartesianIndices(((axes(s)[1] for s in space)...,))
-    depvars = map(pdesys.depvars) do u
+    depvarsdisc = map(depvars) do u
         [Num(Variable{Symbolics.FnType{Tuple{Any}, Real}}(Base.nameof(ModelingToolkit.operation(u.val)),II.I...))(t) for II in indices]
     end
     spacevals = map(y->[Pair(nottime[i],space[i][y.I[i]]) for i in 1:nspace],indices)
@@ -40,18 +42,21 @@ function SciMLBase.symbolic_discretize(pdesys::ModelingToolkit.PDESystem,discret
     ### INITIAL AND BOUNDARY CONDITIONS ###
     # Build symbolic maps for boundaries
     edges = reduce(vcat,[[vcat([Colon() for j in 1:i-1],1,[Colon() for j in i+1:nspace]),
-      vcat([Colon() for j in 1:i-1],size(depvars[1],i),[Colon() for j in i+1:nspace])] for i in 1:nspace])
+      vcat([Colon() for j in 1:i-1],size(depvarsdisc[1],i),[Colon() for j in i+1:nspace])] for i in 1:nspace])
 
     #edgeindices = [indices[e...] for e in edges]
     edgevals = reduce(vcat,[[nottime[i]=>first(space[i]),nottime[i]=>last(space[i])] for i in 1:length(space)])
-    edgevars = [[d[e...] for e in edges] for d in depvars]
-    edgemaps = [spacevals[e...] for e in edges]
-    initmaps = substitute.(pdesys.depvars,[t=>tspan[1]])
+    edgevars = [[d[e...] for e in edges] for d in depvarsdisc]
+
+    bclocs = map(e->substitute.(pdesys.indvars,e),edgevals) # location of the boundary conditions e.g. (t,0.0,y)
+    edgemaps = Dict(bclocs .=> [spacevals[e...] for e in edges])
+    initmaps = substitute.(depvars,[t=>tspan[1]])
 
     # Generate map from variable (e.g. u(t,0)) to discretized variable (e.g. u₁(t))
     subvar(depvar) = substitute.((depvar,),edgevals)
-    depvarmaps = reduce(vcat,[subvar(depvar) .=> edgevars[i] for (i, depvar) in enumerate(pdesys.depvars)])
-    # depvarderivmaps will dictate what to replace the Differential terms with
+    # depvarbcmaps will dictate what to replace the variable terms with in the bcs
+    depvarbcmaps = reduce(vcat,[subvar(depvar) .=> edgevar for (depvar, edgevar) in zip(depvars, edgevars)])
+    # depvarderivbcmaps will dictate what to replace the Differential terms with in the bcs
     if nspace == 1
         # 1D system
         left_weights(j) = DiffEqOperators.calculate_weights(discretization.upwind_order, space[j][1], space[j][1:2])
@@ -60,17 +65,17 @@ function SciMLBase.symbolic_discretize(pdesys::ModelingToolkit.PDESystem,discret
         left_idxs = central_neighbor_idxs(CartesianIndex(2),1)[1:2]
         right_idxs(j) = central_neighbor_idxs(CartesianIndex(length(space[j])-1),1)[end-1:end]
         # Constructs symbolic spatially discretized terms of the form e.g. au₂ - bu₁
-        derivars = [[dot(left_weights(j),depvar[left_idxs]), dot(right_weights(j),depvar[right_idxs(j)])]
-                    for (j, depvar) in enumerate(depvars)]
+        derivars = [[dot(left_weights(1),depvar[left_idxs]), dot(right_weights(1),depvar[right_idxs(1)])]
+                    for depvar in depvarsdisc]
         # Create list of all the symbolic Differential terms evaluated at boundary e.g. Differential(x)(u(t,0))
         subderivar(depvar,s) = substitute.((Differential(s)(depvar),),edgevals)
         # Create map of symbolic Differential terms with symbolic spatially discretized terms
-        depvarderivmaps = reduce(vcat,[subderivar(depvar, s) .=> derivars[i]
-                                       for (i, depvar) in enumerate(pdesys.depvars) for s in nottime])
+        depvarderivbcmaps = reduce(vcat,[subderivar(depvar, s) .=> derivars[i]
+                                       for (i, depvar) in enumerate(depvars) for s in nottime])
    else
         # Higher dimension
         # TODO: Fix Neumann and Robin on higher dimension
-        depvarderivmaps = []
+        depvarderivbcmaps = []
    end
 
     # Generate initial conditions and bc equations
@@ -80,22 +85,23 @@ function SciMLBase.symbolic_discretize(pdesys::ModelingToolkit.PDESystem,discret
         if ModelingToolkit.operation(bc.lhs) isa Sym && ~any(x -> isequal(x, t.val), ModelingToolkit.arguments(bc.lhs))
             # initial condition
             # Assume in the form `u(...) ~ ...` for now
-            push!(u0,vec(depvars[findfirst(isequal(bc.lhs),initmaps)] .=> substitute.((bc.rhs,),spacevals)))
+            push!(u0,vec(depvarsdisc[findfirst(isequal(bc.lhs),initmaps)] .=> substitute.((bc.rhs,),spacevals)))
         else
             # Algebraic equations for BCs
-            i = findfirst(x->occursin(x.val,bc.lhs),first.(depvarmaps))
-
+            i = findfirst(x->occursin(x.val,bc.lhs),first.(depvarbcmaps))
+            bcargs = first(depvarbcmaps[i]).val.arguments
             # Replace Differential terms in the bc lhs with the symbolic spatially discretized terms
             # TODO: Fix Neumann and Robin on higher dimension
-            lhs = nspace == 1 ? substitute(bc.lhs,depvarderivmaps[i]) : bc.lhs
+            lhs = nspace == 1 ? substitute(bc.lhs,depvarderivbcmaps[i]) : bc.lhs
 
             # Replace symbol in the bc lhs with the spatial discretized term
-            lhs = substitute(lhs,depvarmaps[i])
-            rhs = substitute.((bc.rhs,),edgemaps[i])
+            lhs = substitute(lhs,depvarbcmaps[i])
+            rhs = substitute.((bc.rhs,),edgemaps[bcargs])
             lhs = lhs isa Vector ? lhs : [lhs] # handle 1D
             push!(bceqs,lhs .~ rhs)
         end
     end
+
     u0 = reduce(vcat,u0)
     bceqs = reduce(vcat,bceqs)
 
@@ -120,10 +126,10 @@ function SciMLBase.symbolic_discretize(pdesys::ModelingToolkit.PDESystem,discret
         central_neighbor_idxs(II,j) = stencil(j) .+ max(Imin,min(II,Imax))
         central_neighbor_space(II,j) = vec(space[j][map(i->i[j],central_neighbor_idxs(II,j))])
         central_weights(II,j) = DiffEqOperators.calculate_weights(2, space[j][II[j]], central_neighbor_space(II,j))
-        central_deriv(II,j,k) = dot(central_weights(II,j),depvars[k][central_neighbor_idxs(II,j)])
+        central_deriv(II,j,k) = dot(central_weights(II,j),depvarsdisc[k][central_neighbor_idxs(II,j)])
 
-        central_deriv_rules = [(Differential(s)^2)(u) => central_deriv(II,j,k) for (j,s) in enumerate(nottime), (k,u) in enumerate(pdesys.depvars)]
-        valrules = vcat([pdesys.depvars[k] => depvars[k][II] for k in 1:length(pdesys.depvars)],
+        central_deriv_rules = [(Differential(s)^2)(u) => central_deriv(II,j,k) for (j,s) in enumerate(nottime), (k,u) in enumerate(depvars)]
+        valrules = vcat([depvars[k] => depvarsdisc[k][II] for k in 1:length(depvars)],
                         [nottime[j] => space[j][II[j]] for j in 1:nspace])
 
         # TODO: Use rule matching for nonlinear Laplacian
@@ -132,9 +138,9 @@ function SciMLBase.symbolic_discretize(pdesys::ModelingToolkit.PDESystem,discret
         #forward_weights(i,j) = DiffEqOperators.calculate_weights(discretization.upwind_order, 0.0, [space[j][i[j]],space[j][i[j]+1]])
         #reverse_weights(i,j) = DiffEqOperators.calculate_weights(discretization.upwind_order, 0.0, [space[j][i[j]-1],space[j][i[j]]])
         #upwinding_rules = [@rule(*(~~a,(Differential(nottime[j]))(u),~~b) => IfElse.ifelse(*(~~a..., ~~b...,)>0,
-        #                        *(~~a..., ~~b..., dot(reverse_weights(i,j),depvars[k][central_neighbor_idxs(i,j)[1:2]])),
-        #                        *(~~a..., ~~b..., dot(forward_weights(i,j),depvars[k][central_neighbor_idxs(i,j)[2:3]]))))
-        #                        for j in 1:nspace, k in 1:length(pdesys.depvars)]
+        #                        *(~~a..., ~~b..., dot(reverse_weights(i,j),depvarsdisc[k][central_neighbor_idxs(i,j)[1:2]])),
+        #                        *(~~a..., ~~b..., dot(forward_weights(i,j),depvarsdisc[k][central_neighbor_idxs(i,j)[2:3]]))))
+        #                        for j in 1:nspace, k in 1:length(depvars)]
 
         rules = vcat(vec(central_deriv_rules),valrules)
         substitute(eq.lhs,rules) ~ substitute(eq.rhs,rules)
@@ -144,7 +150,7 @@ function SciMLBase.symbolic_discretize(pdesys::ModelingToolkit.PDESystem,discret
     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
-    sys = ODESystem(vcat(eqs,unique(bceqs)),t,vec(reduce(vcat,vec(depvars))),ps,defaults=Dict(defaults))
+    sys = ODESystem(vcat(eqs,unique(bceqs)),t,vec(reduce(vcat,vec(depvarsdisc))),ps,defaults=Dict(defaults))
     sys, tspan
 end
 

test/MOL/MOL_1D_Diffusion.jl

@@ -331,3 +331,54 @@ end
        @test u_approx ≈ exact atol=0.05
     end
 end
+
+@testset "Test 10: linear diffusion, two variables, mixed BCs" begin
+    # Method of Manufactured Solutions
+    u_exact = (x,t) -> exp.(-t) * cos.(x)
+    v_exact = (x,t) -> exp.(-t) * sin.(x)
+
+    # Parameters, variables, and derivatives
+    @parameters t x
+    @variables u(..) v(..)
+    Dt = Differential(t)
+    Dx = Differential(x)
+    Dxx = Dx^2
+
+    # 1D PDE and boundary conditions
+    eqs = [Dt(u(t,x)) ~ Dxx(u(t,x)),
+           Dt(v(t,x)) ~ Dxx(v(t,x))]
+    bcs = [u(0,x) ~ cos(x),
+           v(0,x) ~ sin(x),
+           u(t,0) ~ exp(-t),
+           Dx(u(t,1)) ~ -exp(-t) * sin(1),
+           Dx(v(t,0)) ~ exp(-t),
+           v(t,1) ~ exp(-t) * sin(1)]
+
+    # Space and time domains
+    domains = [t ∈ IntervalDomain(0.0,1.0),
+               x ∈ IntervalDomain(0.0,1.0)]
+
+    # PDE system
+    pdesys = PDESystem(eqs,bcs,domains,[t,x],[u(t,x),v(t,x)])
+
+    # Method of lines discretization
+    l = 100
+    dx = range(0.0,1.0,length=l)
+    order = 2
+    discretization = MOLFiniteDifference([x=>dx],t)
+
+    # Convert the PDE problem into an ODE problem
+    prob = discretize(pdesys,discretization)
+
+    # Solve ODE problem
+    sol = solve(prob,Tsit5(),saveat=0.1)
+
+    x_sol = dx[2:end-1]
+    t_sol = sol.t
+
+    # Test against exact solution
+    for i in 1:length(sol)
+        @test u_exact(x_sol, t_sol[i]) ≈ sol.u[i][1:l-2] atol=0.01
+        @test v_exact(x_sol, t_sol[i]) ≈ sol.u[i][l-1:end] atol=0.01
+    end
+end

test/MOL/MOLtest.jl

@@ -2,17 +2,20 @@ using ModelingToolkit, DiffEqOperators, LinearAlgebra, OrdinaryDiffEq
 
 # Define some variables
 @parameters t x
-@variables u(..)
+@variables u(..) v(..)
 Dt = Differential(t)
 Dxx = Differential(x)^2
-eq  = Dt(u(t,x)) ~ Dxx(u(t,x))
+eqs  = [Dt(u(t,x)) ~ Dxx(u(t,x)), 
+        Dt(v(t,x)) ~ Dxx(v(t,x))]
 bcs = [u(0,x) ~ - x * (x-1) * sin(x),
-       u(t,0) ~ 0.0, u(t,1) ~ 0.0]
+       v(0,x) ~ - x * (x-1) * sin(x),
+       u(t,0) ~ 0.0, u(t,1) ~ 0.0,
+       v(t,0) ~ 0.0, v(t,1) ~ 0.0]
 
 domains = [t ∈ IntervalDomain(0.0,1.0),
            x ∈ IntervalDomain(0.0,1.0)]
 
-pdesys = PDESystem(eq,bcs,domains,[t,x],[u(t,x)])
+pdesys = PDESystem(eqs,bcs,domains,[t,x],[u(t,x),v(t,x)])
 discretization = MOLFiniteDifference([x=>0.1],t)
 prob = discretize(pdesys,discretization) # This gives an ODEProblem since it's time-dependent
 sol = solve(prob,Tsit5())