using NeuralPDE, Flux, ModelingToolkit, GalacticOptim, Optim, DiffEqFlux
using Quadrature, Cuba, CUDA, QuasiMonteCarlo
import ModelingToolkit: Interval, infimum, supremum
@parameters t, x
@variables u1(..), u2(..), v1(..), v2(..)
Dt = Differential(t)
Dtt = Differential(t)^2
Dx = Differential(x)
Dxx = Differential(x)^2
u1_extract(t,x) = (1/(13*exp(8*t)+720*exp(48*x)-96*exp(24*x+4*t))) * (28 * exp(12*x+2*t)*((exp(4*t)+36*exp(24*x))*cos(3*t-5*x)+(5*exp(4*t)-12*exp(24*x))*sin(3*t-5*x)))
v1_extract(t,x) = -(28 * exp(12*x+2*t)*((-5*exp(4*t)+12*exp(24*x))*cos(3*t-5*x)+(exp(4*t)+36*exp(24*x))*sin(3*t-5*x)))/(13*exp(8*t)+720*exp(48*x)-96*exp(24*x+4*t))
u2_extract(t,x) = -(4 * exp(12*x+2*t)*((-43*exp(4*t)+204*exp(24*x))*cos(3*t-5*x)-(7*exp(4*t)-36*exp(24*x))*sin(3*t-5*x)))/(13*exp(8*t)+720*exp(48*x)-96*exp(24*x+4*t))
v2_extract(t,x) = (4 * exp(12*x+2*t)*((-7*exp(4*t)+252*exp(24*x))*cos(3*t-5*x)+(43*exp(4*t)-204*exp(24*x))*sin(3*t-5*x)))/(13*exp(8*t)+720*exp(48*x)-96*exp(24*x+4*t))
eqs = [0. ~ Dx(u1(t,x))-Dtt(v1(t,x))-2 * u1(t,x)^2 * v1(t,x)-2*v1(t,x)^3+4*v1(t,x)*(u2(t,x)^2+v2(t,x)^2)+4*v2(t,x)*u1(t,x)*u2(t,x)+2*v1(t,x)*u2(t,x)^2-2*v1(t,x)*v2(t,x)^2,
0. ~ Dx(v1(t,x))+Dtt(u1(t,x))+2 * v1(t,x)^2 * u1(t,x)+2*u1(t,x)^3-4*u1(t,x)*(u2(t,x)^2+v2(t,x)^2)-4*v2(t,x)*u2(t,x)*v1(t,x)+2*u1(t,x)*u2(t,x)^2-2*u1(t,x)*v2(t,x)^2,
0. ~ Dx(u2(t,x))-2*v2(t,x)*(u2(t,x)^2+v2(t,x)^2-2*u1(t,x)^2-2*v1(t,x)^2)+4*u1(t,x)*u2(t,x)*v1(t,x)+2*v2(t,x)*(u1(t,x)^2-v1(t,x)^2),
0. ~ Dx(v2(t,x))+2*u2(t,x)*(u2(t,x)^2+v2(t,x)^2-2*u1(t,x)^2-2*v1(t,x)^2)-4*u1(t,x)*v2(t,x)*v1(t,x)+2*u2(t,x)*(u1(t,x)^2-v1(t,x)^2)]
bcs = [u1(0,x) ~ u1_extract(0,x),
u2(0,x) ~ u2_extract(0,x),
v1(0,x) ~ v1_extract(0,x),
v2(0,x) ~ v2_extract(0,x),
u1(t,0) ~ u1_extract(t,0),
u2(t,0) ~ u2_extract(t,0),
v1(t,0) ~ v1_extract(t,0),
v2(t,0) ~ v2_extract(t,0),
u1(t,1) ~ u1_extract(t,1),
u2(t,1) ~ u2_extract(t,1),
v1(t,1) ~ v1_extract(t,1),
v2(t,1) ~ v2_extract(t,1)]
# Space and time domains
domains = [t ∈ Interval(0.0,1.0),
x ∈ Interval(0.0,1.0)]
# Neural network
input_ = length(domains)
n = 2
chain =[FastChain(FastDense(input_,n,Flux.σ),FastDense(n,n,Flux.σ),FastDense(n,1)) for _ in 1:4]
###cuda
initθ = map(c -> Float64.(c), DiffEqFlux.initial_params.(chain))
#initθ = map(c -> CuArray(Float64.(c)), DiffEqFlux.initial_params.(chain))
strategy = GridTraining(0.05)
discretization = PhysicsInformedNN(chain,
strategy;
init_params = initθ)
@named pde_system = PDESystem(eqs,bcs,domains,[t,x],[u1(t, x),u2(t, x),v1(t, x),v2(t, x)])
prob = discretize(pde_system,discretization)
sym_prob = symbolic_discretize(pde_system,discretization)
pde_inner_loss_functions = prob.f.f.loss_function.pde_loss_function.pde_loss_functions.contents
bcs_inner_loss_functions = prob.f.f.loss_function.bcs_loss_function.bc_loss_functions.contents
cb = 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
@time res = GalacticOptim.solve(prob,ADAM(0.01);cb=cb,maxiters=2500)
