SWEEP (Seismic Wave Equation Exploration Platform): A Unified Solver Framework for Differentiable Wave Physics

SWEEP is a unified, extensible, and differentiable wave equation solver framework that supports diverse seismic propagation models and plug-and-play components to facilitate advanced gradient-based inversion tasks like full-waveform inversion and least-squares reverse time migration.

The Gist

Imagine you are trying to figure out what's inside a giant, opaque cake without cutting it open. You can't see the layers of chocolate, fruit, or sponge, but you can tap the top with a spoon and listen to the sound it makes. By analyzing how the sound waves bounce around inside, you can build a mental map of the cake's interior.

This is essentially what seismic exploration does, but instead of a cake, we are looking at the Earth's crust to find oil, gas, or understand earthquakes. Scientists send sound waves (seismic waves) underground and record how they bounce back. The challenge is: How do we turn those messy echoes into a clear picture of what's underground?

This is where the paper introduces SWEEP (Seismic Wave Equation Exploration Platform). Think of SWEEP as a "Universal Simulator and Detective Tool" for these underground sound waves.

Here is a breakdown of what SWEEP does, using simple analogies:

1. The Problem: The "Manual Math" Nightmare

Traditionally, to solve these seismic puzzles, scientists had to act like old-school accountants. They had to manually write out complex mathematical formulas (called "adjoint equations") to figure out how to reverse-engineer the underground structure from the sound waves.

• The Analogy: Imagine trying to bake a cake by hand-writing a new recipe every time you want to change the flavor. It's slow, prone to errors, and if you make a tiny mistake in the math, your whole cake (or simulation) fails.

2. The Solution: The "Smart, Self-Correcting" Engine

SWEEP changes the game by using a technology called Automatic Differentiation (AD).

• The Analogy: Instead of manually writing the recipe, SWEEP is like a smart kitchen robot. You tell it, "I want a chocolate cake," and it automatically figures out the exact math needed to get there. If you want to tweak the recipe (change the model), the robot instantly recalculates the steps without you needing to rewrite the whole manual.
• Why it matters: This makes it incredibly easy to use advanced techniques like Full-Waveform Inversion (FWI), which is basically "training" a computer to guess the underground structure by constantly comparing its guesses to real data until it gets it perfect.

3. The "Lego" Architecture (Plug-and-Play)

One of SWEEP's coolest features is its modular design.

• The Analogy: Think of SWEEP as a high-tech Lego set.
  • The Propagator: This is the base plate that moves the waves forward in time.
  • The Source: This is the hand that drops the "pebble" (the sound wave) into the water.
  • The Receiver: This is the net that catches the ripples.
  • The Physics: You can swap out the "physics block" just like swapping a Lego piece. Want to simulate sound in water? Click one block. Want to simulate sound in rock that absorbs energy? Click another. Want to simulate complex rock layers that behave differently in different directions (anisotropy)? Click a third.
• The Benefit: Scientists don't have to rebuild the whole engine every time they want to test a new theory. They just swap the pieces.

4. Speed and Scale: The "Conveyor Belt"

Simulating these waves is computationally heavy. It's like trying to run a marathon while carrying a piano.

• The Analogy: SWEEP is designed to run on super-fast conveyor belts (GPUs and TPUs).
  • Batch Modeling: Instead of simulating one earthquake at a time, SWEEP can simulate 100 earthquakes happening simultaneously on a single conveyor belt. This is like a factory that used to make one car at a time but now has an assembly line that builds 100 cars at once.
  • Multi-GPU: It can spread the work across many computers, like a team of 100 people painting a mural instead of just one person.

5. What Can It Actually Do?

The paper lists many specific "flavors" of physics that SWEEP can handle:

• Acoustic: Simple sound waves (like in water).
• Elastic: Waves that shake and stretch the ground (like in solid rock).
• Viscoacoustic: Waves that lose energy as they travel (like sound getting muffled in a thick fog).
• Anisotropic (VTI/TTI): Waves that travel faster in some directions than others (like light through a prism).
• LSRTM: A super-sharp imaging technique that removes the "blur" from the underground picture.

The Big Picture

Before SWEEP, building a seismic simulator was like building a car from scratch every time you wanted to drive to a new city. You had to forge the metal, design the engine, and tune the brakes manually.

SWEEP is the "Tesla" of seismic simulation. It provides a unified, high-performance, and flexible platform where scientists can focus on the science (solving the mystery of the Earth) rather than the math (worrying about the engine mechanics). It allows researchers to plug in new ideas, run them at lightning speed on powerful computers, and get clearer, more accurate pictures of what lies beneath our feet.

Technical Summary

1. Problem Statement

Seismic exploration relies heavily on solving wave equations for velocity model building and imaging. Traditional workflows for inverse problems (such as Full-Waveform Inversion - FWI, and Least-Squares Reverse Time Migration - LSRTM) require the manual derivation and implementation of adjoint wave equations and gradient expressions using the adjoint-state method. This process is:

