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(y, 2)(u(x, y)) + Differential(x, 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 = [-1.045601725578308 -0.10921033471822739; 1.091493010520935 -0.18767650425434113; … ; -0.46106433868408203 0.22992028295993805; 0.915541410446167 -0.867740273475647], bias = [-0.3490026295185089, -0.5115582942962646, -0.5690100193023682, -0.08262906223535538, 0.6889359951019287, 0.39675045013427734, -0.3616393506526947, 0.4941050410270691, -0.37188026309013367, 0.45623844861984253, -0.5051567554473877, -0.3407018184661865, 0.7008427381515503, 0.6059771180152893, 0.014963648281991482, 0.5286787748336792]), layer_2 = (weight = [0.06755463033914566 0.16899268329143524 … -0.19711469113826752 0.1807904988527298; 0.4115501940250397 0.06671659648418427 … -0.3102651834487915 -0.019414111971855164; … ; -0.3194742500782013 0.06836268305778503 … -0.030815385282039642 -0.14694052934646606; -0.23023533821105957 0.18522308766841888 … -0.22952422499656677 0.10553862899541855], bias = [0.14362940192222595, -0.08819934725761414, 0.23475956916809082, 0.21281558275222778, -0.12065193057060242, -0.10474029183387756, -0.1616690754890442, 0.24242505431175232, -0.18419599533081055, 0.11514776945114136, -0.03171294927597046, -0.12478068470954895, -0.11491352319717407, 0.13035660982131958, 0.16441944241523743, -0.20528900623321533]), layer_3 = (weight = [-0.3666118085384369 0.03553791716694832 … -0.2301514595746994 0.41848769783973694], bias = [-0.011502742767333984]))

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)

334.433801 seconds (1.13 G allocations: 63.464 GiB, 2.90% gc time, 32.98% compilation time)

retcode: MaxIters
u: ComponentVector{Float64}(layer_1 = (weight = [-2.246108250055807 1.772792883950308; 1.8766584116918135 -0.3764065507640711; … ; -1.4145450849739856 -0.8058043208178515; 0.7056130522230699 -0.31418418088797506], bias = [-0.1930480784297757, -1.1698734758061016, -0.5855957732744279, -0.21603834038289826, 0.1839592933647991, 0.283659901506895, -0.5964637308265895, -0.23007897330147245, -0.4967890085472386, 0.5560042471052046, -1.1236828087320831, -0.39665306905589515, -1.1619675794912372, 1.2913611902823703, 0.39628818821212486, 0.20944882302922274]), layer_2 = (weight = [0.13204165050378144 0.06719821742142248 … -0.26013320357533265 -0.33310747416396935; 0.2875670909744126 0.051790480960252995 … -0.3865584724926057 -0.2807309907110464; … ; -0.044373324478931796 -0.4287483074767748 … 0.21280720388556357 0.47034472573429525; -0.40825879195649667 0.4616449278518109 … -0.26554492068454794 0.5498057157800589], bias = [0.16271509657823247, -0.04749907181511072, 0.335002466496306, 0.5451480638667354, 0.5961443900218559, -0.1294079443748471, -1.2955919627606456, 0.2737882953195958, -0.11805831949736428, 0.1674395861694066, 0.0932211714189041, 0.012309427678531512, 0.36598353897043046, 0.45134862150971616, -0.36972772952940625, -0.04831152180674929]), layer_3 = (weight = [0.10155469024303497 -0.6488592554865341 … -0.620215020089187 0.806569657755751], bias = [-0.8100184935455451]))

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.00795623  0.00739334  0.00683437  …  0.0202576  0.0206742  0.0210874
 0.00772117  0.00717064  0.00662395     0.0196112  0.0200153  0.0204159
 0.0074827   0.0069445   0.00641008     0.0189649  0.0193566  0.0197447
 0.00724084  0.00671494  0.00619274     0.018319   0.0186982  0.0190739
 0.00699558  0.00648195  0.00597196     0.0176736  0.0180404  0.0184036
 0.00674697  0.00624557  0.00574775  …  0.0170289  0.0173833  0.0177342
 0.00649501  0.00600582  0.00552012     0.0163851  0.0167271  0.0170657
 0.00623973  0.0057627   0.00528911     0.0157423  0.0160721  0.0163983
 0.00598116  0.00551626  0.00505472     0.0151007  0.0154182  0.0157323
 0.00571933  0.00526652  0.004817       0.0144605  0.0147659  0.0150677
 ⋮                                   ⋱                        ⋮
 0.0077677   0.00783895  0.00790675     0.0145388  0.0146517  0.0147629
 0.00836831  0.0084323   0.00849285     0.0151165  0.0152385  0.0153589
 0.0089727   0.0090294   0.00908264     0.0156962  0.0158274  0.0159569
 0.00958078  0.00963014  0.00967604  …  0.0162778  0.0164182  0.0165569
 0.0101924   0.0102344   0.010273       0.0168613  0.0170109  0.0171588
 0.0108076   0.0108422   0.0108733      0.0174466  0.0176053  0.0177625
 0.0114261   0.0114533   0.0114769      0.0180337  0.0182016  0.018368
 0.012048    0.0120676   0.0120838      0.0186224  0.0187995  0.0189751
 0.012673    0.0126851   0.0126938   …  0.0192127  0.019399   0.0195838

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.