Full FWI inversion via batched per-shot local windows — Marmousi on the SEM solver¶

The SEM twin of notebook 23: a complete frequency-continuation FWI that recovers Marmousi from a smooth start, where every iteration forwards all shots' local model windows in one batched PropSEM call — shot b propagates in its own (B, nz, W) crop. A towed streamer is translation-invariant (fixed cable), so every window has the same shape and the crops stack into one batched tensor; overlapping windows scatter-add onto the full model automatically through autograd. Runs on a single GPU with impl='c' + memory=BoundarySaving() — the shot batch executes as ONE CUDA kernel sequence (per-shot material arrays, batched C-PML band ring), and because the window-local geometry is identical for every shot, the batch is purely a per-shot-model batch.

The only SEM-specific constraint vs the FD version: the window width must tile into spectral elements, (W - 1) % n_order == 0 (the padding below rounds W up to satisfy it).

In [1]:

import numpy as np
import torch
import matplotlib.pyplot as plt
from scipy.ndimage import gaussian_filter

from sweep.signal import ricker
from sweep.datasets import load_marmousi, MARMOUSI_DH
from sweep.propagator import BoundarySaving
from sweep.sem import PropSEM
from sweep.sem.equations.acoustic_potential2d import AcousticPotentialSEM2D
from sweep.sem.prop_c import sem_c_available

dh = MARMOUSI_DH * 2      # 25 m (12.5 m -> 25 m via ::2,::2 downsample) — SEM notebook 17's grid
dt = 0.002                # GLL-limited (element = n_order*dh = 100 m, vmax 4700)
nt = 1800                 # 3.6 s: the deepest far-offset PRIMARIES land at ~3.0-3.3 s,
                          # just behind the 4 km water direct (~3.3 s) — they fit, barely.
                          # Free-surface multiples continue to >5 s and are deliberately
                          # clipped: at these bands they are not signal we want to fit.

nrec, rec_dx = 80, 2      # 80 receivers every 2 cells (50 m) -> ~4 km streamer cable
srcdx, nshot = 25, 20     # source every 25 cells (625 m), 20 shots
src_z, rec_z = 2, 1
pad_x        = 13         # local-window padding around each shot's aperture (cells, 325 m)

n_order       = 4         # GLL polynomial degree (per element)
abcn          = 10 * n_order  # C-PML width in grid cells: ten elements

# One broadband Ricker source; the observed data is modelled once, up front.
# Frequency continuation is then done by low-pass filtering syn & obs to each
# band's cutoff inside the loss — not by re-modelling with a new source.
src_freq, src_delay = 8.0, 0.18
bands = [
    {'f_max': 3.0, 'n_iter': 35, 'lr': 20.0},
    {'f_max': 6.0, 'n_iter': 35, 'lr': 14.0},
    {'f_max': 9.0, 'n_iter': 45, 'lr': 10.0},
]
grad_sigma = 2.0

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
impl = 'c' if (device.type == 'cuda' and sem_c_available()) else 'eager'
print('device:', device, '  impl:', impl)

