UGP

This repository contains implementations of Finite Basis Physics-Informed Neural Networks (FBPINNs) and Full-Waveform Inversion (FWI) for single source multi-frequency acoustic wave propagation problems in time-domain in 1D and 2D.

Installation

It is recommended to use a virtual environment. This project uses uv for package management.

Installing uv

You can install uv on macOS, Linux, and Windows with the following commands:

• macOS and Linux:

  curl -LsSf https://astral.sh/uv/install.sh | sh

• Windows:

  irm https://astral.sh/uv/install.ps1 | iex

Setting up the Environment

1. Create and activate a virtual environment:

   uv venv
   source .venv/bin/activate

   On Windows, use uv venv followed by .venv\Scripts\activate.

2. Install the required dependencies:

   uv pip install -e .

Running the Scripts

You can run the forward models and inversion problems using the following commands :

• In homogeneous medium - water : uv run FM_Homogeneous_2D.py uv run FWI_Homogeneous_2D.py
• In heterogeneous medium - water and sediments : uv run FM_Heterogeneous_2D.py uv run FWI_Heterogeneous_2D.py
• In water with anomaly : uv run FM_Homogeneous_Anomaly_2D.py uv run FWI_Homogeneous_Anomaly_2D.py
• In water-sediment medium with anomaly : uv run FM_Heterogeneous_Anomaly_2D.py uv run FWI_Heterogeneous_Anomaly_2D.py

The scripts will save results, including plots and model checkpoints, in the results directory.

C&G Reviewer Smoke Test

For a quick reproducibility check that does not rerun neural training, use:

make cg-smoke

This rebuilds verification_artifacts/cg_minimal_benchmark_table.csv from the retained minimal CPML benchmark metrics, checks the expected three-seed metric ranges in verification_artifacts/minimal_benchmark.json, cross-checks the aggregate ledger against the retained per-seed ledgers, and verifies SHA-256 fingerprints for the contract, metrics, and audit figure listed in verification_artifacts/cg_repro_manifest.json. It should run in seconds to a few minutes on a laptop and requires only Python's standard library. A passing run reports hashes_checked: 7 and regenerated table hash 18eb4d71791389134566ba252578020700c44cbf925244af30d48a32c9cd0f75.

Running on GPUs

By default the script will use CPU, but if you have NVIDIA GPUs, then run :

• nvcc --version and get compilation tool verisons (useful if you are compiling custom CUDA extensions)
• nvidia-smi and get NVIDIA driver, CUDA runtime and GPU

Uninstall the jaxlib uv pip uninstall jaxlib and install jaxlib wheel for your CUDA version : uv pip install --upgrade "jax==0.4.25" "jaxlib==0.4.25+cuda11.cudnn86" "numpy<2.0" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

Verify CUDA libraries exist and are readable :

ls /usr/local/cuda-11.8/lib64/libcudart.so*
ls /usr/local/cuda-11.8/lib64/libcublas.so*

Managing cuDNN in custom user-space in root directory :

wget --no-check-certificate \
  https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/cudnn-linux-x86_64-8.9.7.29_cuda11-archive.tar.xz

# Inside your home directory, create include/ and lib folders
mkdir -p ~/local/cudnn/include ~/local/cudnn/lib

# Extract only include/ and lib/ contents into ~/.local/cudnn
tar -xJf cudnn-linux-x86_64-*-cuda11-archive.tar.xz \
  --strip-components=1 \
  -C ~/local/cudnn

Export these :

export CUDA_HOME=/usr/local/cuda-11.8
export PATH=$CUDA_HOME/bin:$PATH
export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$HOME/local/cudnn/lib:$LD_LIBRARY_PATH
export CPATH=$HOME/local/cudnn/include:$CPATH

Files in your user-space cuDNN directory : ls -l ~/local/cudnn/lib/libcudnn* you should get something like this :

libcudnn.so -> libcudnn.so.8
libcudnn.so.8 -> libcudnn.so.8.9.1.23
libcudnn.so.8.9.1.23

Test via uv :

uv run python - << 'EOF'
import jax, os
print("nvcc:", os.popen("nvcc --version").read().splitlines()[-1])
print("cuDNN lib loaded:", os.popen("ldd $(which python) | grep cudnn").read())
print("JAX backend:", jax.default_backend())
EOF

Implementation Details

This project uses a slightly modified and extended version of the JAX-based FBPINNs library developed by Ben Moseley, available at https://github.com/benmoseley/FBPINNs. We have adapted the core framework to implement advanced network architectures, adaptive loss weighting schemes, and specialized problem definitions for acoustic wave modeling.

References

• Overview of FWI in Geophysics
• Physics-informed machine learning: from concepts to real-world applications
• On use of Neural Networks for Full Waveform Inversion
• Solving 2-D Helmholtz equation in the rectangular, circular, and elliptical domains using neural networks
• Exploring PINNs
• PINNs for Wave Propagation and Full Waveform Inversion
• Physics-informed neural network for seismic wave inversion in layered semi-infinite domain
• Physics-guided FWI via Encoder-Solver CNNs
• Xavier Initialization
• Is it time to Swish
• Training PINNs with Hard Constraints and Adaptive Weights
• Respecting Casulaity is All you need
• Stability Analysis of PINNs
• Improved PINNSs loss function regularization with a variance-based term
• Learning Overalap Optimization for Domain Decomposition Method
• Improved PINN in mitigating gradient- related failures
• Physics-informed Neural Networks with Fourier Features for Seismic Wavefield Simulation in Time-Domain Nonsmooth Complex Media
• ShaVi
• BPINNs for forward and inverse PDE problems with noisy data
• BPINNs for real-world nonlinear dynamical systems
• Self Adaptive PINNs
• FBPINNs
• Multifidelity domain decomposition-based PINNs and operators for time-dependent problems
• Multilevel domain decomposition-based architectures for PINNs
• ELF-FBPINNs
• Swish
• NAS-PINNs and NAS-PINNv2
• Auto-PINN
• PINN-DARTs
• NAS: A Survey
• A Survey on Computationally Efficient NAS
• Quantitative Seismology - Aki, K., & Richards, P. G. (2002). Quantitative Seismology (2nd ed.). University Science Books.
• Seismic Data Analysis - Yilmaz, Ö. (2001). Seismic Data Analysis (2nd ed.). Society of Exploration Geophysicists.
• Full-waveform inversion in exploration geophysics - Virieux, J., & Operto, S. (2009). An overview of full-waveform inversion in exploration geophystheics. Geophysics, 74(6), WCC1-WCC26.
• Inversion of seismic reflection data - Tarantola, A. (1984). Inversion of seismic reflection data in the acoustic approximation. Geophysics, 49(8), 1259-1266.
• Double Descent
• The Lottery Ticket Hypothesis - think about pruning the PINN
• PIHLoRA
• Transfer Learning in PINNs