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)
\]
with boundary conditions
\[\begin{split}
\begin{align}
u(0, y) &= 0 \\
u(1, y) &= 0 \\
u(x, 0) &= 0 \\
u(x, 1) &= 0 \\
\end{align}
\end{split}\]
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)
\[ \begin{equation}
\frac{\mathrm{d}}{\mathrm{d}y} \frac{\mathrm{d}}{\mathrm{d}y} u\left( x, y \right) + \frac{\mathrm{d}}{\mathrm{d}x} \frac{\mathrm{d}}{\mathrm{d}x} u\left( x, y \right) = - sinpi\left( x \right) sinpi\left( y \right)
\end{equation}
\]
Boundary conditions
bcs = [
u(0, y) ~ 0.0,
u(1, y) ~ 0.0,
u(x, 0) ~ 0.0,
u(x, 1) ~ 0.0
]
\[\begin{split} \begin{align}
u\left( 0, y \right) &= 0 \\
u\left( 1, y \right) &= 0 \\
u\left( x, 0 \right) &= 0 \\
u\left( x, 1 \right) &= 0
\end{align}
\end{split}\]
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{}, Tuple{}, @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{}, Tuple{}, @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.008434183895587921 0.33945852518081665; 0.5721961855888367 -0.49010804295539856; … ; 0.17339718341827393 0.19939455389976501; -0.44173645973205566 -0.7201124429702759], bias = [0.16040171682834625, -0.1962549388408661, 0.024615783244371414, 0.43268173933029175, 0.3687910735607147, -0.12853987514972687, -0.4829976558685303, 0.346865177154541, -0.6579456925392151, 0.11307510733604431, -0.0872335210442543, 0.4116462469100952, 0.24788972735404968, 0.574359655380249, 0.33770737051963806, 0.004282625392079353]), layer_2 = (weight = [0.3634105324745178 0.14344875514507294 … -0.322441041469574 0.35081857442855835; 0.3359600901603699 0.1349341720342636 … 0.14219807088375092 -0.043160926550626755; … ; -0.2209242284297943 -0.32282230257987976 … 0.16128475964069366 -0.32098016142845154; 0.09641334414482117 -0.40283894538879395 … -0.4139650762081146 -0.05651670694351196], bias = [0.2455022931098938, -0.14022457599639893, -0.20586863160133362, 0.20865511894226074, 0.24654892086982727, -0.04985317587852478, 0.12346312403678894, 0.0794956386089325, 0.09374406933784485, -0.0521945059299469, 0.186885803937912, 0.09629356861114502, -0.021031945943832397, 0.1179381012916565, -0.07185527682304382, -0.1561564803123474]), layer_3 = (weight = [-0.22462381422519684 -0.12737667560577393 … -0.39575234055519104 0.1895449012517929], bias = [0.21142923831939697]))
Callback function to record the loss
lossrecord = Float64[]
callback = function (p, l)
push!(lossrecord, l)
return false
end
#1 (generic function with 1 method)
Solve the problem. It may take a long time.
opt = OptimizationOptimJL.LBFGS(linesearch = LineSearches.BackTracking())
@time res = Optimization.solve(prob, opt, callback = callback, maxiters=1000)
1573.473829 seconds (6.43 G allocations: 555.666 GiB, 4.19% gc time, 7.71% compilation time: <1% of which was recompilation)
retcode: Success
u: ComponentVector{Float64}(layer_1 = (weight = [0.45022624328154337 0.8677146250651416; 1.2872670349900603 0.46739050998673043; … ; 0.5511880774426848 0.5132039385152055; -1.3422605022556342 -1.312387279262065], bias = [-0.3143067681088624, -0.761472104273127, 0.16325732777774535, 0.3738241227978687, 0.4812597883419682, -2.0080066833053882, -0.49530008395420183, 0.7281272429957351, 1.1665511953202479, 0.027183913678948228, -0.6167604935490174, 1.0982532102064895, 0.9627216220557481, 0.982983125481668, -0.14889750688397047, 2.1847635799930782]), layer_2 = (weight = [0.40349937773403394 0.251053021469712 … -0.2789566808661767 0.3303051643651855; 0.699313183845359 0.4328091753813831 … 0.484698336244529 -0.7390196401795812; … ; 0.03383503444323888 -0.1442621227168473 … 0.4857601721110854 -0.39895335652312586; -0.02504089573235733 -0.4016277367685223 … -0.524414457660456 0.28477729902462756], bias = [0.2688258166417236, -0.08209244269604407, -0.24057742878185898, 0.19426986638247587, 0.48379328931168997, -0.1158942138792143, 0.3055989086748472, -0.08431649530158702, 0.31532487284286503, -0.27401641654933284, 0.1004376252567366, -0.0865516589729508, -0.07862518254565792, 0.2479886620783946, 0.10631247313362016, -0.11514578950223021]), layer_3 = (weight = [0.17775591081724487 -0.9268885618038049 … -0.024484323170092963 -0.2067676649780672], bias = [0.33421859782513175]))
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}:
7.86776e-6 4.43409e-5 9.51207e-5 … 0.000334656 0.000349774
5.55805e-5 3.62715e-6 4.69237e-5 0.000351078 0.000365983
0.000102478 5.07706e-5 4.39749e-7 0.000367834 0.000382553
0.000148544 9.70737e-5 4.69546e-5 0.000384917 0.000399473
0.000193763 0.000142522 9.26064e-5 0.000402317 0.000416735
0.000238121 0.0001871 0.000137382 … 0.000420024 0.000434328
0.000281602 0.000230795 0.000181267 0.000438031 0.000452244
0.000324194 0.000273595 0.000224251 0.000456328 0.000470473
0.000365884 0.000315486 0.000266322 0.000474906 0.000489005
0.00040666 0.000356458 0.000307469 0.000493756 0.000507831
⋮ ⋱ ⋮
0.00181203 0.00172467 0.00163825 0.0037083 0.00376885
0.00179964 0.00171176 0.00162482 0.00381588 0.00387867
0.00178559 0.00169722 0.00160978 0.00392388 0.00398893
0.00176983 0.001681 0.00159308 … 0.00403226 0.00409957
0.00175232 0.00166307 0.0015747 0.004141 0.00421059
0.00173303 0.00164338 0.0015546 0.00425005 0.00432193
0.00171192 0.00162189 0.00153273 0.0043594 0.00443357
0.00168894 0.00159858 0.00150907 0.004469 0.00454548
0.00166406 0.00157339 0.00148356 … 0.00457882 0.00465762
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.