import numpy as np import torch import matplotlib.pyplot as plt from scipy.ndimage import gaussian_filter from sweep.signal import ricker from sweep.datasets import load_marmousi, MARMOUSI_DH from sweep.propagator import BoundarySaving from sweep.sem import PropSEM from sweep.sem.equations.acoustic_potential2d import AcousticPotentialSEM2D from sweep.sem.prop_c import sem_c_available dh = MARMOUSI_DH * 2 # 25 m (12.5 m -> 25 m via ::2,::2 downsample) — SEM notebook 17's grid dt = 0.002 # GLL-limited (element = n_order*dh = 100 m, vmax 4700) nt = 1800 # 3.6 s: the deepest far-offset PRIMARIES land at ~3.0-3.3 s, # just behind the 4 km water direct (~3.3 s) — they fit, barely. # Free-surface multiples continue to >5 s and are deliberately # clipped: at these bands they are not signal we want to fit. nrec, rec_dx = 80, 2 # 80 receivers every 2 cells (50 m) -> ~4 km streamer cable srcdx, nshot = 25, 20 # source every 25 cells (625 m), 20 shots src_z, rec_z = 2, 1 pad_x = 13 # local-window padding around each shot's aperture (cells, 325 m) n_order = 4 # GLL polynomial degree (per element) abcn = 10 * n_order # C-PML width in grid cells: ten elements # One broadband Ricker source; the observed data is modelled once, up front. # Frequency continuation is then done by low-pass filtering syn & obs to each # band's cutoff inside the loss — not by re-modelling with a new source. src_freq, src_delay = 8.0, 0.18 bands = [ {'f_max': 3.0, 'n_iter': 35, 'lr': 20.0}, {'f_max': 6.0, 'n_iter': 35, 'lr': 14.0}, {'f_max': 9.0, 'n_iter': 45, 'lr': 10.0}, ] grad_sigma = 2.0 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') impl = 'c' if (device.type == 'cuda' and sem_c_available()) else 'eager' print('device:', device, ' impl:', impl)

device: cuda   impl: c

1 — Marmousi, streamer geometry, and per-shot local windows¶

The known ~450 m water layer is reset to 1500 m/s and frozen (marine convention). Each shot's window spans its source + trailing cable plus padding; the fixed cable makes every window the same width W — rounded up so (W - 1) % n_order == 0, the SEM element-tiling constraint.

In [2]:

vp_true = load_marmousi('true')[::2, ::2].astype(np.float32).copy()
vp_init = load_marmousi('smooth')[::2, ::2].astype(np.float32).copy()
nz, nx  = vp_true.shape
assert (nz - 1) % n_order == 0 and (nx - 1) % n_order == 0

# Fix the known water layer to 1500 m/s — the smooth start had blurred it.
n_water = int(np.argmax(vp_true.max(1) > 1501))     # first row below the 1500 water column
vp_init[:n_water] = 1500.0
print(f'water layer: rows 0:{n_water} ({n_water * dh:.0f} m) fixed to 1500 m/s')

# Streamer: source steps to the right, a fixed cable trails to the left.
src_x     = (nrec - 1) * rec_dx + 15 + srcdx * np.arange(nshot)
sources   = np.stack([src_x, np.full(nshot, src_z)], axis=1).astype(np.int64)
receivers = np.zeros((nshot, nrec, 2), dtype=np.int64)
for s in range(nshot):
    receivers[s, :, 0] = src_x[s] - (nrec - 1 - np.arange(nrec)) * rec_dx
    receivers[s, :, 1] = rec_z

# Per-shot local windows: [aperture - pad, aperture + pad], all the same width W —
# rounded UP to the next admissible SEM width ((W-1) % n_order == 0).
lo = np.minimum(src_x, receivers[..., 0].min(1)) - pad_x
hi = np.maximum(src_x, receivers[..., 0].max(1)) + pad_x
W  = int((hi - lo).max())
W += (-(W - 1)) % n_order
x0 = np.clip(lo, 0, nx - W).astype(int)

# Geometry rebased into window-local coordinates. The geometry is translation-
# invariant and no window hits the model edge, so the LOCAL coordinates are the
# same for every shot — the batch differs only in its material crops, which is
# exactly PropSEM's per-shot-model batch.
loc_sources         = sources.copy()
loc_sources[:, 0]  -= x0
loc_receivers       = receivers.copy()
for s in range(nshot):
    loc_receivers[s, :, 0] -= x0[s]
assert (loc_sources[:, 0] == loc_sources[0, 0]).all()
assert (loc_receivers[:, :, 0] == loc_receivers[0, :, 0]).all()

print(f'{nshot} shots, window W = {W} cells ({W * dh / 1000:.2f} km) — all the same shape')

