Defining Systems of PDEs for Physics-Informed Neural Networks (PINNs)

In this example, we will solve the PDE system:

\[\begin{align*} ∂_t^2 u_1(t, x) & = ∂_x^2 u_1(t, x) + u_3(t, x) \, \sin(\pi x) \, ,\\ ∂_t^2 u_2(t, x) & = ∂_x^2 u_2(t, x) + u_3(t, x) \, \cos(\pi x) \, ,\\ 0 & = u_1(t, x) \sin(\pi x) + u_2(t, x) \cos(\pi x) - e^{-t} \, , \end{align*}\]

with the initial conditions:

\[\begin{align*} u_1(0, x) & = \sin(\pi x) \, ,\\ ∂_t u_1(0, x) & = - \sin(\pi x) \, ,\\ u_2(0, x) & = \cos(\pi x) \, ,\\ ∂_t u_2(0, x) & = - \cos(\pi x) \, , \end{align*}\]

and the boundary conditions:

\[\begin{align*} u_1(t, 0) & = u_1(t, 1) = 0 \, ,\\ u_2(t, 0) & = - u_2(t, 1) = e^{-t} \, , \end{align*}\]

with physics-informed neural networks.

Solution

using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL
import ModelingToolkit: Interval, infimum, supremum

@parameters t, x
@variables u1(..), u2(..), u3(..)
Dt = Differential(t)
Dtt = Differential(t)^2
Dx = Differential(x)
Dxx = Differential(x)^2

eqs = [Dtt(u1(t,x)) ~ Dxx(u1(t,x)) + u3(t,x)*sin(pi*x),
       Dtt(u2(t,x)) ~ Dxx(u2(t,x)) + u3(t,x)*cos(pi*x),
       0. ~ u1(t,x)*sin(pi*x) + u2(t,x)*cos(pi*x) - exp(-t)]

bcs = [u1(0,x) ~ sin(pi*x),
       u2(0,x) ~ cos(pi*x),
       Dt(u1(0,x)) ~ -sin(pi*x),
       Dt(u2(0,x)) ~ -cos(pi*x),
       u1(t,0) ~ 0.,
       u2(t,0) ~ exp(-t),
       u1(t,1) ~ 0.,
       u2(t,1) ~ -exp(-t)]


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

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

strategy = QuadratureTraining()
discretization = PhysicsInformedNN(chain, strategy)

