Inverse problem for the wave equation with unknown velocity field

We are going to sovle the wave equation.

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

@parameters x, t
@variables u(..), c(..)

Dₜ = Differential(t)
Dₜ² = Differential(t)^2
Dₓ² = Differential(x)^2

s(x,t) = abs2(x) * sin(x) * cos(t)

eq = Dₜ²(u(x,t)) ~ c(x) * Dₓ²(u(x,t)) + s(x,t)

bcs = [u(x, 0) ~ sin(x),
       Dₜ(u(x, 0)) ~ 0,
       u(0, t) ~ 0,
       u(1, t) ~ sin(1) * cos(t)]

domains = [t ∈ Interval(0.0, 1.0),
           x ∈ Interval(0.0, 1.0)]

@named wave = PDESystem(eq, bcs, domains, [t,x], [u(x,t),c(x)])

\[ \begin{align} \frac{\mathrm{d}}{\mathrm{d}t} \frac{\mathrm{d}}{\mathrm{d}t} u\left( x, t \right) =& c\left( x \right) \frac{\mathrm{d}}{\mathrm{d}x} \frac{\mathrm{d}}{\mathrm{d}x} u\left( x, t \right) + \cos\left( t \right) \sin\left( x \right) \left|x\right|^{2} \end{align} \]

Here the velocity field $c(x)$ is unknown, we will approximate it with a neural network.

pinn = PINN(u = FullyConnected((2,16,16,16,1), sin),
            c = FullyConnected((1,16,16,1), tanh))

sampler = QuasiRandomSampler(500,100)
strategy = NonAdaptiveTraining(1, (10,10,1,1))

NonAdaptiveTraining{Int64, NTuple{4, Int64}}(1, (10, 10, 1, 1))

Next we generate some data of $u(x,t)$. Here we place two sensors at $x=0.1$ and $x=0.5$.

ū(x,t) = sin(x) * cos(t)

x_data = hcat(fill(0.1, 1, 50), fill(0.5, 1, 50))
t_data = repeat(range(0.0, 1.0, length = 50),2)'
input_data = [x_data; t_data]

u_data = ū.(x_data, t_data)

1×100 Matrix{Float64}:
 0.0998334  0.0998126  0.0997503  0.0996464  …  0.275281  0.267213  0.259035

Finally we construct the inverse problem and solve it.

additional_loss(phi, θ) = sum(abs2, phi.u(input_data, θ.u) .- u_data)

prob = Sophon.discretize(wave, pinn, sampler, strategy; additional_loss=additional_loss)

@showprogress res = Optimization.solve(prob, BFGS(), maxiters=1000)

u: ComponentVector{Float64}(u = (layer_1 = (weight = [0.5346858404160654 -0.6531084752875893; 0.5720237000590264 -0.3812813554898803; … ; 0.8226513855910196 0.7704409153269918; 0.6532116354901227 0.4079344744348018], bias = [0.048469596655155135; 0.4338006695648404; … ; -0.19804694306646906; 0.1901725849850755;;]), layer_2 = (weight = [-0.0014480324239770567 0.12765963589144805 … 0.22732804068300885 -0.1557646924020688; -0.19450405609321286 -0.1515531927153182 … 0.06679978301592092 -0.18335304136809147; … ; -0.3744507046998116 0.15232114436292582 … 0.2892327056050837 -0.18823875848954758; -0.2570868133380263 -0.04440660625743631 … 0.23993747235098356 -0.16291802150349288], bias = [0.4899533182269404; -0.08955382019636499; … ; -0.5541062593943046; -0.06397992177375131;;]), layer_3 = (weight = [0.36799270908270953 0.15013127609671678 … 0.40132209681962105 -0.5172700607542375; 0.009194169611074393 0.3107440912328016 … 0.3419179746016028 0.5752019726903708; … ; -0.8235831779601834 -0.6732985792989248 … 0.0400206940056787 -0.5335829616329995; -0.4019953748855882 -0.15407532394596843 … 0.6350055908060517 0.11742282912825225], bias = [-0.22074165098115398; -0.04231273836642446; … ; 0.12849516582256135; -0.3230680202938522;;]), layer_4 = (weight = [0.7816561013470394 0.6312379403044414 … -0.5288206823930085 -0.11758422657641118], bias = [0.07859731222918238;;])), c = (layer_1 = (weight = [0.36063742342770544; -1.440353218984342; … ; -0.24864323426762952; 0.5604506849933061;;], bias = [0.3770390930243978; 0.2882023133982958; … ; -0.9021765556754562; 0.06349328976530866;;]), layer_2 = (weight = [0.1676362537026551 0.17192464695463594 … -0.4915928678341551 0.6537018630967957; -0.517394065055185 -0.34564205172474133 … 0.4287278313542032 0.0536936462057112; … ; 0.603675962102099 -0.45786534869666934 … -0.25687486847904584 0.3890844325468678; -0.2587856772582603 0.17496915064823562 … 0.8394661360249939 -0.5589801973750052], bias = [-0.051218342575141414; -0.34347414322426195; … ; 0.04427150612062857; -0.3399422462554226;;]), layer_3 = (weight = [-0.3264101712534864 -1.3167345724041881 … 0.20760953298617266 0.44591842782037955], bias = [0.48591092426489496;;])))

Let's visualize the predictted solution and inferred velocity

using CairoMakie

ts = range(0, 1; length=100)
xs = range(0, 1; length=100)

u_pred = [pinn.phi.u([x, t], res.u.u)[1] for x in xs, t in ts]
c_pred = [pinn.phi.c([x], res.u.c)[1] for x in xs]

u_true = [ū(x, t) for x in xs, t in ts]
c_true = 1 .+ abs2.(xs) |> vec

axis = (xlabel="x", ylabel="t", title="Analytical Solution")
fig, ax1, hm1 = heatmap(xs, ts, u_true, axis=axis; colormap=:jet)
ax2, hm2= heatmap(fig[1, end+1], xs, ts, u_pred, axis= merge(axis, (;title = "Prediction")); colormap=:jet)
ax3, hm3 = heatmap(fig[1, end+1], xs, ts, abs.(u_true .- u_pred), axis= merge(axis, (;title = "Absolute Error")); colormap=:jet)
Colorbar(fig[:, end+1], hm3)
fig

fig, ax = lines(xs, c_pred)
lines!(ax, xs, c_true)
fig