vp_true = load_marmousi('true')[::2, ::2].astype(np.float32).copy() vp_init = load_marmousi('smooth')[::2, ::2].astype(np.float32).copy() nz, nx = vp_true.shape assert (nz - 1) % n_order == 0 and (nx - 1) % n_order == 0 # Fix the known water layer to 1500 m/s — the smooth start had blurred it. n_water = int(np.argmax(vp_true.max(1) > 1501)) # first row below the 1500 water column vp_init[:n_water] = 1500.0 print(f'water layer: rows 0:{n_water} ({n_water * dh:.0f} m) fixed to 1500 m/s') # Streamer: source steps to the right, a fixed cable trails to the left. src_x = (nrec - 1) * rec_dx + 15 + srcdx * np.arange(nshot) sources = np.stack([src_x, np.full(nshot, src_z)], axis=1).astype(np.int64) receivers = np.zeros((nshot, nrec, 2), dtype=np.int64) for s in range(nshot): receivers[s, :, 0] = src_x[s] - (nrec - 1 - np.arange(nrec)) * rec_dx receivers[s, :, 1] = rec_z # Per-shot local windows: [aperture - pad, aperture + pad], all the same width W — # rounded UP to the next admissible SEM width ((W-1) % n_order == 0). lo = np.minimum(src_x, receivers[..., 0].min(1)) - pad_x hi = np.maximum(src_x, receivers[..., 0].max(1)) + pad_x W = int((hi - lo).max()) W += (-(W - 1)) % n_order x0 = np.clip(lo, 0, nx - W).astype(int) # Geometry rebased into window-local coordinates. The geometry is translation- # invariant and no window hits the model edge, so the LOCAL coordinates are the # same for every shot — the batch differs only in its material crops, which is # exactly PropSEM's per-shot-model batch. loc_sources = sources.copy() loc_sources[:, 0] -= x0 loc_receivers = receivers.copy() for s in range(nshot): loc_receivers[s, :, 0] -= x0[s] assert (loc_sources[:, 0] == loc_sources[0, 0]).all() assert (loc_receivers[:, :, 0] == loc_receivers[0, :, 0]).all() print(f'{nshot} shots, window W = {W} cells ({W * dh / 1000:.2f} km) — all the same shape')

water layer: rows 0:19 (475 m) fixed to 1500 m/s
20 shots, window W = 185 cells (4.62 km) — all the same shape

In [3]:

fig, axes = plt.subplots(2, 1, figsize=(13, 5), constrained_layout=True)
for ax, m, title in zip(axes, [vp_true, vp_init], ['true', 'smooth start (water fixed)']):
    im = ax.imshow(m, cmap='jet', aspect='auto', vmin=1500, vmax=4700,
                   extent=[0, nx * dh / 1000, nz * dh / 1000, 0])
    ax.scatter(src_x * dh / 1000, np.full(nshot, src_z * dh / 1000),
               s=60, c='red', marker='*', edgecolors='k', linewidths=0.5, label=f'{nshot} shots')
    ax.set_title(f'vp {title}')
    ax.set_xlabel('x (km)')
    ax.set_ylabel('z (km)')
    ax.legend(loc='lower right', fontsize=8, framealpha=0.85)
    fig.colorbar(im, ax=ax, label='vp (m/s)', shrink=0.85)
plt.show()

fig, axes = plt.subplots(2, 1, figsize=(13, 5), constrained_layout=True) for ax, m, title in zip(axes, [vp_true, vp_init], ['true', 'smooth start (water fixed)']): im = ax.imshow(m, cmap='jet', aspect='auto', vmin=1500, vmax=4700, extent=[0, nx * dh / 1000, nz * dh / 1000, 0]) ax.scatter(src_x * dh / 1000, np.full(nshot, src_z * dh / 1000), s=60, c='red', marker='*', edgecolors='k', linewidths=0.5, label=f'{nshot} shots') ax.set_title(f'vp {title}') ax.set_xlabel('x (km)') ax.set_ylabel('z (km)') ax.legend(loc='lower right', fontsize=8, framealpha=0.85) fig.colorbar(im, ax=ax, label='vp (m/s)', shrink=0.85) plt.show()

