jaccpot

jaccpot is a JAX-first Fast Multipole Method (FMM) package for N-body gravity and related hierarchical interaction problems. It provides multipole/local expansion operators, near/far-field kernels, and a preset-driven high-level solver API.

Tree construction and traversal artifacts are provided by the companion package yggdrax.

Features

• High-level FastMultipoleMethod API with fast, balanced, and accurate presets
• Configurable expansion basis (solidfmm, cartesian)
• Modular runtime with grouped/dense interaction pathways
• Near-field and far-field execution paths with optional prepared state reuse
• Differentiable gravitational acceleration helper via JAX autodiff

Installation

Install from source:

pip install -e .

yggdrax is not on PyPI yet. Install it from GitHub first (use the latest main, which includes native RadixTree JAX pytree registration):

git clone https://github.com/TobiBu/yggdrax.git
cd yggdrax
pip install -e .
cd ..

Install with development tooling:

pip install -e ".[dev]"

Quick Start

import jax
import jax.numpy as jnp

from jaccpot import FastMultipoleMethod

key = jax.random.PRNGKey(0)
key_pos, key_mass = jax.random.split(key)
positions = jax.random.uniform(key_pos, (1024, 3), minval=-1.0, maxval=1.0)
masses = jax.random.uniform(key_mass, (1024,), minval=0.5, maxval=1.5)

solver = FastMultipoleMethod(preset="balanced", basis="solidfmm")
accelerations = solver.compute_accelerations(positions, masses)
print(accelerations.shape)

For split-step integrators (for example active-particle substeps), you can evaluate only a subset while still using all particles as FMM sources:

active = jnp.asarray([0, 7, 11, 32], dtype=jnp.int32)
state = solver.prepare_state(positions, masses)
active_acc = solver.evaluate_prepared_state(state, target_indices=active)

For ODISSEO-style primitive states (N, 2, 3), you can use the adapter:

from jaccpot import OdisseoFMMCoupler

coupler = OdisseoFMMCoupler(solver, leaf_size=16, max_order=4)
coupler.prepare(primitive_state, masses)  # full source tree
acc_active = coupler.accelerations(primitive_state, active_indices=active)

Development

Run quality gates locally:

black --check .
isort --check-only .
pytest

Or run pre-commit hooks:

pre-commit run --all-files

Coverage is enforced in CI via pytest-cov:

pytest --cov=jaccpot --cov-report=term-missing

Runtime Type Checking

jaccpot can enable package-wide runtime checking for annotated callables using jaxtyping + beartype at import time.

• Disabled by default.
• Enable when needed with:

export JACCPOT_RUNTIME_TYPECHECK=1

Project Structure

• jaccpot/solver.py: preset-first user-facing FMM API
• jaccpot/config.py: config model for solver/runtime knobs
• jaccpot/runtime: execution internals and integration with yggdrax artifacts
• jaccpot/operators: harmonic, translation, and multipole operators
• jaccpot/upward, jaccpot/downward, jaccpot/nearfield: sweep and near-field modules
• tests: unit, integration, and performance checks

CI

GitHub Actions runs:

• formatter checks (black, isort)
• unit/integration tests with coverage threshold
• release build and PyPI publish on version tags

Workflow files:

• .github/workflows/ci.yml
• .github/workflows/release.yml