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 DomainSets: Interval
using IntervalSets: leftendpoint, rightendpoint

@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 = [leftendpoint(d.domain):0.01:rightendpoint(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: 13.283154673327504
pde_losses: [1.1101036313037216, 0.2634039969792587]
bcs_losses: [0.06072130302937153, 0.009619032200011507, 4.7778447944956355, 2.2793789288299457, 2.639542243823676, 1.2569804914952833]
iter: 1000
loss: 0.5441983187297148
pde_losses: [0.22281570193228223, 0.026487928330832133]
bcs_losses: [0.006134272443289203, 0.00016150304656412415, 0.13001218151443023, 0.030351745038346728, 0.17996654151171676, 0.03735182028482027]
iter: 1500
loss: 0.15829958633210858
pde_losses: [0.05863977196551797, 0.010322541939769605]
bcs_losses: [0.0007356799283408988, 0.000146939936299967, 0.013009142923616051, 0.0040638553283016245, 0.06889528140426307, 0.01642629297372281]
iter: 2000
loss: 0.0792966311018617
pde_losses: [0.03452893560241548, 0.004179735542922449]
bcs_losses: [0.0007670275933000074, 3.151603966064018e-5, 0.006535977694828363, 0.0014476701774704713, 0.034527276306775274, 0.009137379511619382]
iter: 2500
loss: 0.056861319998259344
pde_losses: [0.024339930520579547, 0.0033385078034684477]
bcs_losses: [0.00040067280556567816, 7.772525365867592e-5, 0.005285511777665865, 0.0005390592242470957, 0.017463445482402303, 0.004888123847490373]
iter: 3000
loss: 0.03791803957677668
pde_losses: [0.015209685777000924, 0.0038284119430007004]
bcs_losses: [0.00041051849548485124, 0.00012106638273711919, 0.003506516364554817, 0.0004126459946824166, 0.011437899110800936, 0.0035233914850689558]
iter: 3500
loss: 0.028148307028015337
pde_losses: [0.010261819131539292, 0.0019257626363866737]
bcs_losses: [0.00024809189596442145, 0.00015966300298491702, 0.0025845074023270456, 0.0003381562837522987, 0.011077852120244631, 0.002348525023785131]
iter: 4000
loss: 0.023089240590149385
pde_losses: [0.011816438731825394, 0.0019796735335995503]
bcs_losses: [0.00024833534741912827, 0.00017658327611201234, 0.0017838056045811662, 0.0002914374800611908, 0.00831502847502783, 0.0018273342507727454]
iter: 4500
loss: 0.03536790237167595
pde_losses: [0.026006887209134962, 0.0017394913386576027]
bcs_losses: [9.92554532684798e-5, 0.00015571411612082129, 0.0015015042028978492, 0.0001656735219783601, 0.005873918005427626, 0.0012410570594140454]
iter: 5000
loss: 0.015646847982304323
pde_losses: [0.0062352965404908235, 0.0010121419420031984]
bcs_losses: [9.11579534478643e-5, 0.00017229466446398698, 0.0010144640997764112, 0.0001881306984015227, 0.004968426494985367, 0.0012483612828415797]
iter: 5000
loss: 0.012893457179149213
pde_losses: [0.005144073868892938, 0.0011854364238837408]
bcs_losses: [8.521892070995958e-5, 0.00014069193651657945, 0.0009346026900133812, 0.00015742070161782996, 0.0048500960913518115, 0.001174694663009839]

ps[1]

ps[2]