arXiv · 2511.15269
jaxFMM: An Adaptive, GPU-Parallel Implementation of the Fast Multipole Method in JAX
Abstract
We introduce jaxFMM, an open-source, adaptive, highly parallel point-charge Fast Multipole Method implementation for the Laplace kernel written in JAX. It is based on a non-uniform refinement strategy with on-the-fly rotation-based transforms tailored around JAX's just-in-time compiler, which results in extremely concise and simple code. Benchmarks show that the algorithm performs well at moderate accuracies, even for highly non-uniform charge distributions. JaxFMM already massively speeds up stray-field computations in micromagnetics and with JAX features like autodiff, novel applications such as inverse-design problems and machine-learning tasks can be tackled with ease in the future.
Explore related subjects
Keep this discovery
Explore connections, maps & timelines
Robert Kraft, Florian Bruckner, Dieter Suess, Claas Abert. 2025-11-19. jaxFMM: An Adaptive, GPU-Parallel Implementation of the Fast Multipole Method in JAX. https://doi.org/10.1016/j.jcp.2026.115130
Cite the original work for its findings. Save a collection to share your selection of sources.