Merge pull request #400 from akashkgarg/issue-282-nonlinearproblem

creating nonlinear problem with ImplicitFiniteDifference

Author: ChrisRackauckas

Project.toml

@@ -14,6 +14,7 @@ LazyBandedMatrices = "d7e5e226-e90b-4449-9968-0f923699bf6f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 RuntimeGeneratedFunctions = "7e49a35a-f44a-4d26-94aa-eba1b4ca6b47"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
@@ -36,6 +37,7 @@ SciMLBase = "1.11"
 StaticArrays = "0.10, 0.11, 0.12, 1.0"
 SymbolicUtils = "0.11, 0.12"
 julia = "1"
+NonlinearSolve = "0.3.7"
 
 [extras]
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"

docs/src/symbolic_tutorials/mol_heat.md

@@ -7,7 +7,7 @@ In this tutorial we will use the symbolic interface to solve the heat equation.
 ### Dirichlet boundary conditions
 
 ```julia
-using OrdinaryDiffEq, ModelingToolkit, DiffEqOperators
+using OrdinaryDiffEq, ModelingToolkit, DiffEqOperators, DomainSets
 # Method of Manufactured Solutions: exact solution
 u_exact = (x,t) -> exp.(-t) * cos.(x)
 
@@ -59,7 +59,7 @@ savefig("plot.png")
 ### Neumann boundary conditions
 
 ```julia
-using OrdinaryDiffEq, ModelingToolkit, DiffEqOperators
+using OrdinaryDiffEq, ModelingToolkit, DiffEqOperators, DomainSets
 # Method of Manufactured Solutions: exact solution
 u_exact = (x,t) -> exp.(-t) * cos.(x)
 
@@ -114,7 +114,7 @@ savefig("plot.png")
 ### Robin boundary conditions
 
 ```julia
-using OrdinaryDiffEq, ModelingToolkit, DiffEqOperators
+using OrdinaryDiffEq, ModelingToolkit, DiffEqOperators, DomainSets
 # Method of Manufactured Solutions
 u_exact = (x,t) -> exp.(-t) * sin.(x)
 
@@ -204,3 +204,77 @@ anim = @animate for i in 1:length(t)
 end
 gif(anim, "plot.gif",fps=30)
 ```
