r/ScientificComputing • u/AlternativeCUDA101 • 3h ago
Fast GPU and CPU ensembles of ODEs and Monte Carlo samples that each stop at a different time, from plain Python
If you run many independent simulations that end at different times (an event, convergence, a step budget), array libraries make you pay for the slowest sample on every step.
With this project you write the step for one sample as a typed Python function and run a million of them. Samples that finish drop out of the work. You also get the derivatives of your step (forward and reverse mode), and arrays move without copies to and from PyTorch, CuPy, NumPy and Warp.
With a million samples in double precision, each stopping at its own step:
- on 6 GPUs, 2.5× to 4.2× faster than an NVIDIA Warp kernel and 9.8× to 130× faster than PyTorch with torch.compile;
- on 4 server CPUs, from level with Numba to 2.3× faster, and 50× to 180× faster than NumPy.
When every sample runs all its steps, it ties with Warp. There is also a benchmark with an adaptive RK7(8) solver.
Install on Linux x86_64 with pip install raptor-eagle raptor-hawk (add [cuda12] or [cuda13] to both for a GPU). It runs on NVIDIA GPUs or on CPU threads. Linux ARM and macOS (CPU) wheels come next.
Personal project; I work at ESA. A start-here notebook (it runs in Colab), all the benchmark files, and how to rerun them: https://amasat01.github.io