GitHub

Linear parabolic system of PDEs

We can use NeuralPDE to solve the linear parabolic system of PDEs:

\[\begin{aligned} \frac{\partial u}{\partial t} &= a * \frac{\partial^2 u}{\partial x^2} + b_1 u + c_1 w \\ \frac{\partial w}{\partial t} &= a * \frac{\partial^2 w}{\partial x^2} + b_2 u + c_2 w \\ \end{aligned}\]

with initial and boundary conditions:

\[\begin{aligned} u(0, x) &= \frac{b_1 - \lambda_2}{b_2 (\lambda_1 - \lambda_2)} \cdot \cos(\frac{x}{a}) - \frac{b_1 - \lambda_1}{b_2 (\lambda_1 - \lambda_2)} \cdot \cos(\frac{x}{a}) \\ w(0, x) &= 0 \\ u(t, 0) &= \frac{b_1 - \lambda_2}{b_2 (\lambda_1 - \lambda_2)} \cdot e^{\lambda_1t} - \frac{b_1 - \lambda_1}{b_2 (\lambda_1 - \lambda_2)} \cdot e^{\lambda_2t} \\ w(t, 0) &= \frac{e^{\lambda_1}-e^{\lambda_2}}{\lambda_1 - \lambda_2} \\ u(t, 1) &= \frac{b_1 - \lambda_2}{b_2 (\lambda_1 - \lambda_2)} \cdot e^{\lambda_1t} \cdot \cos(\frac{x}{a}) - \frac{b_1 - \lambda_1}{b_2 (\lambda_1 - \lambda_2)} \cdot e^{\lambda_2t} * \cos(\frac{x}{a}) \\ w(t, 1) &= \frac{e^{\lambda_1} \cos(\frac{x}{a})-e^{\lambda_2} \cos(\frac{x}{a})}{\lambda_1 - \lambda_2} \end{aligned}\]

with a physics-informed neural network.

using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimisers,
      OptimizationOptimJL, LineSearches
using Plots
using ModelingToolkit: Interval, infimum, supremum

@parameters t, x
@variables u(..), w(..)
Dxx = Differential(x)^2
Dt = Differential(t)

# Constants
a = 1
b1 = 4
b2 = 2
c1 = 3
c2 = 1
λ1 = (b1 + c2 + sqrt((b1 + c2)^2 + 4 * (b1 * c2 - b2 * c1))) / 2
λ2 = (b1 + c2 - sqrt((b1 + c2)^2 + 4 * (b1 * c2 - b2 * c1))) / 2

# Analytic solution
θ(t, x) = exp(-t) * cos(x / a)
function u_analytic(t, x)
    (b1 - λ2) / (b2 * (λ1 - λ2)) * exp(λ1 * t) * θ(t, x) -
    (b1 - λ1) / (b2 * (λ1 - λ2)) * exp(λ2 * t) * θ(t, x)
end
w_analytic(t, x) = 1 / (λ1 - λ2) * (exp(λ1 * t) * θ(t, x) - exp(λ2 * t) * θ(t, x))

# Second-order constant-coefficient linear parabolic system
eqs = [Dt(u(x, t)) ~ a * Dxx(u(x, t)) + b1 * u(x, t) + c1 * w(x, t),
    Dt(w(x, t)) ~ a * Dxx(w(x, t)) + b2 * u(x, t) + c2 * w(x, t)]

# Boundary conditions
bcs = [u(0, x) ~ u_analytic(0, x),
    w(0, x) ~ w_analytic(0, x),
    u(t, 0) ~ u_analytic(t, 0),
    w(t, 0) ~ w_analytic(t, 0),
    u(t, 1) ~ u_analytic(t, 1),
    w(t, 1) ~ w_analytic(t, 1)]

# Space and time domains
domains = [x ∈ Interval(0.0, 1.0),
    t ∈ Interval(0.0, 1.0)]

# Neural network
input_ = length(domains)
n = 15
chain = [Chain(Dense(input_, n, σ), Dense(n, n, σ), Dense(n, 1)) for _ in 1:2]

strategy = StochasticTraining(500)
discretization = PhysicsInformedNN(chain, strategy)

@named pdesystem = PDESystem(eqs, bcs, domains, [t, x], [u(t, x), w(t, x)])
prob = discretize(pdesystem, discretization)
sym_prob = symbolic_discretize(pdesystem, discretization)

pde_inner_loss_functions = sym_prob.loss_functions.pde_loss_functions
bcs_inner_loss_functions = sym_prob.loss_functions.bc_loss_functions

