jaxFMM: An Adaptive, GPU-Parallel Implementation of the Fast Multipole Method in JAX

TL;DR

jaxFMM is a GPU-accelerated, adaptive Fast Multipole Method implementation in JAX that efficiently handles non-uniform charge distributions, enabling faster computations in micromagnetics and potential applications in inverse design and machine learning.

Contribution

It introduces jaxFMM, a highly parallel, adaptive FMM implementation in JAX with a simple code structure for non-uniform problems.

Findings

Performs well on highly non-uniform charge distributions.

Significantly speeds up stray-field computations in micromagnetics.

Enables applications in inverse design and machine learning.

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, which results in extremely concise and simple code. Benchmarks show that the algorithm performs well 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.