Solving PDE using Julia numerically

using NeuralPDE, Lux, CUDA, Random
using Optimization, OptimizationOptimJL
using OptimizationOptimisers
import ModelingToolkit: Interval

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

# 2D PDE
eq  = Dt(u(x,t)) - Dxx(u(x,t)) ~ -exp(-t)*sin(2*pi*x)+4*pi*pi*exp(-t)*sin(2*pi*x)

# Boundary conditions
bcs = [u(x,0) ~ sin(2*pi*x), u(0,t) ~ 0.0,
       Dx(u(1,t)) ~ 2*pi*exp(-t)]
# Space and time domains
domains = [x ∈ Interval(0.0,1.0),
           t ∈ Interval(0.0,1.0)]

# Neural network
dim = 2 # number of dimensions
chain = Lux.Chain(Dense(dim,20,tanh),Dense(20,20,tanh),Dense(20,20,tanh),Dense(20,20,tanh),Dense(20,1))

ps = Lux.setup(Random.default_rng(), chain)[1]
ps = ps |> Lux.ComponentArray |> gpu .|> Float64

# Discretization
dx = 0.05
discretization = PhysicsInformedNN(chain,GridTraining(dx),init_params = ps)

@named pde_system = PDESystem(eq,bcs,domains,[x,t],[u(x, t)])
prob = discretize(pde_system,discretization)

#Optimizer
opt = OptimizationOptimJL.LBFGS()

using Plots

xs,ts = [infimum(d.domain):dx/10:supremum(d.domain) for d in domains]
analytic_sol_func(x,t) = exp(-t)*sin(2*pi*x)

u_predict = reshape([first(Array(phi(gpu([x,t]), res.u))) for t in ts for x in xs], (length(ts),length(xs)))
u_real = reshape([analytic_sol_func(x,t) for t in ts for x in xs], (length(ts),length(xs)))
diff_u = abs.(u_predict .- u_real)

p1 = plot(xs, ts, u_real, linetype=:contourf,c = :jet,title = "Exact");
p2 = plot(xs, ts, u_predict, linetype=:contourf,c = :jet,title = "Predict");
p3 = plot(xs, ts, diff_u,linetype=:contourf, c = :jet,title = "Abs. Error");

plot(p1,p2,p3)