Solving PDEs with ModelingToolkit and NeuralPDE

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)
_images/52d5f6b198e4d8c989423865391a5c0a509b3a1ca24e999b39fbca78a4942a8b.png

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)
_images/405391c8bf5b3aeb51a4d5aaf915ab0a8e7d58803f56bab3c1aa39ce3451cae0.png

This notebook was generated using Literate.jl.