Metadata-Version: 2.4
Name: jaxFMM
Version: 0.3.1
Summary: Adaptive Fast Multipole Method with Laplace kernel in JAX.
Project-URL: Homepage, https://gitlab.com/jaxfmm/jaxfmm
Project-URL: Documentation, https://jaxfmm.gitlab.io/jaxfmm
Author-email: Robert Kraft <robert.kraft@univie.ac.at>
License: MIT License
        
        Copyright (c) 2025 Robert Kraft
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
License-File: LICENSE
Keywords: FMM,N-body,jax,potential,treecode
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Topic :: Scientific/Engineering :: Physics
Requires-Python: >=3.9
Requires-Dist: jax
Requires-Dist: numpy
Requires-Dist: scipy
Provides-Extra: cuda
Requires-Dist: jax[cuda]; extra == 'cuda'
Provides-Extra: dev
Requires-Dist: matplotlib; extra == 'dev'
Requires-Dist: pytest; extra == 'dev'
Requires-Dist: pyvista; extra == 'dev'
Provides-Extra: docs
Requires-Dist: furo; extra == 'docs'
Requires-Dist: myst-parser; extra == 'docs'
Requires-Dist: sphinx-autodoc-typehints; extra == 'docs'
Requires-Dist: sphinx>=7; extra == 'docs'
Description-Content-Type: text/markdown

# jaxFMM

[![Documentation](https://img.shields.io/badge/docs-jaxfmm.gitlab.io-blue)](https://jaxfmm.gitlab.io/jaxfmm)
[![Paper](https://img.shields.io/badge/DOI-10.1016%2Fj.jcp.2026.115130-blue)](https://doi.org/10.1016/j.jcp.2026.115130)

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](https://docs.jax.dev/en/latest/installation.html).

## Quickstart

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

```python
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):

```python
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](/demos/flavors_tour.py) walks through every engine in a few lines
each; [demos/README.md](/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](https://link.springer.com/article/10.1007/s11227-012-0836-0)),
  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](https://jaxfmm.gitlab.io/jaxfmm);
for algorithmic details, see the [paper](https://doi.org/10.1016/j.jcp.2026.115130); for the
internal architecture, see [ARCHITECTURE.md](/ARCHITECTURE.md).

## Stray Field Evaluation

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

```python
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](/demos/README.md) show the full
pipeline including shape optimization through the solver.

## Citing

jaxFMM is described in our [Journal of Computational Physics paper](https://doi.org/10.1016/j.jcp.2026.115130):

    @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](https://claude.com/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`](/AGENTS.md) (with `CLAUDE.md` as a compatibility import)
carries the working conventions, and [`ARCHITECTURE.md`](/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.

