sweep-solver 0.1.0__tar.gz

sweep_solver-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to do so, subject to the
+following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

sweep_solver-0.1.0/MANIFEST.in

@@ -0,0 +1,5 @@
+include LICENSE
+include README.md
+include build_config.py
+include setup_cuda.py
+recursive-include src/sweep/csrc *.cpp *.cu *.cuh *.h CMakeLists.txt

sweep_solver-0.1.0/PKG-INFO

@@ -0,0 +1,160 @@
+Metadata-Version: 2.1
+Name: sweep-solver
+Version: 0.1.0
+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: 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://deepwave-kaust.github.io/sweep/"><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://deepwave-kaust.github.io/sweep/>
+
+## Install
+
+**From PyPI** — one wheel, any PyTorch version, any Python 3:
+
+```bash
+pip install sweepx
+python -c "import sweep; sweep.precompile()"   # build the CUDA backend now (one-time ~3–5 min)
+```
+
+`sweepx` ships the C++/CUDA *sources*; the compiled backend (`impl='c'`) is compiled
+against **your** torch — only for your GPU's architecture, then cached in
+`~/.cache/torch_extensions`. The `precompile()` line does it up front; drop it and it
+happens automatically on first use of `impl='c'`. No torch/CUDA version lock-in. Needs
+a CUDA GPU + `nvcc >= 12.4` (a system install, your cluster's `module load cuda`, or
+`conda install -c nvidia cuda-toolkit`); the pure-Python **eager** / **JAX** backends
+work without nvcc.
+
+**From source** (a clone):
+
+```bash
+# pure-Python (PyTorch / JAX eager path); impl='c' JIT-compiles on first use
+pip install .
+
+# prebuild the C++/CUDA extension now — skips the first-use compile (needs nvcc)
+SWEEP_BUILD_CUDA=1 pip install -v ".[cuda]" --no-build-isolation
+```
+
+If the prebuild can't auto-detect your GPU, set `TORCH_CUDA_ARCH_LIST` (e.g. `"7.0"`
+V100, `"8.0"` A100, `"8.9"` RTX 6000 Ada) before the second command.
+
+<sub>`sweepx` is the PyPI distribution name; you `import sweep` (the `scikit-learn` → `import sklearn`
+pattern, because the bare name `sweep` is taken on PyPI). `pip install sweep-solver` is equivalent.
+Full install notes are in [the docs](https://deepwave-kaust.github.io/sweep/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,
+                   dev=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: [`examples/notebooks/00_hello_fwi.ipynb`](examples/notebooks/00_hello_fwi.ipynb)
+- **FWI on Marmousi** (acoustic / elastic / multiscale): see [`examples/notebooks/01_*`–`03_*`](examples/notebooks/)
+- **Wavefields, DAS, anisotropic, RTM**: [`examples/notebooks/04_*`–`08_*`](examples/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).

sweep_solver-0.1.0/README.md

@@ -0,0 +1,118 @@
+<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://deepwave-kaust.github.io/sweep/"><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://deepwave-kaust.github.io/sweep/>
+
+## Install
+
+**From PyPI** — one wheel, any PyTorch version, any Python 3:
+
+```bash
+pip install sweepx
+python -c "import sweep; sweep.precompile()"   # build the CUDA backend now (one-time ~3–5 min)
+```
+
+`sweepx` ships the C++/CUDA *sources*; the compiled backend (`impl='c'`) is compiled
+against **your** torch — only for your GPU's architecture, then cached in
+`~/.cache/torch_extensions`. The `precompile()` line does it up front; drop it and it
+happens automatically on first use of `impl='c'`. No torch/CUDA version lock-in. Needs
+a CUDA GPU + `nvcc >= 12.4` (a system install, your cluster's `module load cuda`, or
+`conda install -c nvidia cuda-toolkit`); the pure-Python **eager** / **JAX** backends
+work without nvcc.
+
+**From source** (a clone):
+
+```bash
+# pure-Python (PyTorch / JAX eager path); impl='c' JIT-compiles on first use
+pip install .
+
+# prebuild the C++/CUDA extension now — skips the first-use compile (needs nvcc)
+SWEEP_BUILD_CUDA=1 pip install -v ".[cuda]" --no-build-isolation
+```
+
+If the prebuild can't auto-detect your GPU, set `TORCH_CUDA_ARCH_LIST` (e.g. `"7.0"`
+V100, `"8.0"` A100, `"8.9"` RTX 6000 Ada) before the second command.
+
+<sub>`sweepx` is the PyPI distribution name; you `import sweep` (the `scikit-learn` → `import sklearn`
+pattern, because the bare name `sweep` is taken on PyPI). `pip install sweep-solver` is equivalent.
+Full install notes are in [the docs](https://deepwave-kaust.github.io/sweep/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,
+                   dev=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: [`examples/notebooks/00_hello_fwi.ipynb`](examples/notebooks/00_hello_fwi.ipynb)
+- **FWI on Marmousi** (acoustic / elastic / multiscale): see [`examples/notebooks/01_*`–`03_*`](examples/notebooks/)
+- **Wavefields, DAS, anisotropic, RTM**: [`examples/notebooks/04_*`–`08_*`](examples/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).

sweep_solver-0.1.0/build_config.py

@@ -0,0 +1,258 @@
+import glob
+import inspect
+import os
+import sys
+from distutils import log
+
+try:
+    import packaging.utils as packaging_utils
+except ImportError:
+    packaging_utils = None
+
+try:
+    import setuptools._core_metadata as setuptools_core_metadata
+except ImportError:
+    setuptools_core_metadata = None
+
+
+ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
+PACKAGE_VERSION = "0.1.0"
+
+
+def env_flag_enabled(name):
+    value = os.environ.get(name, "")
+    return value.strip().lower() in {"1", "true", "yes", "on"}
+
+
+def openmp_flags():
+    if sys.platform == "win32":
+        return ["/openmp"]
+    if sys.platform == "darwin":
+        return []
+    return ["-fopenmp"]
+
+
+def configure_cuda_arch_list():
+    """Avoid PyTorch's empty GPU-arch auto-detection on login/CPU nodes."""
+    if os.environ.get("TORCH_CUDA_ARCH_LIST"):
+        return
+
+    arch_list = os.environ.get("SWEEP_CUDA_ARCH_LIST", "7.0")
+    os.environ["TORCH_CUDA_ARCH_LIST"] = arch_list
+    log.warn(
+        "TORCH_CUDA_ARCH_LIST is not set; defaulting to %s. "
+        "Set TORCH_CUDA_ARCH_LIST or SWEEP_CUDA_ARCH_LIST to target other GPUs.",
+        arch_list,
+    )
+
+
+def is_metadata_only_invocation():
+    metadata_commands = {"egg_info", "dist_info", "prepare_metadata_for_build_wheel"}
+    return any(arg in metadata_commands for arg in sys.argv[1:])
+
+
+def patch_packaging_compat():
+    if packaging_utils is None:
+        return
+
+    signature = inspect.signature(packaging_utils.canonicalize_version)
+    if "strip_trailing_zero" in signature.parameters:
+        return
+
+    original = packaging_utils.canonicalize_version
+
+    def canonicalize_version_compat(version, strip_trailing_zero=True):
+        return original(version)
+
+    packaging_utils.canonicalize_version = canonicalize_version_compat
+    if setuptools_core_metadata is not None:
+        setuptools_core_metadata.canonicalize_version = canonicalize_version_compat
+
+
+def get_sources():
+    """Collect C++/CUDA sources for the ``sweep._C`` extension.
+
+    Honours the ``SWEEP_SKIP_CPU`` environment variable: when set to a
+    truthy value (1, true, yes, on), the heavy ``cpu/equations/*`` tree
+    (~19k lines, often the build-time bottleneck) is *excluded* and a tiny
+    stub is linked in its place.  The stub keeps `bindings/module.cpp`
+    linking and routes every call to the CUDA path; attempting to use a
+    CPU tensor raises a clear TORCH_CHECK message.
+
+    This is intended for users who only ever run on CUDA — typically HPC
+    deployments where the CPU C++ path would be dead weight.
+    """
+    cuda_sources = (
+        glob.glob("src/sweep/csrc/cuda/common/**/*.cu", recursive=True)
+        + glob.glob("src/sweep/csrc/cuda/equations/**/*.cu", recursive=True)
+    )
+    binding_sources = ["src/sweep/csrc/bindings/module.cpp"]
+
+    if env_flag_enabled("SWEEP_SKIP_CPU"):
+        log.warn(
+            "SWEEP_SKIP_CPU=1: skipping cpu/equations/* (~19k LoC); linking "
+            "cpu_binding_stub.cpp instead. CPU tensors will raise a clear "
+            "error at call time."
+        )
+        cpu_sources = ["src/sweep/csrc/cpu/cpu_binding_stub.cpp"]
+    else:
+        cpu_sources = glob.glob("src/sweep/csrc/cpu/**/*.cpp", recursive=True)
+        # Defensive: don't accidentally include the stub if it's globbed
+        cpu_sources = [
+            s for s in cpu_sources
+            if not s.endswith("cpu_binding_stub.cpp")
+        ]
+
+    return cpu_sources + cuda_sources + binding_sources
+
+
+def _check_ninja_on_path():
+    """Print an actionable note if torch's ninja-binary probe will fail.
+
+    Torch's ``is_ninja_available()`` shells out to ``ninja --version`` on
+    PATH, NOT to the bundled Python ``ninja`` package.  If the conda env is
+    not activated (e.g. invoking ``/path/to/envs/X/bin/python setup.py``
+    directly), the ninja binary at ``<env>/bin/ninja`` is invisible and
+    torch silently falls back to the *slow* distutils sequential build —
+    ~6× slower in practice.  Loud-warn now, instead of having the user
+    discover it 20 minutes into a serial compile.
+    """
+    import shutil
+
+    if shutil.which("ninja") is not None:
+        return  # binary on PATH — torch will use ninja, all good
+
+    # Bundled ninja package?  Tell the user how to expose it.
+    bundled = None
+    try:
+        import ninja as _ninja_pkg
+        bundled = os.path.join(os.path.dirname(_ninja_pkg.__file__), "..", "..", "..", "..", "bin", "ninja")
+        bundled = os.path.normpath(bundled)
+        if not os.path.exists(bundled):
+            # Try the conda env layout
+            python_exec = sys.executable
+            env_bin = os.path.dirname(python_exec)
+            candidate = os.path.join(env_bin, "ninja")
+            bundled = candidate if os.path.exists(candidate) else None
+    except ImportError:
+        pass
+
+    msg = (
+        "ninja binary not found on PATH.  Torch will fall back to the slow "
+        "distutils sequential build (~6× slower).  "
+    )
+    if bundled and os.path.exists(bundled):
+        msg += (
+            f"A ninja binary is bundled at {bundled}; either activate your "
+            "conda env (`conda activate <env>`) so PATH includes it, or run "
+            f"PATH='{os.path.dirname(bundled)}:$PATH' python setup.py ..."
+        )
+    else:
+        msg += "Install ninja-build (`apt install ninja-build` or `pip install ninja`)."
+    log.warn(msg)
+    print(f"WARNING: {msg}", file=sys.stderr, flush=True)
+
+
+def make_build_extension(BuildExtension):
+    def emit(message):
+        print(message, file=sys.stderr, flush=True)
+        log.info(message)
+
+    class SweepBuildExtension(BuildExtension):
+        def run(self):
+            self.verbose = max(getattr(self, "verbose", 1), 2)
+            _check_ninja_on_path()
+            for ext in self.extensions:
+                sources = list(getattr(ext, "sources", []))
+                emit(f"Building CUDA extension '{ext.name}' with {len(sources)} source files")
+                for index, source in enumerate(sources, start=1):
+                    emit(f"  [{index}/{len(sources)}] {source}")
+            super().run()
+
+        def build_extensions(self):
+            self.verbose = max(getattr(self, "verbose", 1), 2)
+            emit("Starting C++/CUDA compilation")
+            super().build_extensions()
+            emit("Finished C++/CUDA compilation")
+
+    return SweepBuildExtension
+
+
+def build_ext_kwargs(build_cuda=None):
+    """Return setup() kwargs for the optional AOT C++/CUDA extension.
+
+    The default distribution is **JIT** (see ``sweep/_jit.py``): one ``py3-none``
+    wheel ships the C++/CUDA sources and compiles ``sweep._C`` against the user's
+    own torch on first use — so this returns NO ``ext_modules`` and every dep
+    comes from ``pyproject.toml``. The ``SWEEP_BUILD_CUDA=1`` path is kept only
+    for building optional pre-compiled fast-path wheels (e.g. a GitHub release),
+    never for the PyPI wheel.
+    """
+    if build_cuda is None:
+        build_cuda = env_flag_enabled("SWEEP_BUILD_CUDA")
+
+    kwargs = {
+        "ext_modules": [],
+        "cmdclass": {},
+    }
+
+    if not build_cuda:
+        return kwargs
+
+    try:
+        from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+    except ImportError as exc:
+        if is_metadata_only_invocation():
+            log.warn(
+                "Skipping CUDA extension setup during metadata generation because PyTorch "
+                "is not installed in the current build environment."
+            )
+            return kwargs
+
+        raise RuntimeError(
+            "Building sweep with SWEEP_BUILD_CUDA=1 requires PyTorch to be installed first, "
+            "because the CUDA extension uses torch.utils.cpp_extension. "
+            "In a pure JAX environment, install without SWEEP_BUILD_CUDA or install PyTorch "
+            "before building the CUDA extension."
+        ) from exc
+
+    SweepBuildExtension = make_build_extension(BuildExtension)
+    omp_flags = openmp_flags()
+    configure_cuda_arch_list()
+
+    # Optional extra nvcc flags (e.g. -DELASTIC3D_LB_MINBLOCKS=6 to retune a
+    # forward launch_bounds without editing kernel source).  Space-separated.
+    extra_nvcc = os.environ.get("SWEEP_EXTRA_NVCC", "").split()
+
+    kwargs["ext_modules"] = [
+        CUDAExtension(
+            name="sweep._C",
+            sources=get_sources(),
+            include_dirs=[
+                os.path.join(ROOT_DIR, "src/sweep/csrc"),
+                os.path.join(ROOT_DIR, "src/sweep/csrc/bindings"),
+                os.path.join(ROOT_DIR, "src/sweep/csrc/shared"),
+                os.path.join(ROOT_DIR, "src/sweep/csrc/cuda"),
+                os.path.join(ROOT_DIR, "src/sweep/csrc/cuda/common"),
+                os.path.join(ROOT_DIR, "src/sweep/csrc/cuda/equations"),
+            ],
+            extra_compile_args={
+                "cxx": ["-O3", "-Wno-attributes", *omp_flags],
+                "nvcc": [
+                    "-O3",
+                    "--use_fast_math",
+                    "--threads=16",
+                    "-Xcompiler=-Wno-deprecated-declarations",
+                    *extra_nvcc,
+                ],
+            },
+            # RPATH so the shipped wheel resolves libtorch/libc10 against the
+            # USER's torch (auditwheel --exclude keeps those libs external).
+            # Belt-and-suspenders: sweep always imports torch before sweep._C.
+            extra_link_args=[*omp_flags, "-Wl,-rpath,$ORIGIN/../torch/lib"],
+        )
+    ]
+    kwargs["cmdclass"] = {
+        "build_ext": SweepBuildExtension.with_options(use_ninja=True)
+    }
+    return kwargs

sweep_solver-0.1.0/pyproject.toml

@@ -0,0 +1,74 @@
+[build-system]
+requires = ["setuptools>=64", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "sweep-solver"
+version = "0.1.0"
+description = "Seismic Wave Equation Exploration Platform — equations, propagators, operators."
+readme = "README.md"
+requires-python = ">=3.9"
+license = {text = "MIT"}
+authors = [
+    {name = "Shaowen Wang", email = "shaowen.wang@kaust.edu.sa"},
+]
+keywords = ["fwi", "geophysics", "seismic", "wave-equation", "propagator", "cuda", "pytorch", "jax"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: POSIX :: Linux",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Scientific/Engineering :: Physics",
+]
+# JIT distribution: one py3-none wheel ships the C++/CUDA *sources* and compiles
+# sweep._C against the user's OWN torch on first use (see sweep/_jit.py), so a
+# single wheel works with any torch version + any Python 3. torch is unpinned.
+# The first-use compile needs nvcc from a CUDA toolkit (system / `module load
+# cuda` / `conda install -c nvidia cuda-nvcc`) — the pip `nvidia-cuda-nvcc-cu12`
+# wheel ships only ptxas, not the nvcc frontend, so it can't be relied on.
+dependencies = [
+    "numpy>=1.20",
+    "scipy>=1.7",
+    "torch",
+    "ninja",
+]
+
+[project.optional-dependencies]
+cuda  = ["torch"]
+torch = ["torch"]
+jax   = ["jax", "jaxlib"]
+all   = ["torch", "jax", "jaxlib"]
+# Benchmark dataset loaders (sweep.datasets). Embedded demo models need
+# nothing extra; parsing is numpy-only. This extra only adds the HTTP client
+# used to fetch the downloadable full-size benchmarks.
+datasets = ["requests>=2.25", "tqdm>=4.60"]
+
+[project.urls]
+Homepage = "https://github.com/DeepWave-KAUST/sweep"
+Issues   = "https://github.com/DeepWave-KAUST/sweep/issues"
+
+[project.scripts]
+sweep = "sweep.cli:main"
+sweep-datasets = "sweep.datasets.cli:main"
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+
+[tool.setuptools.packages.find]
+where = ["src"]
+include = ["sweep*", "geophyai*"]
+exclude = ["sweep.csrc*"]
+
+# Ship the C++/CUDA sources inside the wheel — sweep._C is JIT-compiled from them
+# on first use (csrc is data, not an importable package, hence excluded above).
+[tool.setuptools.package-data]
+sweep = [
+    "csrc/**/*.cu", "csrc/**/*.cuh", "csrc/**/*.cpp",
+    "csrc/**/*.h", "csrc/**/*.hpp", "csrc/**/CMakeLists.txt",
+]

sweep_solver-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

sweep_solver-0.1.0/setup.py

@@ -0,0 +1,15 @@
+import os
+import sys
+
+from setuptools import setup
+
+ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
+if ROOT_DIR not in sys.path:
+    sys.path.insert(0, ROOT_DIR)
+
+from build_config import build_ext_kwargs, patch_packaging_compat
+
+
+patch_packaging_compat()
+
+setup(**build_ext_kwargs())

sweep_solver-0.1.0/setup_cuda.py

@@ -0,0 +1,16 @@
+import os
+import sys
+
+from setuptools import setup
+
+ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
+if ROOT_DIR not in sys.path:
+    sys.path.insert(0, ROOT_DIR)
+
+from build_config import build_ext_kwargs, patch_packaging_compat
+
+
+patch_packaging_compat()
+os.environ["SWEEP_BUILD_CUDA"] = "1"
+
+setup(**build_ext_kwargs(build_cuda=True))

sweep_solver-0.1.0/src/geophyai/__init__.py

@@ -0,0 +1,4 @@
+import sys
+import sweep
+
+sys.modules['geophyai'] = sweep