+
+### Stationary Problems
+
+```julia
+using ModelingToolkit,DiffEqOperators,LinearAlgebra,DomainSets
+using ModelingToolkit: Differential
+using DifferentialEquations
+
+@parameters x
+@variables u(..)
+Dxx = Differential(x)^2
+
+eq = Dxx(u(x)) ~ 0
+dx = 0.1
+
+bcs = [u(0) ~ 1,
+        u(1) ~ 2]
+
+# Space and time domains
+domains = [x ∈ Interval(0.0,1.0)]
+
+pdesys = PDESystem([eq],bcs,domains,[x],[u(x)])
+
+# Note that we pass in `nothing` for the time variable `t` here since we
+# are creating a stationary problem without a dependence on time, only space.
+discretization = MOLFiniteDifference([x=>dx], nothing, centered_order=2)
+prob = discretize(pdesys,discretization)
+sol = solve(prob)
+
+using Plots
+xs = domains[1].domain.lower:dx:domains[1].domain.upper
+plot(xs, sol.u)
+
+```
+
+```julia
+using ModelingToolkit,DiffEqOperators,LinearAlgebra,DomainSets
+using ModelingToolkit: Differential
+using DifferentialEquations
+
+@parameters x y
+@variables u(..)
+Dxx = Differential(x)^2
+Dyy = Differential(y)^2
+
+eq = Dxx(u(x, y)) + Dyy(u(x, y))~ 0
+dx = 0.1
+dy = 0.1
+
+bcs = [u(0,y) ~ 0.0,
+       u(1,y) ~ y,
+       u(x,0) ~ 0.0,
+       u(x,1) ~ x]
+
+# Space and time domains
+domains = [x ∈ Interval(0.0,1.0),
+           y ∈ Interval(0.0,1.0)]
+
+pdesys = PDESystem([eq],bcs,domains,[x,y],[u(x,y)])
+
+# Note that we pass in `nothing` for the time variable `t` here since we
+# are creating a stationary problem without a dependence on time, only space.
+discretization = MOLFiniteDifference([x=>dx,y=>dy], nothing, centered_order=2)
+
+prob = discretize(pdesys,discretization)
+sol = solve(prob)
+
+using Plots
+xs,ys = [infimum(d.domain):dx:supremum(d.domain) for d in domains]
+u_sol = reshape(sol.u, (length(xs),length(ys)))
+plot(xs, ys, u_sol, linetype=:contourf,title = "solution")
+```
+![solution](https://user-images.githubusercontent.com/620651/121951649-0b701e00-cd10-11eb-8af8-b203f7f13abc.png)
+

src/MOLFiniteDifference/MOL_discretization.jl

@@ -4,6 +4,7 @@ import DomainSets
 # Method of lines discretization scheme
 
 @enum GridAlign center_align edge_align
+
 struct MOLFiniteDifference{T,T2} <: DiffEqBase.AbstractDiscretization
     dxs::T
     time::T2
@@ -45,10 +46,13 @@ function SciMLBase.symbolic_discretize(pdesys::ModelingToolkit.PDESystem,discret
     pdeeqs = pdesys.eqs isa Vector ? pdesys.eqs : [pdesys.eqs]
     t = discretization.time
     # Get tspan
-    tdomain = pdesys.domain[findfirst(d->isequal(t.val, d.variables),pdesys.domain)]
-    @assert tdomain.domain isa DomainSets.Interval
-    tspan = (tdomain.domain.lower,tdomain.domain.upper)
-    
+    tspan = nothing
+    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
+
     depvar_ops = map(x->operation(x.val),pdesys.depvars)
 
     u0 = []
@@ -66,7 +70,7 @@ function SciMLBase.symbolic_discretize(pdesys::ModelingToolkit.PDESystem,discret
         # Read the independent variables,
         # ignore if the only argument is [t]
         allindvars = Set(filter(xs->!isequal(xs,[t]),map(arguments,depvars)))
-        allnottime = Set(filter(!isempty,map(u->filter(x->!isequal(x,t.val),arguments(u)),depvars)))
+        allnottime = Set(filter(!isempty,map(u->filter(x-> t == nothing || !isequal(x,t.val),arguments(u)),depvars)))
         if isempty(allnottime)
             push!(alleqs,eq)
             push!(alldepvarsdisc,depvars)
@@ -87,7 +91,7 @@ function SciMLBase.symbolic_discretize(pdesys::ModelingToolkit.PDESystem,discret
             space = map(nottime) do x
                 xdomain = pdesys.domain[findfirst(d->isequal(x, d.variables),pdesys.domain)]
                 dx = discretization.dxs[findfirst(dxs->isequal(x, dxs[1].val),discretization.dxs)][2]
-                dx isa Number ? (xdomain.domain.lower:dx:xdomain.domain.upper) : dx
+                dx isa Number ? (DomainSets.infimum(xdomain.domain):dx:DomainSets.supremum(xdomain.domain)) : dx
             end
             dxs = map(nottime) do x        
                 dx = discretization.dxs[findfirst(dxs->isequal(x, dxs[1].val),discretization.dxs)][2]
@@ -112,7 +116,9 @@ function SciMLBase.symbolic_discretize(pdesys::ModelingToolkit.PDESystem,discret
             space_indices = CartesianIndices(((axes(s)[1] for s in space)...,))
             grid_indices = CartesianIndices(((axes(g)[1] for g in grid)...,))
             depvarsdisc = map(depvars) do u
-                if isequal(arguments(u),[t])
+                if t == nothing
+                    [Num(Variable{Real}(Base.nameof(operation(u)),II.I...)) for II in grid_indices]
+                elseif isequal(arguments(u),[t])
                     [u for II in grid_indices]
                 else
                     [Num(Variable{Symbolics.FnType{Tuple{Any}, Real}}(Base.nameof(operation(u)),II.I...))(t) for II in grid_indices]
@@ -134,7 +140,10 @@ function SciMLBase.symbolic_discretize(pdesys::ModelingToolkit.PDESystem,discret
 
             bclocs = map(e->substitute.(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]])
+            initmaps = depvars
+            if t != nothing
+                initmaps = substitute.(depvars,[t=>tspan[1]])
+            end
 
             # Generate map from variable (e.g. u(t,0)) to discretized variable (e.g. u₁(t))
             subvar(depvar) = substitute.((depvar,),edgevals)
@@ -179,7 +188,7 @@ function SciMLBase.symbolic_discretize(pdesys::ModelingToolkit.PDESystem,discret
             for bc in pdesys.bcs
                 bcdepvar = first(get_depvars(bc.lhs, depvar_ops))
                 if any(u->isequal(operation(u),operation(bcdepvar)),depvars)
-                    if operation(bc.lhs) isa Sym && !any(x -> isequal(x, t.val), arguments(bc.lhs))
+                    if t != nothing && operation(bc.lhs) isa Sym && !any(x -> isequal(x, t.val), arguments(bc.lhs))
                         # initial condition
                         # Assume in the form `u(...) ~ ...` for now
                         i = findfirst(isequal(bc.lhs),initmaps)
@@ -334,7 +343,7 @@ function SciMLBase.symbolic_discretize(pdesys::ModelingToolkit.PDESystem,discret
             push!(alldepvarsdisc,reduce(vcat,depvarsdisc))
         end
     end
-    u0 = reduce(vcat,u0)
+    u0 = !isempty(u0) ? reduce(vcat,u0) : u0
     bceqs = reduce(vcat,bceqs)
     alleqs = reduce(vcat,alleqs)
     alldepvarsdisc = unique(reduce(vcat,alldepvarsdisc))
@@ -343,12 +352,25 @@ 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(alleqs,unique(bceqs)),t,vec(reduce(vcat,vec(alldepvarsdisc))),ps,defaults=Dict(defaults))
-    sys, tspan
+    if t == nothing
+        # At the time of writing, NonlinearProblems require that the system of equations be in this form:
+        # 0 ~ ...
+        # Thus, before creating a NonlinearSystem we normalize the equations s.t. the lhs is zero.
+        eqs = map(eq -> 0 ~ eq.rhs - eq.lhs, vcat(alleqs,unique(bceqs)))
+        sys = NonlinearSystem(eqs,vec(reduce(vcat,vec(alldepvarsdisc))),ps,defaults=Dict(defaults))
+        return sys, nothing
+    else
+        sys = ODESystem(vcat(alleqs,unique(bceqs)),t,vec(reduce(vcat,vec(alldepvarsdisc))),ps,defaults=Dict(defaults))
+        return sys, tspan
+    end
 end
 
 function SciMLBase.discretize(pdesys::ModelingToolkit.PDESystem,discretization::DiffEqOperators.MOLFiniteDifference)
     sys, tspan = SciMLBase.symbolic_discretize(pdesys,discretization)
-    simpsys = structural_simplify(sys)
-    prob = ODEProblem(simpsys,Pair[],tspan)
+    if tspan == nothing
+        return prob = NonlinearProblem(sys, ones(length(sys.states)))
+    else
+        simpsys = structural_simplify(sys)
+        return prob = ODEProblem(simpsys,Pair[],tspan)
+    end
 end

test/MOL/MOL_NonlinearProblem.jl

@@ -0,0 +1,114 @@
+# 1D diffusion problem
+
+# Packages and inclusions
+using ModelingToolkit,DiffEqOperators,LinearAlgebra,Test
+using ModelingToolkit: Differential
+using NonlinearSolve
+using DomainSets
+
+
+# Handrolled discrete laplace problem
+@testset "Hand rolled 1D Laplace" begin
+    @variables a b c d
+    eqs = [0 ~ a - 1,
+           0 ~ a + c - 2*b,
+           0 ~ b + d - 2*c,
+           0 ~ d - 1]
+    ns = NonlinearSystem(eqs, [a, b, c, d], [])
+    f = eval(generate_function(ns, [a, b, c, d])[2])
+    prob = NonlinearProblem(ns, zeros(4), [])
+    sol = NonlinearSolve.solve(prob, NewtonRaphson())
+    @test sol.u ≈ ones(4)
+end
+
+# Laplace's Equation, same as above but with MOL discretization
+@testset "1D Laplace - constant solution" begin
+    @parameters x
+    @variables u(..)
+    Dxx = Differential(x)^2
+
+    eq = Dxx(u(x)) ~ 0
+    dx = 1/3
+
+    bcs = [u(0) ~ 1,
+           u(1) ~ 1]
+
+    # Space and time domains
+    domains = [x ∈ Interval(0.0,1.0)]
+
+    pdesys = PDESystem([eq],bcs,domains,[x],[u(x)])
+    discretization = MOLFiniteDifference([x=>dx], nothing, centered_order=2)
+    prob = discretize(pdesys,discretization)
+    sol = NonlinearSolve.solve(prob, NewtonRaphson())
+
+    @test sol.u ≈ ones(4)
+end
+
+# Laplace's Equation, linear solution
+@testset "1D Laplace - linear solution" begin
+    @parameters x
+    @variables u(..)
+    Dxx = Differential(x)^2
+
+    eq = Dxx(u(x)) ~ 0
+    dx = 0.1
+
+    bcs = [u(0) ~ 1,
+           u(1) ~ 2]
+
+    # Space and time domains
+    domains = [x ∈ Interval(0.0,1.0)]
+
+    pdesys = PDESystem([eq],bcs,domains,[x],[u(x)])
+    discretization = MOLFiniteDifference([x=>dx], nothing, centered_order=2)
+    prob = discretize(pdesys,discretization)
+    sol = NonlinearSolve.solve(prob, NewtonRaphson())
+
+    @test sol.u ≈ 1.0:0.1:2.0
+end
+
+# 2D heat
+@testset "2D heat equation" begin
+    @parameters x y
+    @variables u(..)
+    Dxx = Differential(x)^2
+    Dyy = Differential(y)^2
+
+    eq = Dxx(u(x, y)) + Dyy(u(x, y))~ 0
+    dx = 0.1
+    dy = 0.1
+
+    bcs = [u(0,y) ~ x*y,
+           u(1,y) ~ x*y,
+           u(x,0) ~ x*y,
+           u(x,1) ~ x*y]
+
+
+    # Space and time domains
+    domains = [x ∈ Interval(0.0,1.0),
+               y ∈ Interval(0.0,1.0)]
+
+    pdesys = PDESystem([eq],bcs,domains,[x,y],[u(x,y)])
+
+    # Note that we pass in `nothing` for the time variable `t` here since we
+    # are creating a stationary problem without a dependence on time, only space.
+    discretization = MOLFiniteDifference([x=>dx,y=>dy], nothing, centered_order=2)
+
+    prob = discretize(pdesys,discretization)
+    sol = NonlinearSolve.solve(prob, NewtonRaphson())
+    xs,ys = [infimum(d.domain):dx:supremum(d.domain) for d in domains]
+    u_sol = reshape(sol.u, (length(xs),length(ys)))
+
+    # test boundary
+    @test all(abs.(u_sol[:,1]) .< eps(Float32))
+    @test all(abs.(u_sol[1,:]) .< eps(Float32))
+    @test u_sol[:,end] ≈ 0:dy:1.0
+    @test u_sol[end,:] ≈ 0:dx:1.0
+
+    # test interior with finite differences
+    interior = CartesianIndices((axes(xs)[1], axes(ys)[1]))[2:end-1,2:end-1]
+    fd = map(interior) do I
+        abs(u_sol[(I - CartesianIndex(1, 0))] + u_sol[(I + CartesianIndex(1,0))] - 2*u_sol[I]) < eps(Float32)
+    end
+    @test all(fd)
+end

test/runtests.jl

@@ -43,6 +43,7 @@ if GROUP == "All" || GROUP == "MOLFiniteDifference"
     @time @safetestset "MOLFiniteDifference Interface: 2D Diffusion" begin include("MOL/MOL_2D_Diffusion.jl") end
     @time @safetestset "MOLFiniteDifference Interface: 1D HigherOrder" begin include("MOL/MOL_1D_HigherOrder.jl") end
     @time @safetestset "MOLFiniteDifference Interface: 1D Partial DAE" begin include("MOL/MOL_1D_PDAE.jl") end
+    @time @safetestset "MOLFiniteDifference Interface: Stationary Nonlinear Problems" begin include("MOL/MOL_NonlinearProblem.jl") end
 end
 
 if GROUP == "All" || GROUP == "Misc"