Metadata-Version: 2.1
Name: sweep-solver
Version: 0.3.1
Summary: Seismic Wave Equation Exploration Platform — equations, propagators, operators.
Author-email: Shaowen Wang <shaowen.wang@kaust.edu.sa>
License: MIT
Project-URL: Homepage, https://github.com/DeepWave-KAUST/sweep
Project-URL: Issues, https://github.com/DeepWave-KAUST/sweep/issues
Keywords: fwi,geophysics,seismic,wave-equation,propagator,cuda,pytorch,jax
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: POSIX :: Linux
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Physics
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy>=1.20
Requires-Dist: scipy>=1.7
Requires-Dist: torch
Requires-Dist: ninja
Provides-Extra: cuda
Requires-Dist: torch; extra == "cuda"
Provides-Extra: cuda12
Requires-Dist: nvidia-cufft-cu12; extra == "cuda12"
Provides-Extra: cuda13
Requires-Dist: nvidia-cufft<13,>=12; extra == "cuda13"
Provides-Extra: torch
Requires-Dist: torch; extra == "torch"
Provides-Extra: jax
Requires-Dist: jax; extra == "jax"
Requires-Dist: jaxlib; extra == "jax"
Provides-Extra: all
Requires-Dist: torch; extra == "all"
Requires-Dist: jax; extra == "all"
Requires-Dist: jaxlib; extra == "all"
Provides-Extra: datasets
Requires-Dist: requests>=2.25; extra == "datasets"
Requires-Dist: tqdm>=4.60; extra == "datasets"

<p align="center">
  <picture>
    <source media="(prefers-color-scheme: dark)" srcset="docs/assets/logo/sweep-icon-dark.svg">
    <img src="docs/assets/logo/sweep-icon-light.svg" alt="SWEEP" width="180">
  </picture>
</p>

<h1 align="center">SWEEP</h1>

<p align="center">
  <a href="https://sweepx.deepwave.group/solver/"><img alt="Docs" src="https://img.shields.io/badge/docs-online-blue?logo=readthedocs&logoColor=white"></a>
  <a href="https://opensource.org/licenses/MIT"><img alt="License: MIT" src="https://img.shields.io/badge/License-MIT-yellow.svg"></a>
  <a href="https://pytorch.org"><img alt="PyTorch" src="https://img.shields.io/badge/PyTorch-2.0%2B-EE4C2C?logo=pytorch&logoColor=white"></a>
</p>

<p align="center">English | <a href="README.zh-CN.md">中文</a></p>

**Seismic Wave Equation Exploration Platform** — a differentiable framework for seismic wave-equation modeling, migration, and full-waveform inversion. One API, 20+ equations (acoustic / elastic / VTI / TTI / DAS), PyTorch and JAX backends, eager and compiled CUDA paths.

📖 **Documentation**: <https://sweepx.deepwave.group/solver/>

## Install

```bash
pip install sweepx
```

The wheel ships prebuilt CUDA cores, so the compiled backend (`impl='c'`) works right
after install with any PyTorch version: no nvcc, no compiler, no build step. It needs
an NVIDIA GPU and driver; the core for your torch's CUDA version is picked
automatically (CUDA 12: V100 and newer; CUDA 13: T4 and newer, driver >= 580). The
eager PyTorch and JAX backends are pure Python.

**From source**: a clone has no prebuilt core, so the CUDA core is compiled locally
once (with an nvcc matching your torch's CUDA version) and cached:

```bash
pip install .
python -m sweep.build   # optional: compile it now instead of on the first impl='c' call
```

<sub>`sweepx` (Python >= 3.10) is the PyPI name and also installs the `sweep-agent` companion;
you `import sweep`. `pip install sweep-solver` installs the solver alone (Python >= 3.9).
GPU coverage, custom cores and developer builds are covered in the
[installation guide](https://sweepx.deepwave.group/solver/getting-started/installation/).</sub>

## Hello SWEEP

One shot, one receiver, one `.backward()` — read off the velocity-model gradient for a single trace:

```python
import numpy as np
import torch
from sweep.equations import Acoustic
from sweep.propagator.torch import PropTorch
from sweep.signal import ricker

shape = (96, 128)
dh, dt, nt = 10.0, 0.002, 800
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

vp_true = np.full(shape, 1500.0, dtype=np.float32)
vp_true[shape[0] // 2:, :] = 2500.0
vp_init = np.full(shape, 1500.0, dtype=np.float32)

solver = PropTorch(Acoustic(device=device), shape=shape, dh=dh, dt=dt,
                   device=device, pml_type="cpmlr", use_ckpt=False)

t = np.arange(nt) * dt
wavelet = ricker(t - 0.14, f=10.0).astype(np.float32)
sources   = np.array([[shape[1] // 4, shape[0] // 2]], dtype=np.int64)
receivers = np.array([[[3 * shape[1] // 4, shape[0] // 2]]], dtype=np.int64)

with torch.no_grad():
    obs = solver(wavelet, sources, receivers, models=[torch.tensor(vp_true, device=device)])

vp_t = torch.tensor(vp_init, device=device, requires_grad=True)
pred = solver(wavelet, sources, receivers, models=[vp_t])
(0.5 * (pred - obs).pow(2).sum()).backward()

print("vp gradient shape:", tuple(vp_t.grad.shape))
```

Swap `Acoustic` for `Elastic`, `AcousticVTI`, `ElasticTTI`, ... — the surrounding code is unchanged.

## Notebooks & examples

- **Hello SWEEP** — forward / backward / 5-line FWI loop: [`docs/notebooks/00_hello_fwi.ipynb`](docs/notebooks/00_hello_fwi.ipynb)
- **FWI on Marmousi** (acoustic / elastic / multiscale): see [`docs/notebooks/01_*`–`03_*`](docs/notebooks/)
- **Wavefields, DAS, anisotropic, RTM**: [`docs/notebooks/04_*`–`08_*`](docs/notebooks/)
- **Production scripts** (multi-GPU, MPI shot parallelism, multi-shot batching): under [`examples/`](examples/)

## Citing

```bibtex
@misc{wang2026sweep,
  title  = {{SWEEP} ({S}eismic {W}ave {E}quation {E}xploration {P}latform):
            A Unified Solver Framework for Differentiable Wave Physics},
  author = {Wang, Shaowen and Alkhalifah, Tariq},
  year   = {2026},
  eprint = {2604.14189},
  archivePrefix = {arXiv},
  url    = {https://arxiv.org/abs/2604.14189},
}
```

## License

MIT — see [LICENSE](LICENSE).