@named pdesystem = PDESystem(eqs,bcs,domains,[t,x],[u1(t, x),u2(t, x),u3(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)
    println("loss: ", l)
    println("pde_losses: ", map(l_ -> l_(p), pde_inner_loss_functions))
    println("bcs_losses: ", map(l_ -> l_(p), bcs_inner_loss_functions))
    return false
end

res = Optimization.solve(prob,BFGS(); callback = callback, maxiters=5000)

phi = discretization.phi
3-element Vector{NeuralPDE.Phi{Lux.Chain{NamedTuple{(:layer_1, :layer_2, :layer_3), Tuple{Lux.Dense{true, typeof(NNlib.sigmoid_fast), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(NNlib.sigmoid_fast), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, NamedTuple{(:layer_1, :layer_2, :layer_3), Tuple{NamedTuple{(), Tuple{}}, NamedTuple{(), Tuple{}}, NamedTuple{(), Tuple{}}}}}}:
 NeuralPDE.Phi{Lux.Chain{NamedTuple{(:layer_1, :layer_2, :layer_3), Tuple{Lux.Dense{true, typeof(NNlib.sigmoid_fast), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(NNlib.sigmoid_fast), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, NamedTuple{(:layer_1, :layer_2, :layer_3), Tuple{NamedTuple{(), Tuple{}}, NamedTuple{(), Tuple{}}, NamedTuple{(), Tuple{}}}}}(Chain(), (layer_1 = NamedTuple(), layer_2 = NamedTuple(), layer_3 = NamedTuple()))
 NeuralPDE.Phi{Lux.Chain{NamedTuple{(:layer_1, :layer_2, :layer_3), Tuple{Lux.Dense{true, typeof(NNlib.sigmoid_fast), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(NNlib.sigmoid_fast), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, NamedTuple{(:layer_1, :layer_2, :layer_3), Tuple{NamedTuple{(), Tuple{}}, NamedTuple{(), Tuple{}}, NamedTuple{(), Tuple{}}}}}(Chain(), (layer_1 = NamedTuple(), layer_2 = NamedTuple(), layer_3 = NamedTuple()))
 NeuralPDE.Phi{Lux.Chain{NamedTuple{(:layer_1, :layer_2, :layer_3), Tuple{Lux.Dense{true, typeof(NNlib.sigmoid_fast), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(NNlib.sigmoid_fast), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, NamedTuple{(:layer_1, :layer_2, :layer_3), Tuple{NamedTuple{(), Tuple{}}, NamedTuple{(), Tuple{}}, NamedTuple{(), Tuple{}}}}}(Chain(), (layer_1 = NamedTuple(), layer_2 = NamedTuple(), layer_3 = NamedTuple()))

Direct Construction via symbolic_discretize

One can take apart the pieces and reassemble the loss functions using the symbolic_discretize interface. Here is an example using the components from symbolic_discretize to fully reproduce the discretize optimization:

using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL
import ModelingToolkit: Interval, infimum, supremum

@parameters t, x
@variables u1(..), u2(..), u3(..)
Dt = Differential(t)
Dtt = Differential(t)^2
Dx = Differential(x)
Dxx = Differential(x)^2

eqs = [Dtt(u1(t,x)) ~ Dxx(u1(t,x)) + u3(t,x)*sin(pi*x),
       Dtt(u2(t,x)) ~ Dxx(u2(t,x)) + u3(t,x)*cos(pi*x),
       0. ~ u1(t,x)*sin(pi*x) + u2(t,x)*cos(pi*x) - exp(-t)]

bcs = [u1(0,x) ~ sin(pi*x),
       u2(0,x) ~ cos(pi*x),
       Dt(u1(0,x)) ~ -sin(pi*x),
       Dt(u2(0,x)) ~ -cos(pi*x),
       u1(t,0) ~ 0.,
       u2(t,0) ~ exp(-t),
       u1(t,1) ~ 0.,
       u2(t,1) ~ -exp(-t)]

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

# Neural network
input_ = length(domains)
n = 15
chain =[Lux.Chain(Dense(input_,n,Lux.σ),Dense(n,n,Lux.σ),Dense(n,1)) for _ in 1:3]
@named pdesystem = PDESystem(eqs,bcs,domains,[t,x],[u1(t, x),u2(t, x),u3(t, x)])

strategy = NeuralPDE.QuadratureTraining()
discretization = PhysicsInformedNN(chain, strategy)
sym_prob = NeuralPDE.symbolic_discretize(pdesystem, discretization)

pde_loss_functions = sym_prob.loss_functions.pde_loss_functions
bc_loss_functions = sym_prob.loss_functions.bc_loss_functions

callback = function (p, l)
    println("loss: ", l)
    println("pde_losses: ", map(l_ -> l_(p), pde_loss_functions))
    println("bcs_losses: ", map(l_ -> l_(p), bc_loss_functions))
    return false
end

loss_functions =  [pde_loss_functions;bc_loss_functions]

function loss_function(θ,p)
    sum(map(l->l(θ) ,loss_functions))
end

f_ = OptimizationFunction(loss_function, Optimization.AutoZygote())
prob = Optimization.OptimizationProblem(f_, sym_prob.flat_init_params)

res = Optimization.solve(prob,OptimizationOptimJL.BFGS(); callback = callback, maxiters=5000)
u: ComponentVector{Float64}(depvar = (u1 = (layer_1 = (weight = [1.2598805207166084 0.312221155791691; 0.5991973097018655 -3.1370855708509375; … ; -0.0897172321101584 -1.246478313475155; -2.298061639364025 -1.7628062933046438], bias = [0.4827680415744099; 1.5065065347905968; … ; -0.6191023319214995; -3.262200851411611;;]), layer_2 = (weight = [-1.8897127592176763 -1.0171121895690165 … 1.2765696003615532 -2.3756166927447424; -0.9701765339340646 1.7332740398418294 … 0.5490493744476563 0.6716780243473668; … ; 0.6167311560389348 0.7747379780711173 … 0.9455367327235564 0.29767161929791736; 2.6915001254283766 -1.5771293990897783 … -1.6798352702784602 -0.1796716229124783], bias = [-0.015330828894222106; 0.04595237707608215; … ; 1.177268967817394; 0.3921305255801789;;]), layer_3 = (weight = [-2.281465897521867 -1.0543138264298755 … 1.190768150058002 5.24810726630908], bias = [1.332795998322404;;])), u2 = (layer_1 = (weight = [2.7585984175398672 1.2658718166943468; -0.8562290372640357 3.0131786431434677; … ; -1.8369712113466707 -0.23463759507858956; 0.9194471061161975 -1.7758062344467382], bias = [2.705308907214537; 2.0493305315522856; … ; -0.909287578749412; 0.081151784192459;;]), layer_2 = (weight = [0.6648019135550823 0.6086275376773939 … 0.8667479718346237 -2.7023157974871355; 1.3242452943521743 -0.4115547233113682 … -0.48223737352069596 1.3942489348197808; … ; 0.167863895993335 0.2317097723696383 … 0.5307978778504792 -0.2587623061555858; -0.1668280072512618 0.5725235925290392 … 0.2433584591650294 -0.8409120576719155], bias = [0.4105719460776892; -0.18337804220398332; … ; 0.6214976201777749; 0.7460385052686629;;]), layer_3 = (weight = [2.2956443857238615 -2.4010033184719304 … 1.3442200427080557 1.6735165030801273], bias = [-0.9685823737254043;;])), u3 = (layer_1 = (weight = [0.031307350538772726 0.34562828967355635; -0.6078539157133697 -0.06311492828018693; … ; -0.6070488059055709 -0.531524656965756; -1.3111339199739425 -0.31953541727203766], bias = [1.1707520393501158; 0.6675672446174573; … ; 0.2356769355731398; 0.598749909846506;;]), layer_2 = (weight = [-0.26280977482997075 -0.6679585087836086 … -0.22022721445500137 -0.6072690835864913; -0.15519243099113036 -0.7795138866467479 … -0.3981726011481543 -0.6437885441845802; … ; -0.494273873774728 -0.12879890874105232 … -0.8050634541326316 -0.8109402258040336; 0.07050135252913804 -0.44405374930006913 … -0.29095243887567573 -0.1140977813646072], bias = [-0.6888208368373587; -1.1039487712144542; … ; -0.7509493969286664; -0.16149304554017657;;]), layer_3 = (weight = [-1.329062916291498 -4.072938955595914 … -1.8183028723435832 -2.397609342721431], bias = [5.0266838927965205;;]))))

Solution Representation

Now let's perform some analysis for both the symbolic_discretize and discretize APIs:

using Plots

phi = discretization.phi
ts,xs = [infimum(d.domain):0.01:supremum(d.domain) for d in domains]

minimizers_ = [res.u.depvar[sym_prob.depvars[i]] for i in 1:3]

analytic_sol_func(t,x) = [exp(-t)*sin(pi*x), exp(-t)*cos(pi*x), (1+pi^2)*exp(-t)]
u_real  = [[analytic_sol_func(t,x)[i] for t in ts for x in xs] for i in 1:3]
u_predict  = [[phi[i]([t,x],minimizers_[i])[1] for t in ts  for x in xs] for i in 1:3]
diff_u = [abs.(u_real[i] .- u_predict[i] ) for i in 1:3]
for i in 1:3
    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");
    plot(p1,p2,p3)
    savefig("sol_u$i")
end

sol_uq1

sol_uq2

sol_uq3

Notice here that the solution is represented in the OptimizationSolution with u as the parameters for the trained neural network. But, for the case where the neural network is from Lux.jl, it's given as a ComponentArray where res.u.depvar.x corresponds to the result for the neural network corresponding to the dependent variable x, i.e. res.u.depvar.u1 are the trained parameters for phi[1] in our example. For simpler indexing, you can use res.u.depvar[:u1] or res.u.depvar[Symbol(:u,1)] as shown here.

Subsetting the array also works, but is inelegant.

(If param_estim == true, then res.u.p are the fit parameters)

If Flux.jl is used, then subsetting the array is required. This looks like:

init_params = [Flux.destructure(c)[1] for c in chain]
acum =  [0;accumulate(+, length.(init_params))]
sep = [acum[i]+1 : acum[i+1] for i in 1:length(acum)-1]
minimizers_ = [res.minimizer[s] for s in sep]

Note: Solving Matrices of PDEs

Also, in addition to vector systems, we can use the matrix form of PDEs:

using ModelingToolkit, NeuralPDE
@parameters x y
@variables u[1:2,1:2](..)
@derivatives Dxx''~x
@derivatives Dyy''~y

# Initial and boundary conditions
bcs = [u[1](x,0) ~ x, u[2](x,0) ~ 2, u[3](x,0) ~ 3, u[4](x,0) ~ 4]

# matrix PDE
eqs  = @. [(Dxx(u_(x,y)) + Dyy(u_(x,y))) for u_ in u] ~ -sin(pi*x)*sin(pi*y)*[0 1; 0 1]

size(eqs)
(2, 2)