callback = function (p, l)
    if p.iter % 500 == 0
        println("iter: ", p.iter)
        println("loss: ", l)
        println("pde_losses: ", map(l_ -> l_(p.u), pde_inner_loss_functions))
        println("bcs_losses: ", map(l_ -> l_(p.u), bcs_inner_loss_functions))
    end
    return false
end

res = solve(prob, OptimizationOptimisers.Adam(1e-2); maxiters = 5000, callback)

phi = discretization.phi

# Analysis
ts, xs = [infimum(d.domain):0.01:supremum(d.domain) for d in domains]
depvars = [:u, :w]
minimizers_ = [res.u.depvar[depvars[i]] for i in 1:length(chain)]

analytic_sol_func(t, x) = [u_analytic(t, x), w_analytic(t, x)]
u_real = [[analytic_sol_func(t, x)[i] for t in ts for x in xs] for i in 1:2]
u_predict = [[phi[i]([t, x], minimizers_[i])[1] for t in ts for x in xs] for i in 1:2]
diff_u = [abs.(u_real[i] .- u_predict[i]) for i in 1:2]
ps = []
for i in 1:2
    p1 = plot(ts, xs, u_real[i], linetype = :contourf, title = "u$i, analytic")
    p2 = plot(ts, xs, u_predict[i], linetype = :contourf, title = "predict")
    p3 = plot(ts, xs, diff_u[i], linetype = :contourf, title = "error")
    push!(ps, plot(p1, p2, p3))
end
iter: 500
loss: 1.7167671221769436
pde_losses: [0.35955257144058506, 0.13022429671159852]
bcs_losses: [0.011055790343847058, 0.0005433573109808677, 0.569451349441112, 0.19704164555117307, 0.18797390193510027, 0.056034532213962734]
iter: 1000
loss: 0.18692367522782471
pde_losses: [0.08606397325364208, 0.01951037659450489]
bcs_losses: [0.006445291690481083, 0.0016813031450510995, 0.02685426553712979, 0.011871819966290045, 0.022909847568795706, 0.009567551187314177]
iter: 1500
loss: 0.07168562289706792
pde_losses: [0.029639557600507387, 0.01016453520629952]
bcs_losses: [0.0024327716065268147, 0.00043874961550646166, 0.007433986455374967, 0.004488887324721494, 0.010286438911272682, 0.005027863621170717]
iter: 2000
loss: 0.03639350179215006
pde_losses: [0.014039495378085545, 0.00645134444324304]
bcs_losses: [0.0010133335925061016, 0.0003135000195673327, 0.002091931612152612, 0.0017006011719617076, 0.006965335227274094, 0.00302041669199628]
iter: 2500
loss: 0.02696918620319554
pde_losses: [0.012685984312105448, 0.003978030687715619]
bcs_losses: [0.0003940189654414142, 4.48255530463332e-5, 0.0011140547559155454, 0.000990623471999216, 0.005248331437906526, 0.001876426807840671]
iter: 3000
loss: 0.01444819804723309
pde_losses: [0.007012071610050898, 0.0024326128342307698]
bcs_losses: [0.00022486563798978723, 1.7695308120258917e-5, 0.000639318315435444, 0.0009233701755863071, 0.003533581653208711, 0.0011689183671796629]
iter: 3500
loss: 0.015518082468218387
pde_losses: [0.006096947301853994, 0.0037376167498057197]
bcs_losses: [0.00011038475456107152, 3.435406428126275e-5, 0.0005690792211963005, 0.0005354942930674767, 0.003201742116680345, 0.0009425568284614059]
iter: 4000
loss: 0.011558913746885528
pde_losses: [0.005742030035309972, 0.001569468150585766]
bcs_losses: [0.0001286123489053328, 2.7296150237732625e-5, 0.00040654121817287534, 0.00041031407827816173, 0.0018106504257582746, 0.000774361199997924]
iter: 4500
loss: 0.008593244876890626
pde_losses: [0.002777628199684725, 0.0016725633944506403]
bcs_losses: [3.9164805691471554e-5, 3.2697439437220684e-7, 0.00037312308685235273, 0.0003306916036623714, 0.0018255445920676012, 0.0006599144619675411]
iter: 5000
loss: 0.008944989174796938
pde_losses: [0.006170243721448381, 0.0023806396047505018]
bcs_losses: [2.8782516221149686e-5, 0.00010716255866999601, 0.0003559635623837615, 0.00039736596141239945, 0.0013709204807985675, 0.00043197828156926206]
iter: 5000
loss: 0.006303144386988414
pde_losses: [0.0024478100879705557, 0.00139045425197305]
bcs_losses: [5.827184885935512e-5, 1.1199526795857598e-6, 0.0003048440074957438, 0.00030702084218067917, 0.0018349514362775014, 0.0005206888182025267]
ps[1]
Example block output
ps[2]
Example block output