jaxFMM 0.3.3

Adaptive Fast Multipole Method with Laplace kernel in JAX.

jaxFMM

jaxFMM is an open source implementation of the Fast Multipole Method in JAX. The goal is to offer an easily readable/maintainable FMM implementation with good performance that runs on CPU/GPU and supports autodiff. This is enabled through JAX's just-in-time compiler.

📖 Documentation: https://jaxfmm.gitlab.io/jaxfmm

Installation

jaxFMM depends only on JAX and can be installed from pypi or by downloading the source as follows:

pip install jaxfmm

If you want to run jaxFMM on GPUs, the easiest way is to use NVIDIA CUDA and cuDNN from pip wheels by instead typing:

pip install jaxfmm[cuda]

Using a custom, self-installed CUDA with jax is described in the JAX documentation.

Quickstart

The simplest entry point is a certified evaluator: pick a target accuracy, get a function.

import jaxfmm

fmm = jaxfmm.plan(pts, err_tol=1e-3)   # picks + certifies a configuration on first call
pot = fmm(chrgs)                       # fast evaluations from then on

For explicit control, each engine offers the same two-call pattern (setup once per geometry, evaluate per charge vector):

from jaxfmm import kifmm

s = kifmm.setup(pts, p=3)              # tuned tree + operators for this point cloud
fmm = kifmm.compile_evaluator(s)       # frozen, jitted: fmm(chrgs) -> potential

The flavors tour walks through every engine in a few lines each; demos/README.md describes the full demo set, including two inverse-design examples that differentiate straight through the FMM.

Engines

jaxFMM ships several FMM engines on two tree families, sharing one near-field module:

engine	tree	best at
flexible (jaxfmm.flex)	non-uniform 2^N-ary, arbitrary box shapes	adaptive clouds, shape derivatives
KIFMM (jaxfmm.kifmm)	2:1-balanced cubic octree	general workhorse; potentials and fields (−∇φ); pluggable kernels
uniform spherical-harmonic (jaxfmm.uniform)	cubic octree, precomputed per-level operators	raw speed on space-filling clouds

All engines support periodic boundary conditions (periodic_axes=, hierarchical lattice operators, validated against the NaCl Madelung constant) and mixed-precision fast paths (f32 kernels with f64 accumulation) where the extra noise provably sits below the truncation error — accuracy-neutral defaults, opt-out via m2l_mixed/p2p_mixed.

Features

• Laplace kernel with real solid-harmonic bases computed via recurrence relations; rotation-based O(p^3) M2M/M2L/L2L on the flexible tree (following Goude & Engblom), equivalent-density (kernel-independent) and precomputed-operator variants on the cubic tree.
• Pluggable kernels on the KIFMM engine (jaxfmm.kernels): bring five callables and a symmetry flag; Laplace and Yukawa ship built in (kifmm.setup(pts, kernel=jaxfmm.kernels.yukawa(2.0))).
• End-to-end jit compilation; autodiff through charges and geometry (guarded moving-geometry drivers rebuild capped tree topology inside jit — see jaxfmm.driver).
• Periodic boundary conditions in 1/2/3 axes for all engines.
• Finite-element sources: P0 volume + P1 surface charge FMM with singular (Duffy) self-term correction (jaxfmm.fem).
• Micromagnetic stray-field evaluation for P1-FEM meshes (jaxfmm.apps.mag.strayfield, see below), Oersted fields from currents (jaxfmm.apps.mag.oersted).
• Scales to 10^9 sources on a single 80 GB GPU.

For the full API reference and guide, see the documentation; for algorithmic details, see the paper; for the internal architecture, see ARCHITECTURE.md.

Stray Field Evaluation

jaxFMM is primarily developed for rapid stray field evaluation in finite-element micromagnetics:

from jaxfmm.apps.mag.strayfield import init_strayfield, eval_strayfield

h = init_strayfield(verts, tets, Ms)   # geometry-fixed setup
H = eval_strayfield(m, **h)            # stray field for a magnetization state

eval_strayfield supports both a lumped (fast, per-timestep) and an energy-consistent autodiff mode (mode="diff"); the field-text demos show the full pipeline including shape optimization through the solver.

Citing

jaxFMM is described in our Journal of Computational Physics paper:

@article{kraft2026jaxfmm,
    title = {jaxFMM: An Adaptive, GPU-Parallel Implementation of the Fast Multipole Method in JAX},
    author = {Robert Kraft and Florian Bruckner and Dieter Suess and Claas Abert},
    journal = {Journal of Computational Physics},
    pages = {115130},
    year = {2026},
    doi = {10.1016/j.jcp.2026.115130},
}

TODOs

jaxFMM is primarily developed for fast stray-field evaluation in finite-element micromagnetics, which explains the feature set and design decisions above. Contributions are always welcome; current directions in development:

• symmetry-class M2L application (validated prototype: ~3.7× on non-Laplace kernels, 17× operator-table memory) and further per-kernel performance work.
• multi-RHS evaluation (several charge vectors per tree traversal).
• distributed parallelism via jax.sharding.
• volume FMM (continuous source densities) — prototype exists, on hold.

AI-assisted development

jaxFMM is co-developed with Claude Code: large parts of the library were written, benchmarked and refactored in human-directed agent sessions, with every change gated by the test suite and reviewed by the maintainers. The repo is set up so agents are first-class contributors: AGENTS.md (with CLAUDE.md as a compatibility import) carries the working conventions, and ARCHITECTURE.md is the maintained internal map — including the hard-won invariants and measured dead ends (§7/§8) that keep both humans and agents from re-deriving or re-breaking things. If you contribute with an agent, point it there first.