2 — Solver (impl='c' + batched boundary saving) and the batched local-window forward¶

One PropSEM built for the window shape (nz, W). models=[crops] with crops of shape (B, nz, W) gives every shot its own material model in a single batched CUDA kernel sequence; memory=BoundarySaving() stores the batched C-PML band ring (nt, B, nband) instead of the save-all histories, so the backward stays affordable at any batch size.

In [4]:

solver = PropSEM(AcousticPotentialSEM2D, shape=(nz, W), dh=dh, dt=dt, nt=nt,
                 n_order=n_order, free_surface=('top',), abcn=abcn,
                 memory=BoundarySaving(), f_dominant=src_freq,
                 device=device.type, impl=impl)
print('SEM window mesh:', solver.mesh.n_elements, 'elements,',
      solver.mesh.n_global, 'GLL nodes  (per shot; batch B =', nshot, ')')


def batched_syn(vp_full, wavelet):
    """Forward every shot at once, each in its own local model window (B, nz, W)."""
    crops = torch.stack([vp_full[:, x0[s]:x0[s] + W] for s in range(nshot)])
    return solver(wavelet, loc_sources, loc_receivers, models=[crops])


def lowpass(x, f_max, taper_hz=0.5):
    """Differentiable zero-phase low-pass along the time axis (dim 1)."""
    freqs   = torch.fft.rfftfreq(x.shape[1], d=dt, device=x.device)
    fade_lo = max(f_max - taper_hz, 0.5 * f_max)
    ramp    = 0.5 * (1 + torch.cos(np.pi * (freqs - fade_lo) / (f_max - fade_lo)))
    mask    = torch.where(freqs <= fade_lo, torch.ones_like(freqs),
                          torch.where(freqs >= f_max, torch.zeros_like(freqs), ramp))
    spec = torch.fft.rfft(x, dim=1) * mask.view(1, -1, 1, 1)
    return torch.fft.irfft(spec, n=x.shape[1], dim=1)


# Model the observed data ONCE with the broadband source.
t         = np.arange(nt, dtype=np.float32) * dt
wavelet   = ricker(t - src_delay, f=src_freq).astype(np.float32)
vp_true_t = torch.tensor(vp_true, device=device)
with torch.no_grad():
    obs = batched_syn(vp_true_t, wavelet).detach()
print('observed data modelled once:', tuple(obs.shape))

solver = PropSEM(AcousticPotentialSEM2D, shape=(nz, W), dh=dh, dt=dt, nt=nt, n_order=n_order, free_surface=('top',), abcn=abcn, memory=BoundarySaving(), f_dominant=src_freq, device=device.type, impl=impl) print('SEM window mesh:', solver.mesh.n_elements, 'elements,', solver.mesh.n_global, 'GLL nodes (per shot; batch B =', nshot, ')') def batched_syn(vp_full, wavelet): """Forward every shot at once, each in its own local model window (B, nz, W).""" crops = torch.stack([vp_full[:, x0[s]:x0[s] + W] for s in range(nshot)]) return solver(wavelet, loc_sources, loc_receivers, models=[crops]) def lowpass(x, f_max, taper_hz=0.5): """Differentiable zero-phase low-pass along the time axis (dim 1).""" freqs = torch.fft.rfftfreq(x.shape[1], d=dt, device=x.device) fade_lo = max(f_max - taper_hz, 0.5 * f_max) ramp = 0.5 * (1 + torch.cos(np.pi * (freqs - fade_lo) / (f_max - fade_lo))) mask = torch.where(freqs <= fade_lo, torch.ones_like(freqs), torch.where(freqs >= f_max, torch.zeros_like(freqs), ramp)) spec = torch.fft.rfft(x, dim=1) * mask.view(1, -1, 1, 1) return torch.fft.irfft(spec, n=x.shape[1], dim=1) # Model the observed data ONCE with the broadband source. t = np.arange(nt, dtype=np.float32) * dt wavelet = ricker(t - src_delay, f=src_freq).astype(np.float32) vp_true_t = torch.tensor(vp_true, device=device) with torch.no_grad(): obs = batched_syn(vp_true_t, wavelet).detach() print('observed data modelled once:', tuple(obs.shape))

