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)

retcode: Failure
u: ComponentVector{Float64}(u = (layer_1 = (weight = [-0.8921287459550206 0.7989881761893487; -0.5228512106467211 0.8560274612659605; … ; -0.5794440984615244 -0.4602065019207065; 0.3037544825372233 1.0474285852968888], bias = [-0.40590582866185665; -0.08968827560700116; … ; -0.24111223570596665; 0.34605417264086186;;]), layer_2 = (weight = [-0.4486917122122122 0.21471189404687527 … -0.36717911915531504 0.14242493239744353; 0.06275578958343292 0.5020313616847977 … 0.08096499069580651 -0.05678770175471716; … ; 0.09508709630405675 0.4129714551549028 … -0.17768916459994244 -0.5196775422182645; -0.5654433061416609 0.523406111929351 … -0.08264854513975974 -0.023481611176312332], bias = [-0.3732038376356896; 0.4916632115406574; … ; 0.4707253435712153; 0.5921978667335741;;]), layer_3 = (weight = [-0.465031929772273 -0.15051144220176163 … -0.5977896513977837 0.040259973607291555; -0.4898963786101251 -0.5369307882848194 … 0.4838300400530062 -0.4171593463991003; … ; -0.15568406998503176 -0.38233189798128553 … -0.339946666713513 0.5084470663948238; 0.6839979900485638 -0.16162835083297042 … 0.6042019529619278 0.6555745967369081], bias = [0.14644820461440558; -0.09605165688398047; … ; 0.06218079628334218; -0.5815319316297307;;]), layer_4 = (weight = [0.546562074826394 -0.11617270436841351 … 0.09071031377445016 -0.4052191401544157], bias = [0.07118739623866194;;])), c = (layer_1 = (weight = [-0.5336699777084094; -2.472170996707888; … ; 1.1157597122215257; -0.8130433762499777;;], bias = [-0.05156375042110989; -0.28206632343901633; … ; -0.7183752911244874; 0.43562064551783813;;]), layer_2 = (weight = [-0.2211828633572243 -0.7587923095832849 … -0.6893815517363661 0.7545893628785777; -0.5537221509190258 -0.34508005375465883 … 0.02351893061155048 0.4303703362151079; … ; -0.4416699316635114 0.04180222279993116 … -0.5371212137246701 0.715821218772597; -0.0969838130302924 -0.5478252428162321 … -0.39868059116332266 0.6164994783488134], bias = [0.2608039073774514; 0.07707739399901076; … ; -0.04092750027289554; 0.03467304440008583;;]), layer_3 = (weight = [-0.9843140779704637 -0.6118163475430344 … -0.378530026111043 0.5192519979771147], bias = [0.10541845830814549;;])))

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