jz-fmm documentation
jz-fmm (JAX z-order tree, Fast Multipole Method) is a GPU-native FMM implementation for fast and differentiable gravitational N-body simulations. It builds on jz-tree for GPU-native tree construction, traversal, and distributed communication. Its high-level interface is written in JAX, while performance-critical operations use a CUDA backend through JAX’s foreign function interface.
Features
GPU-native force evaluation: the fast multipole method on one or many GPUs, including multi-host execution, with isolated boundary conditions.
Differentiable simulations: gradients with respect to particle positions, velocities, and masses through force calculations and time integration.
Reproducible forces: bit-perfect reproducibility of repeated FMM calculations with the same inputs and configuration on the same hardware.
Reversible time integration: bit-perfect backwards integration with the integer-lattice integrator, alongside a standard floating-point integrator.
Configurable accuracy: selectable multipole orders in two and three dimensions, with single- and double-precision calculations.
Extensibility: custom external potentials and support for adding new interaction kernels beyond gravitational N-body applications.
Convenient Python/JAX interface: compose force calculations, integrators, and external potentials to customize N-body simulations in a few lines of Python, while retaining JIT-compiled GPU performance.
Reference calculations: direct summation for checking approximate forces.
Start with the Getting started, see Multi-GPU guide for distributed execution, or consult JAX compatibility and helpers for function-specific support.
Performance and scaling
The benchmarks below show force accuracy versus evaluation time and weak scaling across multiple GPUs. jz-fmm delivers world-class performance and demonstrates the potential of a fully GPU-native approach to N-body simulations.
Accuracy and performance. Comparison with GADGET-4 (G4) and PKDGRAV3 for a single force evaluation in a Hernquist sphere with \(4 \times 10^7\) particles. Labels give the opening angle; jz-fmm curves show different multipole expansion orders. All codes were benchmarked on one node with four NVIDIA A100 GPUs and a 32-core CPU; a fairer hardware allocation for the CPU-only GADGET-4 comparison would use approximately 4–8 times as many CPU cores.
Multi-GPU scaling. Weak scaling from one to 64 GPUs for uniform particle distributions. In the GPU-saturating regime, efficiency decreases by less than a factor of two from one to 64 devices.