Investigating symbolic_discretize with the PhysicsInformedNN Discretizer for the 1-D Burgers' Equation

Let's consider the Burgers' equation:

\[\begin{gather*} ∂_t u + u ∂_x u - (0.01 / \pi) ∂_x^2 u = 0 \, , \quad x \in [-1, 1], t \in [0, 1] \, , \\ u(0, x) = - \sin(\pi x) \, , \\ u(t, -1) = u(t, 1) = 0 \, , \end{gather*}\]

with Physics-Informed Neural Networks. Here is an example of using the low-level API:

using ModelingToolkit, NeuralPDE, SciMLBase, Lux, Optimization, OptimizationOptimJL, LineSearches
using Optim: BFGS
using DomainSets: Interval
using IntervalSets: leftendpoint, rightendpoint

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

#2D PDE
eq = Dt(u(t, x)) + u(t, x) * Dx(u(t, x)) - (0.01 / pi) * Dxx(u(t, x)) ~ 0

# Initial and boundary conditions
bcs = [u(0, x) ~ -sin(pi * x),
    u(t, -1) ~ 0.0,
    u(t, 1) ~ 0.0,
    u(t, -1) ~ u(t, 1)]

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

# Neural network
chain = Chain(Dense(2, 16, σ), Dense(16, 16, σ), Dense(16, 1))
strategy = QuadratureTraining(; abstol = 1e-6, reltol = 1e-6, batch = 200)

indvars = [t, x]
depvars = [u(t, x)]
@named pde_system = PDESystem(eq, bcs, domains, indvars, depvars)

discretization = PhysicsInformedNN(chain, strategy)
sym_prob = symbolic_discretize(pde_system, discretization)

phi = sym_prob.phi

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.u), pde_loss_functions))
    println("bcs_losses: ", map(l_ -> l_(p.u), bc_loss_functions))
    return false
end

loss_functions = [pde_loss_functions; bc_loss_functions]

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

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

res = solve(prob, BFGS(linesearch = BackTracking()); maxiters = 3000)

retcode: Success
u: ComponentVector{Float64}(layer_1 = (weight = [-21.530009570115507 -6.075101049179681; -9.484093018754384 4.260660828972828; … ; 21.43737536648253 6.138854447711346; -6.573079873552515 8.54253307299108], bias = [-17.526184684218695, -7.479786617348753, -3.938009311722886, -2.8670955096530366, -1.076952165227065, -0.13945001254653128, 8.545019015254915, -10.735277984996548, 4.365273801043351, 1.0945705334756008, 7.696853496405436, 3.7522224615317454, -4.989783390800272, 3.121961571316755, 14.844884800114395, -8.90309607359854]), layer_2 = (weight = [-1.785790038311257 -0.6834409565653528 … -1.4270347142960242 0.12181196497671108; -3.201305579152684 -0.9440437201660503 … 0.2279606471918537 1.8648594775768255; … ; -1.0039534534107022 -0.8200706271527969 … 5.00597180837146 -1.8639270791274032; -5.037436010193577 -0.8527045994578661 … -3.3565773258367733 1.570170063287898], bias = [-1.1499510536818167, 3.0838776616368717, -2.3443785208748533, -2.360564689280173, 4.148674440655619, 3.5534566734740274, -2.9592705373940627, -7.9343826179426165, -3.3649173307989892, 6.352260289924138, 1.9575356835040272, 3.397840843171905, -0.021293120281407688, 2.289611469317595, -1.3261138382051445, -1.4502236052084754]), layer_3 = (weight = [-0.6950970386978397 0.7608728743990999 … -3.4649851118022794 0.010599606920642343], bias = [0.7210230972518267]))

And some analysis:

using Plots

ts, xs = [leftendpoint(d.domain):0.01:rightendpoint(d.domain) for d in domains]
u_predict_contourf = reshape([first(phi([t, x], res.u)) for t in ts for x in xs],
    length(xs), length(ts))
plot(ts, xs, u_predict_contourf, linetype = :contourf, title = "predict")

u_predict = [[first(phi([t, x], res.u)) for x in xs] for t in ts]
p1 = plot(xs, u_predict[3], title = "t = 0.1");
p2 = plot(xs, u_predict[11], title = "t = 0.5");
p3 = plot(xs, u_predict[end], title = "t = 1");
plot(p1, p2, p3)