Allen-Cahn Equation with Sequential Training

In this tutorial we are going to solve the Allen-Cahn equation with periodic boundary condition from $t=0$ to $t=1$. The traning process is split into four stages, namely $t\in [0,0.25]$, $t\in [0.0,0.5]$, $t\in [0.0,0.75]$ and $t\in [0.0, 1.0]$.

using ModelingToolkit, IntervalSets
using Sophon
using Optimization, OptimizationOptimJL, Zygote

@parameters t, x
@variables u(..)
Dₓ = Differential(x)
Dₓ² = Differential(x)^2
Dₜ = Differential(t)

eq = Dₜ(u(x, t)) - 0.0001 * Dₓ²(u(x, t)) + 5 * u(x,t) * (abs2(u(x,t)) - 1.0) ~ 0.0

domain = [x ∈ -1.0..1.0, t ∈ 0.0..0.25]

bcs = [u(x,0) ~ x^2 * cospi(x),
       u(-1,t) ~ u(1,t)]

@named allen = PDESystem(eq, bcs, domain, [x, t], [u(x, t)])

\[ \begin{align} - 0.0001 \frac{\mathrm{d}}{\mathrm{d}x} \frac{\mathrm{d}}{\mathrm{d}x} u\left( x, t \right) + \frac{\mathrm{d}}{\mathrm{d}t} u\left( x, t \right) + 5 u\left( x, t \right) \left( -1 + \left|u\left( x, t \right)\right|^{2} \right) =& 0 \end{align} \]

Then we define the neural net, the sampler, and the training strategy.

chain = FullyConnected(2, 1, tanh; hidden_dims=16, num_layers=4)
pinn = PINN(chain)
sampler = QuasiRandomSampler(500, (300, 100))
strategy = NonAdaptiveTraining(1, (50, 1))
prob = Sophon.discretize(allen, pinn, sampler, strategy)

OptimizationProblem. In-place: true
u0: ComponentVector{Float64}(layer_1 = (weight = [-0.4123325049877167 0.18085990846157074; -2.003528356552124 0.3318171203136444; … ; -1.328696608543396 -0.4236811697483063; -1.7262861728668213 0.8828585743904114], bias = [0.0; 0.0; … ; 0.0; 0.0;;]), layer_2 = (weight = [-0.5651610493659973 -0.09537572413682938 … 0.6447268724441528 0.5269163250923157; 0.15611079335212708 -0.2845761179924011 … 0.6084874868392944 -0.18998602032661438; … ; -0.5248222351074219 -0.033358607441186905 … -0.05383634567260742 0.17499686777591705; 0.05618054419755936 0.632692277431488 … 0.2576780617237091 -0.2994765639305115], bias = [0.0; 0.0; … ; 0.0; 0.0;;]), layer_3 = (weight = [-0.37522342801094055 -0.22957171499729156 … -0.44984138011932373 0.46322518587112427; 0.23610979318618774 -0.18304824829101562 … -0.38943901658058167 -0.3059847056865692; … ; -0.34790030121803284 -0.36346709728240967 … 0.008519652299582958 0.281686395406723; -0.2465709149837494 0.5674788951873779 … -0.5011789202690125 -0.0596284456551075], bias = [0.0; 0.0; … ; 0.0; 0.0;;]), layer_4 = (weight = [-0.19919660687446594 -0.40056362748146057 … -0.5478575229644775 0.09956255555152893; -0.3926990032196045 -0.1641024798154831 … -0.6989591717720032 0.21101084351539612; … ; 0.20064270496368408 -0.460193008184433 … 0.10334718227386475 0.03867632523179054; -0.7065118551254272 -0.6821295619010925 … -0.6051058769226074 0.593750536441803], bias = [0.0; 0.0; … ; 0.0; 0.0;;]), layer_5 = (weight = [0.18680620193481445 0.07704775035381317 … 0.06331620365381241 -0.5933051705360413], bias = [0.0;;]))

We solve the equation sequentially in time.

function train(allen, prob, sampler, strategy)
    bfgs = BFGS()
    res = Optimization.solve(prob, bfgs; maxiters=2000)

    for tmax in [0.5, 0.75, 1.0]
        allen.domain[2] = t ∈ 0.0..tmax
        data = Sophon.sample(allen, sampler)
        prob = remake(prob; u0=res.u, p=data)
        res = Optimization.solve(prob, bfgs; maxiters=2000)
    end
    return res
end

res = train(allen, prob, sampler, strategy)

u: ComponentVector{Float64}(layer_1 = (weight = [-1.1992511260286487 0.5536127834164969; -3.2747182873453964 0.2924714885638415; … ; -1.700834155450946 -0.8325988150115164; -2.729066279090629 1.2675146069205947], bias = [-1.1470844184048485; -0.2869701573202819; … ; 0.06930836106725166; -0.007949014948434358;;]), layer_2 = (weight = [-0.53054274721773 0.24454881012690793 … 0.7538309797738761 0.9323294234673759; 0.1692544226896694 -0.33317210808574316 … 0.43795397733398805 0.17146149514204606; … ; -0.5414666794354713 0.3273027265334088 … -0.04651744440990589 1.7794252323936262; 0.3764977829558713 0.566050667197219 … 0.20943414912156952 -1.7565361414078298], bias = [0.3401488182442814; 0.7075370384259545; … ; 0.46133879533327377; 0.043258930648781474;;]), layer_3 = (weight = [0.2172718510384985 -0.011223553777859116 … -0.39844721595440097 0.09862413332310603; 0.39959332992459823 -0.5928626768880424 … -0.3747041561480196 -0.23492186893216518; … ; -0.6048261792171523 -0.7290943656210735 … 0.07776120507547045 0.11560101388242093; -0.2961644653358888 0.32639105750369696 … -0.12021049449032517 -0.35810919613138836], bias = [-0.3926677485223336; 0.09249829650889098; … ; 0.029064453826893308; -0.11333842278175466;;]), layer_4 = (weight = [-0.5844714879306302 -0.1893341124824068 … -0.5861800696242588 -0.1932791586375084; -0.18734864515715668 -0.41072828078197726 … -0.7320551886898377 -0.027989488636508835; … ; 0.28607710640344136 -0.5819389898050473 … -0.17052390483826899 -0.12821015535326638; -1.3374592687404065 -1.2160589571836942 … -1.0353052183547633 0.6294490207654713], bias = [0.08786238439291946; -0.10921176136585342; … ; 0.029562743104981682; 0.3917391948113058;;]), layer_5 = (weight = [-0.13644872473100467 0.3122978718985719 … 0.03969812566648266 -0.4351041696946015], bias = [-1.2257549195206854;;]))

Let's plot the result.

using CairoMakie

phi = pinn.phi
xs, ts = [infimum(d.domain):0.01:supremum(d.domain) for d in allen.domain]
axis = (xlabel="t", ylabel="x", title="Prediction")
u_pred = [sum(pinn.phi([x, t], res.u)) for x in xs, t in ts]
fig, ax, hm = heatmap(ts, xs, u_pred', axis=axis)