Here is the equation I’m trying to solve and I’m surely doing something wrong - but I don’t know what exactly.
Could you please help me to understand where I’m thinking wrong about it?
using NeuralPDE, Lux, LuxCUDA, Random, ComponentArrays, CUDA
using Optimization
using OptimizationOptimisers
import ModelingToolkit: Interval
device = gpu_device()
println("Setting Parameters")
@parameters t x y
@variables u(..)
Dx = Differential(x)
Dy = Differential(y)
Dxx = Differential(x)^2
Dyy = Differential(y)^2
Dt = Differential(t)
Dtt = Differential(t)^2
t_min = 0.0
t_max = 2.0
x_min = 0.0
x_max = 2.0
x_center = (x_min + x_max)/2
y_min = 0.0
y_max = 2.0
y_center = (y_min + y_max)/2
c = 1.0
first_maxiter = 250
#retrain_maxiter = first_maxiter
eq_scale=10000
bc_scale=1000
ic_scale=1000
println("Setting PDE")
# 2D PDE
eq = eq_scale*Dtt(u(t, x, y)) ~ eq_scale*(Dx(c*c*Dx(u(t, x, y))) + Dy(c*c*Dy(u(t, x, y))))
println("Setting IC and BC")
# analytic_sol_func(t, x, y) = exp(x + y) * cos(x + y + 4t)
# # Initial and boundary conditions
# bcs = [u(t_min, x, y) ~ analytic_sol_func(t_min, x, y),
# u(t, x_min, y) ~ analytic_sol_func(t, x_min, y),
# u(t, x_max, y) ~ analytic_sol_func(t, x_max, y),
# u(t, x, y_min) ~ analytic_sol_func(t, x, y_min),
# u(t, x, y_max) ~ analytic_sol_func(t, x, y_max)]
open_bc = [
bc_scale*u(t,x_min,y) ~ bc_scale*c*Dx(u(t,x_min,y)),
bc_scale*u(t,x_max,y) ~ -bc_scale*c*Dx(u(t,x_max,y)),
bc_scale*u(t,x,y_min) ~ bc_scale*c*Dx(u(t,x,y_min)),
bc_scale*u(t,x,y_max) ~ -bc_scale*c*Dx(u(t,x,y_max))
]
# impulse source IC
ic_func(x,y) = (x-x_center+0.2)*exp(-12*((x-x_center)^2+(y-y_center)^2))
ic = [
ic_scale*u(t_min,x,y) ~ ic_scale*ic_func(x, y),
ic_scale*Dt(u(t_min,x,y)) ~ 0.0
]
all_conditions = append!(open_bc, ic)
println("Setting Space and Domain")
# Space and time domains
domains = [t ∈ Interval(t_min, t_max),
x ∈ Interval(x_min, x_max),
y ∈ Interval(y_min, y_max)]
println("Setting NN")
# Neural network
inner = 25
chain = Chain(Dense(3, inner, Lux.σ),
Dense(inner, inner, Lux.σ),
Dense(inner, inner, Lux.σ),
Dense(inner, inner, Lux.σ),
Dense(inner, 1))
println("Prep PINN and cast to F64 GPU")
strategy = GridTraining(0.05)
ps, st = Lux.setup(Random.default_rng(), chain)
ps = ps |> ComponentArray |> device .|> Float64
println("Setting discretization using PINN")
discretization = PhysicsInformedNN(chain,
strategy,
init_params = ps)
println("Putting it all together into PDESystem")
@named pde_system = PDESystem(eq, all_conditions, domains, [t, x, y], [u(t, x, y)])
problem = discretize(pde_system, discretization)
symprob = symbolic_discretize(pde_system, discretization)
println("Start to solve using Optimization.Adam for 2500 max iterations")
callback = function (p, l)
println("Current loss is: $l")
return false
end
res = Optimization.solve(problem, Adam(0.01); callback = callback, maxiters = first_maxiter)
#println("Remake a problem with warmed-up approximation from the above iterations")
#prob = remake(problem, u0 = res.u)
#res = Optimization.solve(problem, Adam(0.001); callback = callback, maxiters = 2500)
println("Inspect the solution")
phi = discretization.phi
ts, xs, ys = [infimum(d.domain):0.1:supremum(d.domain) for d in domains]
#u_real = [analytic_sol_func(t, x, y) for t in ts for x in xs for y in ys]
u_predict = [first(Array(phi(device([t, x, y]), res.u))) for t in ts for x in xs for y in ys]
using Plots
using Printf
function plot_(res)
# Animate
anim = @animate for (i, t) in enumerate(t_min:0.05:t_max)
@info "Animating frame $i..."
#u_real = reshape([analytic_sol_func(t, x, y) for x in xs for y in ys],
# (length(xs), length(ys)))
u_predict = reshape([first(Array(phi(device([t, x, y]), res.u))) for x in xs for y in ys],
length(xs), length(ys))
#u_error = abs.(u_predict .- u_real)
title = @sprintf("predict, t = %.3f", t)
p1 = plot(xs, ys, u_predict, st = :surface, label = "", title = title)
#title = @sprintf("real")
#p2 = plot(xs, ys, u_real, st = :surface, label = "", title = title)
#title = @sprintf("error")
#p3 = plot(xs, ys, u_error, st = :contourf, label = "", title = title)
plot(p1,
#p2,
#p3
)
end
gif(anim, "3pde.gif", fps = 10)
end
println("First plot")
plot_(res)
println("Look at '3pde.gif'")
function plot_wave(res)
anim = @animate for (i,t) in enumerate(t_min:0.05:t_max)
@info "Animating frame $i..."
u_predict = reshape([Array(phi(device([t, x, y]), res.u))[1] for x in xs for y in ys],
length(xs), length(ys))
surface(xs,ys,u_predict,layout=(1,2),size=(640,320),
xlabel="x",ylabel="y",
#zaxis=((-0.1,0.1),"u(x,y)"),
zlabel="u(x,y)",
color=:redsblues,alpha=0.66,
#clims=(-0.1,0.1),
colorbar=:none,
title="Wave equation",dpi=100 )
contour!(xs,ys,u_predict,levels=24,aspect_ratio=1,
subplot=2,
xlabel="x",ylabel="y",
color=:redsblues,
#clims=(-0.1,0.1),
colorbar=:none,
title=@sprintf("t = %.2f",t) )
end
closeall()
gif(anim, "wave_pde.gif", fps = 10)
end
println("Second plot")
plot_wave(res)
println("Look at 'wave_pde.gif'")