DEEP_FWI_GIPE

Reference implementation for "Reliability-calibrated deep residual full-waveform inversion using geometry-invariant physics encoding: synthetic validation and zero-shot Marmousi-2 testing" — Deepak Kumar, Jayant Nath Tripathi and Laxmidhar Behera, submitted to Computers & Geosciences.

What this code does

Neural networks that map shot gathers to velocity models memorise the acquisition geometry they were trained on, and their uncertainty estimates become unreliable once the data drift from the training distribution. This repository implements a pipeline that addresses both problems and runs entirely on a laptop.

The network never sees shot gathers. Variable acquisition geometries are mapped into a fixed model-space representation computed by classical physics operators — a geometry-invariant physics encoding (GIPE) of ten channels:

Channel	Content
0	c0 — smoothed starting model
1	c_ADMM — reweighted-L1 / total-variation ADMM-FWI prior
2	misfit gradient evaluated at c_ADMM
3	misfit gradient evaluated at c0 (RTM-like structure image)
4–9	coverage tensor: k_min, k_max, wavenumber fill fraction, orientation entropy, orientation gap, illumination

The learned stage is not a replacement for FWI. It is a calibrated residual corrector applied to the classical prior. A six-member heterogeneous ensemble (U-Net, residual CNN, attention U-Net; two seeds each) predicts a residual together with a per-pixel variance, and is calibrated by physics-conditioned conformal prediction.

For deployment beyond the calibration distribution, a held-out-shot physics audit rescales the intervals without any ground truth: shots that the inversion never used are simulated through samples of the predictive distribution, and the interval width is set where the data-space coverage of those shots peaks.

Requirements

Python 3.11 or newer, and:

torch >= 2.0
deepwave >= 0.0.27
scikit-fmm
numpy
scipy
matplotlib

Install with:

python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

Wave propagation runs on CPU. Ensemble training optionally uses the Apple MPS backend or a CUDA GPU if one is available; nothing in this repository requires a GPU.

Developed and tested on macOS 15 with Python 3.14, PyTorch 2.13 and Deepwave 0.0.27, on an Apple M3 Pro (11 cores, 18 GB unified memory).

Quick test

A self-contained test that requires no data generation and runs in a few seconds:

python quicktest.py

It loads the six trained ensemble members from ENSEMBLE_RESULTS/, applies them to the two cached corpus instances in examples/, and verifies that the ensemble improves on its classical prior, that the conformal 90% intervals are close to nominal, and that the predicted standard deviation ranks the error. Both examples come from acquisition families (D and E) that never appear in training, so the test also exercises the zero-shot geometry result of the paper.

Expected output:

Loading ensemble members ...
6 ensemble members loaded
  architectures: attn, rescnn, unet
  conformal quantile q_hat(0.9) = 1.0269

instance        fam   shots  recs    prior      ens    gain   cov@90    rho
---------------------------------------------------------------------------
inst_0800.npz   D         4    84    183.5    103.3   43.7%    0.954   0.74
inst_0904.npz   E         7    66    285.5    162.4   43.1%    0.955   0.82

QUICK TEST PASSED: ...

prior and ens are RMSE in m/s against the true model; gain is the reduction achieved by the learned stage; cov@90 is the empirical coverage of the calibrated 90% intervals; rho is the Spearman correlation between the predicted sigma and the absolute error. The test exits with status 0 on success and 1 on failure.

Repository layout

quicktest.py            the quick test described above
examples/               two cached corpus instances (families D and E)
ENSEMBLE_RESULTS/       six trained members + channel normalisation stats
CALIBRATION_RESULTS/    conformal quantiles and calibration metrics
*.py                    pipeline scripts (flat; they import one another)

The scripts are kept in a single flat directory because they import one another as top-level modules and resolve their output directories relative to their own location. Each script writes to its own *_RESULTS/ directory and can be re-run independently; expensive stages checkpoint per instance or per condition and skip completed work on restart.

Reproducing the paper

The stages below are listed in dependency order. Wall-clock figures refer to the laptop described above; the complete matrix is roughly 90 hours.

1. Synthetic corpus (~33 h, two parallel shards)