• Complex and Error-Prone: It demands rigorous attention to temporal consistency, boundary conditions, and source terms.
• Inflexible: Maintaining code quality and correctness across different physical models (acoustic, elastic, anisotropic, attenuative) is difficult.
• Performance vs. Flexibility Trade-off: Existing Automatic Differentiation (AD) tools face a dilemma. Frameworks like PyTorch offer high development efficiency but often suffer from lower runtime performance. Conversely, CUDA-based implementations offer high performance but lack development flexibility and are difficult to maintain.

2. Methodology

The authors propose SWEEP, a unified and extensible wave equation solver library that bridges the gap between ease of use and high performance by leveraging JAX and PyTorch.

Core Architecture

• Unified Backend: SWEEP supports both PyTorch and JAX. JAX is highlighted for its ability to combine AD capabilities with Just-In-Time (JIT) compilation via the XLA compiler, achieving performance comparable to hand-written CUDA code while retaining Pythonic flexibility.
• Modular Design: The solver is decomposed into three primary modules:
  1. Propagator: Handles time-stepping, memory allocation, and wavefield updates.
  2. Source: Manages wavefield injection with flexible configurations.
  3. Receiver: Records wavefields at specific locations/components.
• RNN-like Abstraction: The time-stepping process is abstracted to resemble a Recurrent Neural Network (RNN). While the propagation is not performed by neural networks, this structural analogy provides a unified interface for integrating various wave equations and facilitates seamless integration with AD frameworks.
• User Interface: To implement a new wave equation, users only need to define the single time-step update behavior (the step function). The framework automatically infers model shapes, handles memory allocation, and orchestrates the full simulation.

Advanced Features

• Batch Modeling: Supports simultaneous computation of multiple shots (sources) and multiple velocity models in a single kernel launch, significantly reducing overhead compared to sequential processing.
• High-Performance Computing (HPC): Designed for GPUs and TPUs. It supports multi-device execution via DistributedDataParallel (PyTorch) or pmap (JAX), enabling efficient shot-based parallelism.
• Differentiability: Native support for automatic differentiation allows for the seamless implementation of gradient-based optimization methods (FWI, LSRTM) without manual adjoint derivation.

3. Key Contributions

• Unified Solver Framework: A single codebase that supports a wide spectrum of wave physics, including:
  • Acoustic (1st and 2nd order)
  • Elastic (Velocity-Stress)
  • Viscoacoustic (Attenuative)
  • Anisotropic (VTI, TTI, Decoupled VTI/TTI)
  • Coupled Acoustic-Elastic
  • Pseudo-Elastic
• Differentiable Physics: The framework enables the direct computation of gradients with respect to model parameters, facilitating advanced inversion techniques like FWI and LSRTM without manual adjoint coding.
• Extensibility: The "plug-and-play" architecture allows for easy integration of custom loss functions, neural network-based velocity representations, and multi-scale inversion strategies.
• Performance Optimization: By utilizing JAX's XLA compiler, the framework achieves near-native GPU performance while maintaining the rapid development cycle of Python.

4. Results and Validation

The paper validates the framework through extensive numerical examples demonstrating the solver's capability across various physical scenarios:

• Acoustic & Viscoacoustic: Successfully modeled attenuation and dispersion effects in viscoacoustic media, showing clear differences between phase-shifted and amplitude-damped wavefields.
• Anisotropy: Demonstrated the generation of wavefields for VTI and TTI media with varying anisotropic parameters ($\eta, \epsilon, \delta$). The results correctly reproduced known artifacts (e.g., diamond patterns in VTI) and the effectiveness of decoupled equations in suppressing degenerate qSV artifacts.
• Elastic & Coupled: Validated the elastic wave equation and the Acoustic-Elastic Coupled (AEC) equation in multi-layer models, correctly capturing P-wave and S-wave interactions.
• LSRTM: Formulated the linearized Born approximations for both acoustic and elastic isotropic media, as well as the acoustic-elastic coupled case, providing the mathematical foundation for least-squares migration within the framework.
• Pseudo-Elastic: Showed that the pseudo-elastic equation can generate wavefields similar to acoustic equations in fluid-like layers while avoiding dispersion issues found in standard elastic solvers with low S-wave velocities.

5. Significance

SWEEP represents a significant advancement in computational geophysics by democratizing access to high-performance, differentiable wave physics.

• Scientific Impact: It lowers the barrier to entry for developing complex inversion workflows, allowing researchers to focus on high-level scientific challenges (e.g., designing better loss functions or neural architectures) rather than low-level numerical implementation.
• Technological Impact: It resolves the long-standing trade-off between development speed and computational efficiency in seismic modeling. By unifying PyTorch and JAX, it future-proofs seismic solvers for emerging hardware (TPUs) and AI-driven workflows.
• Community Value: As an open repository with a clean, concise codebase, it serves as a standardized platform for the community to share, compare, and extend wave equation solvers, accelerating innovation in seismic imaging and inversion.