Solving PDEs with ModelingToolkit and NeuralPDE

Solving Poisson PDE Systems (https://docs.sciml.ai/NeuralPDE/stable/tutorials/pdesystem/)

\partial^{2}_{x}u(x,y) + \partial^{2}_{y}u(x,y) = -\sin (\pi x) \sin (\pi y)

(1)

with boundary conditions

\begin{align} u(0, y) &= 0 \\ u(1, y) &= 0 \\ u(x, 0) &= 0 \\ u(x, 1) &= 0 \\ \end{align}

(2)

where

$x ∈ [0, 1], y ∈ [0, 1]$

using NeuralPDE
using Lux
using Optimization
using OptimizationOptimJL
using ModelingToolkit
using DomainSets
using LineSearches
using Plots

2D PDE

@parameters x y
@variables u(..)
Dxx = Differential(x)^2
Dyy = Differential(y)^2
eq  = Dxx(u(x, y)) + Dyy(u(x, y)) ~ -sinpi(x) * sinpi(y)

Differential(x, 2)(u(x, y)) + Differential(y, 2)(u(x, y)) ~ -sinpi(x)*sinpi(y)

Boundary conditions

bcs = [
    u(0, y) ~ 0.0,
    u(1, y) ~ 0.0,
    u(x, 0) ~ 0.0,
    u(x, 1) ~ 0.0
]

4-element Vector{Symbolics.Equation}:
 u(0, y) ~ 0.0
 u(1, y) ~ 0.0
 u(x, 0) ~ 0.0
 u(x, 1) ~ 0.0

Space domains

domains = [
    x ∈ DomainSets.Interval(0.0, 1.0),
    y ∈ DomainSets.Interval(0.0, 1.0)
]

2-element Vector{Symbolics.VarDomainPairing}:
 Symbolics.VarDomainPairing(x, 0.0 .. 1.0)
 Symbolics.VarDomainPairing(y, 0.0 .. 1.0)

Build a neural network for the PDE solver. Input: 2 dimensions. Hidden layers: 16 neurons * 2 layers. Output: single output u(x, y)

dim = 2
chain = Lux.Chain(Dense(dim, 16, Lux.σ), Dense(16, 16, Lux.σ), Dense(16, 1))

Chain(
    layer_1 = Dense(2 => 16, σ),                  # 48 parameters
    layer_2 = Dense(16 => 16, σ),                 # 272 parameters
    layer_3 = Dense(16 => 1),                     # 17 parameters
)         # Total: 337 parameters,
          #        plus 0 states.

Discretization method usesPhysicsInformedNN() (PINN).

dx = 0.05
discretization = PhysicsInformedNN(chain, QuadratureTraining(; batch = 200, abstol = 1e-6, reltol = 1e-6))

NeuralPDE.PhysicsInformedNN{Lux.Chain{@NamedTuple{layer_1::Lux.Dense{typeof(NNlib.σ), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Lux.Dense{typeof(NNlib.σ), Int64, Int64, Nothing, Nothing, Static.True}, layer_3::Lux.Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, NeuralPDE.QuadratureTraining{Float64, Integrals.CubatureJLh}, Nothing, Nothing, NeuralPDE.Phi{LuxCore.StatefulLuxLayerImpl.StatefulLuxLayer{Val{true}, Lux.Chain{@NamedTuple{layer_1::Lux.Dense{typeof(NNlib.σ), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Lux.Dense{typeof(NNlib.σ), Int64, Int64, Nothing, Nothing, Static.True}, layer_3::Lux.Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}, layer_3::@NamedTuple{}}}}, typeof(NeuralPDE.numeric_derivative), Bool, Nothing, Nothing, Nothing, Base.RefValue{Int64}, Base.Pairs{Symbol, Union{}, Nothing, @NamedTuple{}}}(Lux.Chain{@NamedTuple{layer_1::Lux.Dense{typeof(NNlib.σ), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Lux.Dense{typeof(NNlib.σ), Int64, Int64, Nothing, Nothing, Static.True}, layer_3::Lux.Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}((layer_1 = Dense(2 => 16, σ), layer_2 = Dense(16 => 16, σ), layer_3 = Dense(16 => 1)), nothing), NeuralPDE.QuadratureTraining{Float64, Integrals.CubatureJLh}(Integrals.CubatureJLh(0), 1.0e-6, 1.0e-6, 1000, 200), nothing, nothing, NeuralPDE.Phi{LuxCore.StatefulLuxLayerImpl.StatefulLuxLayer{Val{true}, Lux.Chain{@NamedTuple{layer_1::Lux.Dense{typeof(NNlib.σ), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Lux.Dense{typeof(NNlib.σ), Int64, Int64, Nothing, Nothing, Static.True}, layer_3::Lux.Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}, layer_3::@NamedTuple{}}}}(LuxCore.StatefulLuxLayerImpl.StatefulLuxLayer{Val{true}, Lux.Chain{@NamedTuple{layer_1::Lux.Dense{typeof(NNlib.σ), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Lux.Dense{typeof(NNlib.σ), Int64, Int64, Nothing, Nothing, Static.True}, layer_3::Lux.Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}, layer_3::@NamedTuple{}}}(Lux.Chain{@NamedTuple{layer_1::Lux.Dense{typeof(NNlib.σ), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Lux.Dense{typeof(NNlib.σ), Int64, Int64, Nothing, Nothing, Static.True}, layer_3::Lux.Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}((layer_1 = Dense(2 => 16, σ), layer_2 = Dense(16 => 16, σ), layer_3 = Dense(16 => 1)), nothing), nothing, (layer_1 = NamedTuple(), layer_2 = NamedTuple(), layer_3 = NamedTuple()), nothing, Val{true}())), NeuralPDE.numeric_derivative, false, nothing, nothing, nothing, NeuralPDE.LogOptions(50), Base.RefValue{Int64}(1), true, false, Base.Pairs{Symbol, Union{}, Nothing, @NamedTuple{}}())

Build the PDE system and discretize it.

@named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x, y)])
prob = discretize(pde_system, discretization)

OptimizationProblem. In-place: true
u0: ComponentVector{Float64}(layer_1 = (weight = [0.37146589159965515 -0.5519679188728333; -0.06716541200876236 -0.6592244505882263; … ; 0.8543078899383545 0.45701122283935547; -0.14638188481330872 -0.9327391982078552], bias = [-0.2038401961326599, 0.28063225746154785, -0.27332448959350586, 0.36746546626091003, 0.14847500622272491, 0.28951311111450195, -0.6703715920448303, -0.6976760029792786, 0.43775123357772827, 0.424797922372818, -0.2157108634710312, 0.5180280804634094, -0.2626146674156189, 0.189740389585495, -0.6522026658058167, 0.22016257047653198]), layer_2 = (weight = [-0.09551253914833069 -0.3402182161808014 … -0.25514090061187744 -0.049054697155952454; 0.15580157935619354 0.2309907227754593 … -0.13612739741802216 -0.2858422100543976; … ; -0.1595686823129654 -0.38449645042419434 … -0.3608781099319458 0.02025916799902916; 0.29261091351509094 -0.1312141865491867 … -0.33025839924812317 -0.3215297758579254], bias = [0.05825051665306091, 0.010248452425003052, -0.221941739320755, 0.24332138895988464, -0.1420469880104065, 0.09089159965515137, 0.04712700843811035, 0.1027546226978302, -0.024579346179962158, 0.039833009243011475, -0.05245399475097656, -0.08214247226715088, -0.08433565497398376, 0.24101674556732178, 0.15472674369812012, 0.13643532991409302]), layer_3 = (weight = [-0.21647284924983978 0.04427662491798401 … -0.057972367852926254 0.021405113860964775], bias = [0.022627204656600952]))

Callback function to record the loss

lossrecord = Float64[]
callback = function (p, l)
    push!(lossrecord, l)
    return false
end

#2 (generic function with 1 method)

Solve the problem. You can increase maxiters to get a better solution, but it will take more time.

opt = OptimizationOptimJL.LBFGS(linesearch = LineSearches.BackTracking())
@time res = Optimization.solve(prob, opt, callback = callback, maxiters=100)

345.793477 seconds (1.20 G allocations: 68.939 GiB, 1.88% gc time, 34.26% compilation time: <1% of which was recompilation)

retcode: MaxIters
u: ComponentVector{Float64}(layer_1 = (weight = [0.645271180442747 -1.7550527332441908; 0.18123461998312976 -0.5200619247378403; … ; 1.3338696724585148 0.8639004175565185; -0.35565703615634864 -1.2595380546761603], bias = [0.1336027806899353, 0.5006923537203092, -0.9036141043512831, 0.40317961781710376, 0.28723422800425336, -1.4441326075029863, -0.21255473420075865, -0.37473236784017244, 0.480255539439811, 0.43142040853366354, 0.8280907261125402, 0.193245076306045, -0.24026699381623312, 0.09568576282389048, -0.2705962480500785, -0.02602575961302319]), layer_2 = (weight = [-0.09220391268290588 -0.33744004122603893 … -0.3130644260405646 0.02334825854647941; 0.104562966249892 0.2377938216978144 … 0.08002678629254312 -0.09362888404639504; … ; -0.17934646319264347 -0.3880318318005241 … -0.3994769814141501 -0.02095323927172857; 0.37737852967642166 -0.42309744197118965 … -0.5120706050678893 0.1065353972316012], bias = [0.03801486809235068, 0.32767668148158074, -0.16287203031724304, 0.12247725801363635, -0.3950313862863262, 0.2551006142486201, 0.006435036145001657, 0.11520975670823014, -0.13054509130639552, -0.023264239187654526, 0.014168370021076838, -0.056612436962498354, -0.16165866746980173, 0.35175667851731024, 0.08766321730983528, 0.13374352348475876]), layer_3 = (weight = [0.008804012436776476 1.0803123864603252 … -0.22283978896626705 1.600275597341363], bias = [-0.04901198655156345]))

plot(lossrecord, xlabel="Iters", yscale=:log10, ylabel="Loss", lab=false)

Plot the predicted solution of the PDE and compare it with the analytical solution to see the relative error.

xs, ys = [DomainSets.infimum(d.domain):dx/10:DomainSets.supremum(d.domain) for d in domains]
analytic_sol_func(x,y) = (sinpi(x)*sinpi(y))/(2pi^2)

phi = discretization.phi
u_predict = reshape([first(phi([x, y], res.u)) for x in xs for y in ys], (length(xs), length(ys)))
u_real = reshape([analytic_sol_func(x, y) for x in xs for y in ys], (length(xs), length(ys)))
diff_u = abs.(u_predict .- u_real)

201×201 Matrix{Float64}:
 0.000263799  0.00101265   0.00175165   …  0.0242257  0.0244178  0.0246051
 0.000428854  0.000318054  0.0010552       0.0239195  0.024111   0.0242978
 0.00110872   0.000363827  0.000371379     0.0236127  0.0238037  0.0239901
 0.00177575   0.00103294   0.000299751     0.0233055  0.0234959  0.0236818
 0.00242989   0.00168925   0.00095814      0.0229979  0.0231877  0.0233733
 0.00307111   0.00233269   0.00160374   …  0.0226901  0.0228794  0.0230645
 0.00369937   0.00296325   0.00223653      0.022382   0.0225708  0.0227555
 0.00431464   0.00358089   0.00285646      0.022074   0.0222622  0.0224465
 0.0049169    0.00418558   0.00346351      0.0217659  0.0219536  0.0221375
 0.00550614   0.0047773    0.00405766      0.021458   0.0216452  0.0218286
 ⋮                                      ⋱                        ⋮
 0.00994296   0.00955874   0.00917878      0.0303814  0.0308324  0.0312863
 0.010177     0.00978445   0.00939621      0.0308562  0.0313155  0.0317777
 0.0104098    0.0100089    0.00961243      0.0313333  0.031801   0.0322715
 0.0106413    0.0102322    0.00982746   …  0.0318127  0.0322888  0.0327677
 0.0108716    0.0104542    0.0100413       0.0322943  0.0327788  0.0332662
 0.0111005    0.0106749    0.0102539       0.032778   0.0332711  0.0337668
 0.0113282    0.0108944    0.0104653       0.0332639  0.0337654  0.0342696
 0.0115545    0.0111127    0.0106755       0.0337517  0.0342618  0.0347745
 0.0117796    0.0113297    0.0108845    …  0.0342415  0.0347602  0.0352814

p1 = plot(xs, ys, u_real, linetype=:contourf, title = "analytic");
p2 = plot(xs, ys, u_predict, linetype=:contourf, title = "predicted");
p3 = plot(xs, ys, diff_u, linetype=:contourf, title = "error");
plot(p1, p2, p3)

This notebook was generated using Literate.jl.