SEM window mesh: 2970 elements, 47965 GLL nodes  (per shot; batch B = 20 )

observed data modelled once: (20, 1800, 80, 1)

3 — Frequency-continuation FWI loop¶

The observed data is modelled once with the broadband source. Each band then low-passes syn and obs to its cutoff (3 → 6 → 9 Hz) inside the loss and Adam-descends the full model. The gradient is Gaussian-smoothed, the water layer is frozen, and overlapping windows scatter-add back onto vp through autograd. (The cutoffs are 3/4 of notebook 23's 4/8/12 Hz: this grid is 2× coarser, and 9 Hz is where the slowest Marmousi zones still keep ~5 GLL points per wavelength — the order-4 SEM accuracy rule.)

In [5]:

# Freeze the fixed water layer; ramp the gradient in just below the seabed.
water_mask = np.ones(nz, dtype=np.float32)
water_mask[:n_water]            = 0.0
water_mask[n_water:n_water + 4] = np.linspace(0, 1, 4) ** 2

vp                 = torch.tensor(vp_init, device=device, requires_grad=True)
losses, band_edges = [], []

for band in bands:
    f_max, n_iter = band['f_max'], band['n_iter']
    opt    = torch.optim.Adam([vp], lr=band['lr'])
    obs_lp = lowpass(obs, f_max)                       # filter the observed data once per band
    for _ in range(n_iter):
        opt.zero_grad()
        syn_lp = lowpass(batched_syn(vp, wavelet), f_max)
        loss   = 0.5 * ((syn_lp - obs_lp) ** 2).sum()
        loss.backward()

        grad    = gaussian_filter(vp.grad.detach().cpu().numpy(), grad_sigma)
        vp.grad = torch.tensor(grad * water_mask[:, None], device=device)
        opt.step()
        with torch.no_grad():
            vp.clamp_(1500.0, 5000.0)
        losses.append(float(loss.detach()))

    band_edges.append(len(losses))
    print(f'  <= {f_max:.0f} Hz: {n_iter} iters, loss -> {losses[-1]:.3e}')

vp_inv = vp.detach().cpu().numpy()

# Freeze the fixed water layer; ramp the gradient in just below the seabed. water_mask = np.ones(nz, dtype=np.float32) water_mask[:n_water] = 0.0 water_mask[n_water:n_water + 4] = np.linspace(0, 1, 4) ** 2 vp = torch.tensor(vp_init, device=device, requires_grad=True) losses, band_edges = [], [] for band in bands: f_max, n_iter = band['f_max'], band['n_iter'] opt = torch.optim.Adam([vp], lr=band['lr']) obs_lp = lowpass(obs, f_max) # filter the observed data once per band for _ in range(n_iter): opt.zero_grad() syn_lp = lowpass(batched_syn(vp, wavelet), f_max) loss = 0.5 * ((syn_lp - obs_lp) ** 2).sum() loss.backward() grad = gaussian_filter(vp.grad.detach().cpu().numpy(), grad_sigma) vp.grad = torch.tensor(grad * water_mask[:, None], device=device) opt.step() with torch.no_grad(): vp.clamp_(1500.0, 5000.0) losses.append(float(loss.detach())) band_edges.append(len(losses)) print(f' <= {f_max:.0f} Hz: {n_iter} iters, loss -> {losses[-1]:.3e}') vp_inv = vp.detach().cpu().numpy()

  <= 3 Hz: 35 iters, loss -> 1.711e+02

  <= 6 Hz: 35 iters, loss -> 2.533e+04

  <= 9 Hz: 45 iters, loss -> 2.769e+05

4 — Result: recovered model and misfit¶

In [6]:

# Score below the fixed water layer, within the streamer's shallow-resolvable depth.
illum = np.zeros(nx, dtype=bool)
for s in range(nshot):
    illum[x0[s]:x0[s] + W] = True

z_max = 100                       # 2.5 km on the 25 m grid
mask  = np.zeros((nz, nx), dtype=bool)
mask[n_water:z_max, illum] = True

def rmse(v):
    return float(np.sqrt(((v - vp_true)[mask] ** 2).mean()))

dp_inv  = (vp_inv  - vp_init)[mask]
dp_true = (vp_true - vp_init)[mask]
pcc     = float(dp_inv @ dp_true / (np.linalg.norm(dp_inv) * np.linalg.norm(dp_true)))
print(f'model RMSE below seabed : {rmse(vp_init):.1f} -> {rmse(vp_inv):.1f} m/s')
print(f'perturbation correlation: {pcc:.3f}')

fig, axes = plt.subplots(4, 1, figsize=(13, 9),
                         gridspec_kw={'height_ratios': [1, 1, 1, 0.8]},
                         constrained_layout=True)
for ax, m, title in zip(axes[:3], [vp_init, vp_inv, vp_true],
                        ['initial (smooth)', 'inverted (SEM batched local-window FWI)', 'true']):
    im = ax.imshow(m, cmap='jet', aspect='auto', vmin=1500, vmax=4700,
                   extent=[0, nx * dh / 1000, nz * dh / 1000, 0])
    ax.set_title(title)
    ax.set_ylabel('z (km)')
fig.colorbar(im, ax=axes[:3], label='vp (m/s)', shrink=0.6)

axes[3].semilogy(losses, 'k')
for e in band_edges[:-1]:
    axes[3].axvline(e, color='r', ls='--', lw=0.8)
axes[3].set_title('misfit')
axes[3].set_xlabel('iteration')
axes[3].set_ylabel(r'$\frac{1}{2}\,\|syn-obs\|^2$')
plt.show()

# Score below the fixed water layer, within the streamer's shallow-resolvable depth. illum = np.zeros(nx, dtype=bool) for s in range(nshot): illum[x0[s]:x0[s] + W] = True z_max = 100 # 2.5 km on the 25 m grid mask = np.zeros((nz, nx), dtype=bool) mask[n_water:z_max, illum] = True def rmse(v): return float(np.sqrt(((v - vp_true)[mask] ** 2).mean())) dp_inv = (vp_inv - vp_init)[mask] dp_true = (vp_true - vp_init)[mask] pcc = float(dp_inv @ dp_true / (np.linalg.norm(dp_inv) * np.linalg.norm(dp_true))) print(f'model RMSE below seabed : {rmse(vp_init):.1f} -> {rmse(vp_inv):.1f} m/s') print(f'perturbation correlation: {pcc:.3f}') fig, axes = plt.subplots(4, 1, figsize=(13, 9), gridspec_kw={'height_ratios': [1, 1, 1, 0.8]}, constrained_layout=True) for ax, m, title in zip(axes[:3], [vp_init, vp_inv, vp_true], ['initial (smooth)', 'inverted (SEM batched local-window FWI)', 'true']): im = ax.imshow(m, cmap='jet', aspect='auto', vmin=1500, vmax=4700, extent=[0, nx * dh / 1000, nz * dh / 1000, 0]) ax.set_title(title) ax.set_ylabel('z (km)') fig.colorbar(im, ax=axes[:3], label='vp (m/s)', shrink=0.6) axes[3].semilogy(losses, 'k') for e in band_edges[:-1]: axes[3].axvline(e, color='r', ls='--', lw=0.8) axes[3].set_title('misfit') axes[3].set_xlabel('iteration') axes[3].set_ylabel(r'$\frac{1}{2}\,\|syn-obs\|^2$') plt.show()

model RMSE below seabed : 337.5 -> 266.7 m/s
perturbation correlation: 0.631

---

Download this notebook — 44_sem_batched_local_window_fwi.ipynb · or view on GitHub