python build_corpus.py --shard 0 --of 2 &
python build_corpus.py --shard 1 --of 2 &

Writes 1000 instances to ~/DeepFWI_corpus/. Each instance draws a random layered, faulted or salt-bearing model, a random acquisition from six families, random wavelet error and noise, then runs the full classical chain and caches the GIPE tensor. Splits by index: 0–599 train, 600–699 validation, 700–799 calibration, 800–999 test. Families D, E and F never occur below index 600.

2. Ensemble training (~10 min per member on MPS)

python train_ensemble.py --all

3. Conformal calibration and test evaluation (~10 min)

python calibrate_ensemble.py

Produces CALIBRATION_RESULTS/calibration_metrics.json, which every downstream script reads for the conformal quantile.

4. Marmousi-2 chain

Place vp.segy, vs.segy and density.segy in the repository root, then run in order:

python marmousi_prepare.py      # decimate to 10 m and 20 m grids
python marmousi_fd_data.py      # 32-shot marine survey, 6 s records
python marmousi_fwi_pilot.py    # multiscale FWI, 3-12 Hz      (~62 min)
python marmousi_admm_prior.py   # RW-L1/TV refinement          (~25 min)
python marmousi_gipe.py         # GIPE channels                (~10 min)
python marmousi_zeroshot.py     # zero-shot ensemble
python audit_marmousi.py        # held-out-shot physics audit

The Marmousi-2 model is distributed by the Allied Geophysical Laboratories and is not redistributed here; see Martin et al. (2006).

5. Robustness experiments

python fullline_sweep.py              # 11 full-line conditions, ~1.6 h each
python fullline_audit.py              # audit spine over those conditions
python fullline_audit_sensitivity.py  # audit hyperparameter sensitivity
python sweep_crops.py                 # 36 crop-scale conditions
python sweep_shots.py                 # shot-count axis
python sweep_audit.py                 # crop-scale audit spine

6. Baselines and ablations

python fixedgeo_data.py && python fixedgeo_train.py   # fixed-geometry gather net
python srccond_data.py  && python srccond_train.py    # source-conditioned net
python build_corpus_noaug.py                          # no-augmentation corpus
python train_ensemble.py --all --data-dir ~/DeepFWI_corpus_noaug --run-tag _noaug
python calibrate_ensemble.py --suffix _noaug --cal-dir ~/DeepFWI_corpus_noaug
python train_ensemble.py --all --no-ck                # coverage-channel ablation

7. Elastic physics-mismatch experiment

python marmousi_elastic_data.py
python elastic_fwi.py && python elastic_admm.py && python elastic_gipe.py
python elastic_zeroshot.py && python elastic_audit.py

8. Figures and summary

python final_figures.py     # publication layouts
python results_summary.py   # aggregates every metrics file

Note that restyle_figures.py and final_figures.py both write MARMOUSI_AUDIT_RESULTS/Fig_Audit.png. final_figures.py uses the audit selection rule reported in the paper (plateau-right over a 13-point tau grid) and should be run last.

Key results

Reproduced by the pipeline above:

Experiment	Result
Synthetic test set (200 instances)	prior 205 m/s, ensemble 128 m/s (−38%)
Per-family RMSE	111–144 m/s; never-trained D/E/F at 130/122/137 m/s
Zero-shot Marmousi-2, full 17 km line	prior 354 m/s, ensemble 304 m/s
Raw conformal coverage on Marmousi	0.415 at nominal 0.90
After the held-out-shot audit	tau 4.0 (oracle 4.24), coverage 0.888
Fixed-geometry gather baseline	335 m/s at its own geometry, ~495 m/s off it
Source-conditioned gather baseline	377 m/s, worse than the prior on every family

Determinism

Every random draw is seeded deterministically by instance index or condition tag, so any single stage can be reproduced in isolation without re-running the ones before it. Corpus instance i uses seed 10000 + i; sweep conditions seed from a CRC32 of their tag.

Licence

MIT — see LICENSE.

Citation

If you use this code, please cite the accompanying paper. Machine-readable metadata is in CITATION.cff.

Contact

Deepak Kumar — deepak.kumar@igf.edu.pl Institute of Geophysics, Polish Academy of Sciences, Warsaw, Poland