Full-Waveform Inversion: recover a wave speed (2nd-order time)

Download .py All tutorials

The inverse sibling of the vibrating membrane. A wave

\[u_{tt} = c^2\,\Delta u,\qquad u = 0 \text{ on } \partial\Omega,\qquad u(0)=\sin(\pi x)\sin(\pi y),\quad u_t(0)=0,\]

travels through a medium whose speed \(c^2\) is unknown. Given the observed displacement history \(u_\text{obs}(t)\) — a "seismogram" — we recover \(c^2\) by differentiating the time integration itself.

The parameter rides the weak form; fem.solve() differentiates the march

c2  = jno.np.parameter((1,), name="c2")          # the unknown coefficient
fem = jno.fem([ui.tt * vi + c2 * (ui.x * vi.x + ui.y * vi.y),
               u(xb, yb) - 0.0, u(xi0, yi0) - u0, ui0.t - 0.0])

For a second-order form fem.solve() returns the trajectory marched with the energy-conserving trapezoidal (θ=½) rule, reducing \(M_2\ddot u + K u = 0\) to the first-order augmented block in \(y=[u,\,v{=}u_t]\). When a coefficient is a jno.np.parameter the block is re-formed from the parameter each step, and the gradient flows through the whole scan back to \(c^2\) — no custom_root, no hand-written adjoint. The same mechanism recovers a density on the ui.tt term or a shear modulus in a vector (elastodynamic) form: the machinery behind full-waveform inversion and elastography.

c2.initialize(jax.nn.initializers.constant(1.0))     # start at the wrong speed
c2.optimizer(optax.adam(5e-2))
crux = jno.core([(fem.solve() - u_obs).mse], domain=jno.domain.from_array({"_": np.zeros((1, 1))}))
crux.solve(220)                                       # fit c² to the seismogram

What to notice

• A wrong speed makes the wave oscillate at the wrong frequency, so the misfit is sharply informative — the optimizer has a clean gradient to follow.
• The recovered trajectory lands back on top of the data; \(c^2\) is recovered to well under 1%.
• Second-order soft modes need float64 (see the ringing cantilever); jno.fem warns if a u_tt form is assembled without jax_enable_x64.

Result

The wave at the wrong starting speed (dotted) oscillates too slowly; after fitting through the differentiable fem.solve(), the recovered trajectory (dashed) coincides with the observed seismogram (solid) and \(c^2\) is recovered to \(\approx0\%\).

Full script

"""Full-waveform inversion: recover a wave speed from a trajectory through a 2nd-order ``u_tt`` solve.

The inverse sibling of the vibrating membrane. A wave obeys

    u_tt = c² Δu ,    u = 0 on the boundary,    u(t=0) = sin(πx) sin(πy) ,   u_t(t=0) = 0 ,

and the medium's speed ``c²`` is *unknown*. Given the observed displacement history ``u_obs(t)`` (a
"seismogram"), recover ``c²`` by differentiating the **time integration itself**: for a second-order
weak form ``fem.solve()`` returns the trajectory ``u(save_ts)`` marched with the energy-conserving
trapezoidal (θ=½) rule, and the gradient flows through every step of the augmented ``[u, v=u_t]`` block
back to the parameter. ``crux`` then fits it to the data. This is the mechanism behind full-waveform
inversion and elastography — the same call recovers a density on the ``u_tt`` term or a shear modulus
in a vector (elastodynamic) form.

A wrong speed makes the wave oscillate at the wrong frequency, so the misfit is sharply informative;
the recovered trajectory lands back on top of the data. (Second-order soft modes need float64 — see
the ringing-cantilever tutorial — so we opt into x64 up front.)
"""

import jax

jax.config.update("jax_enable_x64", True)

import jax.numpy as jnp  # noqa: E402
import numpy as np  # noqa: E402
import optax  # noqa: E402
from shapely.geometry import box  # noqa: E402

import jno  # noqa: E402
from jno.utils.solver.backend_blocks import _block_time_grid, _default_transient_integrate  # noqa: E402

c2_true, c2_guess = 2.0, 1.0  # the unknown wave speed² to recover, and a deliberately wrong start

# --- forward wave u_tt = c² Δu on a clamped square, plucked into its fundamental mode ---
d = jno.domain(box(0.0, 0.0, 1.0, 1.0), mesh_size=0.14, time=(0.0, 1.5, 60))
u, phi = d.fem_symbols()
xi, yi, ti = d.variable("interior", split=True)
xb, yb, _ = d.variable("boundary", split=True)
xi0, yi0, ti0 = d.variable("initial", split=True)
ui, vi = u.bind(x=xi, y=yi, t=ti), phi.bind(x=xi, y=yi, t=ti)
ui0 = u.bind(x=xi0, y=yi0, t=ti0)
u0 = jno.np.sin(np.pi * xi0) * jno.np.sin(np.pi * yi0)
c2 = jno.np.parameter((1,), name="c2")  # the unknown coefficient
fem = jno.fem([ui.tt * vi + c2 * (ui.x * vi.x + ui.y * vi.y), u(xb, yb) - 0.0, u(xi0, yi0) - u0, ui0.t - 0.0])

# --- the "observed" seismogram: the forward wave at the true speed (the data) ---
blk = fem.operator
ts = np.asarray(_block_time_grid(blk))
u_obs = np.asarray(_default_transient_integrate(blk, {"c2": c2_true}, ts))

# --- recover c² from the data through the differentiable transient solve ---
c2.dtype(jnp.float64)
c2.initialize(jax.nn.initializers.constant(c2_guess))  # start at the wrong speed
c2.optimizer(optax.adam(5e-2))
crux = jno.core([(fem.solve() - u_obs).mse], domain=jno.domain.from_array({"_": np.zeros((1, 1))}))
crux.solve(220)
rec = float(np.asarray(crux.eval([c2])).reshape(-1)[0])

print(f"\nFull-waveform inversion (u_tt = c² Δu):  recovered c² = {rec:.4f}  (truth {c2_true})")
print(f"  started at c² = {c2_guess}   ->   rel-err = {abs(rec - c2_true) / c2_true:.2%}")
assert abs(rec - c2_true) / c2_true < 0.02, f"wave speed not recovered: c²={rec:.4f} (truth {c2_true})"