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 ModelingToolkit, NeuralPDE, SciMLBase, Lux, Optimization, OptimizationOptimisers,
      OptimizationOptimJL, LineSearches
using Optimisers: Adam
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, 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: 10.982685911703612
pde_losses: [1.3652888891849846, 0.4865205201902776]
bcs_losses: [0.037204786395999855, 0.005512299471114463, 4.2695956362405605, 1.5056538965349162, 1.7246750761495793, 0.6945632684686072]
iter: 1000
loss: 0.3242614929810155
pde_losses: [0.07081701679494287, 0.030239227130667998]
bcs_losses: [0.003753525947941748, 0.001471384187749971, 0.01862967858193809, 0.021150700944111856, 0.06607867660113735, 0.01384817434219347]
iter: 1500
loss: 0.10064767722374734
pde_losses: [0.0378490647191137, 0.011426545624316784]
bcs_losses: [0.0009523975900995343, 9.14863207567324e-5, 0.003911466786054005, 0.005627392257924458, 0.0399275053447993, 0.009108356002185672]
iter: 2000
loss: 0.05882999183421685
pde_losses: [0.023435439073419724, 0.005401622905562114]
bcs_losses: [0.00033257208901539786, 7.717664029307092e-5, 0.000896240229650189, 0.0018770025314022512, 0.019654479030845416, 0.004609803991527654]
iter: 2500
loss: 0.04607771012194292
pde_losses: [0.02524757517319116, 0.004267542425800789]
bcs_losses: [0.0002249476160356975, 2.5184416468160364e-5, 0.0007425634698138174, 0.0013518826788456483, 0.010531590253587903, 0.0034384876655773346]
iter: 3000
loss: 0.024169148722242892
pde_losses: [0.01187065296302505, 0.003121789511990958]
bcs_losses: [0.00019333263593628746, 3.9235840426214156e-5, 0.00047196062862972835, 0.0004780593541419917, 0.006955614186313215, 0.0020752717444719442]
iter: 3500
loss: 0.0236979550425193
pde_losses: [0.01256425781464024, 0.0019222694335662772]
bcs_losses: [0.00012116743365617789, 2.5845175512302643e-5, 0.0003004608818270589, 0.00030963187788618187, 0.005134861489305766, 0.0014176316415715107]
iter: 4000
loss: 0.019585428767150942
pde_losses: [0.009477033774846487, 0.0016477855408582056]
bcs_losses: [0.0001035026290532427, 6.510080544877178e-6, 0.00024210918620064787, 0.00034720468533772287, 0.0035195459430642077, 0.0015901768170969728]
iter: 4500
loss: 0.01904088815172481
pde_losses: [0.011325068909069905, 0.0023046058184401598]
bcs_losses: [0.00015577938875085268, 8.17230819753032e-6, 0.00019506771605853483, 0.0003229533003652332, 0.004236159278131618, 0.0010117966361790623]
iter: 5000
loss: 0.01091586549238897
pde_losses: [0.00496021339466068, 0.001657880318521619]
bcs_losses: [7.133141988892377e-5, 1.2170817441367597e-5, 0.00018024947261885136, 0.00010898353659821181, 0.0032812380461647063, 0.0009618479686524701]
iter: 5000
loss: 0.01077153447585036
pde_losses: [0.004707703352921515, 0.001953057806632979]
bcs_losses: [7.969395786372834e-5, 1.4917239298604327e-5, 0.00012728976186550384, 0.00011262682695324404, 0.0030097975015477203, 0.001086233972684234]
ps[1]
Example block output
ps[2]
Example block output