API Reference

This reference lists the supported interfaces for configuring and running jz-fmm. Lower-level implementation functions are intentionally omitted; users interested in those details can consult the source code directly.

Particle data

The core data structures describe simulation particles and the potential, force, and tidal-field values calculated for them.

class jzfmm.data.LocalExpansion(values, dim=3)

Bases: object

Potential and spatial derivatives evaluated at a set of positions.

Parameters:

• values (Array) – Array with potential followed by its spatial derivatives along the last axis.

• dim (int) – Number of spatial dimensions.

fphi()

Returns force components followed by potential.

potential()

Returns the potential.

force()

Returns the force, or negative potential gradient.

tide()

Returns the independent components of the negative Hessian.

class jzfmm.data.Particles(*, pos, mass, vel, loc=None, num=None, num_total=None)

Bases: object

Particle data used by simulations.

num and num_total are only required for multi-GPU execution and may remain None for single-GPU particle data.

Parameters:

• pos (Array) – Position array of shape (size, dim).

• mass (Array) – Particle masses, either scalar or an array of shape (size,).

• vel (Array) – Velocity array with the same shape as pos.

• loc (LocalExpansion | None) – Optional local expansion evaluated at the particle positions.

• num (Array | None) – Number of filled entries on the local device for padded, multi-GPU data.

• num_total (int | None) – Total particle count across all devices for multi-GPU data.

Configuration

Configuration objects select the interaction kernel, opening criterion, force solver, units, and integration method. They are static arguments to compiled JAX functions, so modifying them generally triggers recompilation. Config variables are not differentiable and should not be used as optimization parameters for jax.grad() or jax.vjp().

Configurations are composable dataclasses: specialized objects can be nested to describe a complete simulation. For example, this configuration combines a custom tree, interaction kernel, opening criterion, and integrator:

import jzfmm
from jztree.config import TreeConfig

cfg = jzfmm.SimConfig(
    force=jzfmm.FMMConfig(
        tree=TreeConfig(max_leaf_size=64),
        kernel=jzfmm.PlummerKernel(softening=0.01),
        opening=jzfmm.OpeningByAngle(theta=0.7),
        p=6,
    ),
    integrator=jzfmm.KDKConfig(),
)

Most configurations have sensible defaults, so only parameters relevant to a particular simulation need to be changed. Individual settings can conveniently be modified after construction, including settings in nested configs:

cfg = jzfmm.SimConfig()
cfg.force.kernel.softening = 0.1
cfg.force.p = 6

# Replace an entire nested config to select a different force solver.
cfg.force = jzfmm.DirectSummationConfig(
    kernel=jzfmm.PlummerKernel(softening=0.1)
)

Config hashes include the values of nested configuration objects. Configs should therefore not be modified after being added to a dictionary, set, or other hash-based container.

class jzfmm.config.KernelConfig

Bases: object

Base class for radial interaction-kernel configurations.

kind_id()

Returns the backend identifier of the kernel.

params(dtype=jnp.float32)

Returns kernel parameters in the requested dtype.

class jzfmm.config.PlummerKernel(softening=0.001)

Bases: KernelConfig

Plummer-softened inverse-distance kernel.

Implements \(K(r)=-1/\sqrt{r^2+\epsilon^2}\).

Parameters:

softening (float) – Softening length \(\epsilon\).

kind_id()

Returns the backend identifier of the kernel.

params(dtype=jnp.float32)

Returns kernel parameters in the requested dtype.

class jzfmm.config.Plummer2DKernel(softening=0.001)

Bases: KernelConfig

Two-dimensional Plummer-softened logarithmic kernel.

Implements \(K(r)=\frac{1}{2}\log(r^2+\epsilon^2)\).

Parameters:

softening (float) – Softening length \(\epsilon\).

kind_id()

Returns the backend identifier of the kernel.

params(dtype=jnp.float32)

Returns kernel parameters in the requested dtype.

class jzfmm.config.SoftenedDistanceKernel(softening=0.001)

Bases: KernelConfig

Softened distance kernel.

Implements \(K(r)=\sqrt{r^2+\epsilon^2}\).

Parameters:

softening (float) – Softening length \(\epsilon\).

kind_id()

Returns the backend identifier of the kernel.

params(dtype=jnp.float32)

Returns kernel parameters in the requested dtype.

class jzfmm.config.OpeningCriterionConfig

Bases: object

Base class for FMM opening-criterion configurations.

kind_id()

Returns the backend identifier of the opening criterion.

params(dtype=jnp.float32)

Returns criterion parameters in the requested dtype.

class jzfmm.config.OpeningByAngle(theta=0.8)

Bases: OpeningCriterionConfig

Geometric opening criterion based on an opening angle.

Parameters:

theta (float) – Maximum opening angle. Smaller values increase accuracy and computational cost.

kind_id()

Returns the backend identifier of the opening criterion.

params(dtype=jnp.float32)

Returns criterion parameters in the requested dtype.

class jzfmm.config.PotentialField

Bases: object

Base class for external potential fields.

Assign an instance to SimConfig.external_potential to add its acceleration during time integration. Subclasses normally implement potential(); acceleration() obtains its negative gradient with autodiff. Built-in fields are provided by jzfmm.external_potential.

The external contribution is applied during integration and is not included in the jzfmm.data.LocalExpansion returned for particle self-interactions.

potential(x, t=0., cfg=None)

Evaluates the potential at positions x.

acceleration(x, t=0., cfg=None)

Evaluates acceleration as the negative potential gradient.

class jzfmm.config.UnitConfig(pos_in_kpc=1.0, vel_in_kmps=1.0, mass_in_msol=1.0)

Bases: object

Defines simulation units relative to common astrophysical units.

Parameters:

• pos_in_kpc (float) – Length represented by one simulation position unit in kpc.

• vel_in_kmps (float) – Speed represented by one simulation velocity unit in km/s.

• mass_in_msol (float) – Mass represented by one simulation mass unit in solar masses.

G()

Returns the gravitational constant in simulation units.

class jzfmm.config.IntegratorConfig

Bases: object

Base class for time-integrator configurations.

class jzfmm.config.DKDConfig

Bases: IntegratorConfig

Drift-kick-drift leapfrog integrator configuration.

class jzfmm.config.KDKConfig

Bases: IntegratorConfig

Kick-drift-kick leapfrog integrator configuration.

class jzfmm.config.DKDLatticeConfig(dx=0.0001, dv=0.0001, int_dtype=<class 'jax.numpy.int32'>)

Bases: IntegratorConfig

Drift-kick-drift integrator using integer phase-space coordinates.

Parameters:

• dx (float) – Position lattice spacing.

• dv (float) – Velocity lattice spacing.

• int_dtype (type) – Integer dtype used for lattice coordinates. Both jax.numpy.int32 and jax.numpy.int64 are supported, independently of the floating-point format used elsewhere.

class jzfmm.config.DirectSummationConfig(kernel=<factory>, kahan_summation=True, remove_self_interaction=True)

Bases: object

Configures direct pair summation.

Parameters:

• kernel (KernelConfig) – Radial interaction kernel.

• kahan_summation (bool) – Whether to use compensated summation.

• remove_self_interaction (bool) – Whether to exclude each particle’s interaction with itself.

class jzfmm.config.FMMConfig(tree=<factory>, kernel=<factory>, p=5, opening=<factory>, alloc_fac_ilist=64.0, alloc_fac_comm_nodes=1.5, alloc_fac_comm_particles=1.5, kahan_summation=False, remove_self_interaction=True)

Bases: object

Configures fast-multipole force evaluation.

Parameters:

• tree (TreeConfig) – Tree construction configuration.

• kernel (KernelConfig) – Radial interaction kernel.

• p (int) – Multipole expansion order.

• opening (OpeningCriterionConfig) – Node-opening criterion.

• alloc_fac_ilist (float) – Interaction-list entries allocated per leaf-node buffer entry; the total capacity is approximately this factor times the allocated number of leaf nodes.

• alloc_fac_comm_nodes (float) – Node communication-buffer capacity as a multiple of the local node-buffer size.

• alloc_fac_comm_particles (float) – Particle communication-buffer capacity as a multiple of the local particle-buffer size.

• kahan_summation (bool) – Whether to use compensated summation where available.

• remove_self_interaction (bool) – Whether to exclude each particle’s interaction with itself.

class jzfmm.config.SimConfig(force=<factory>, units=<factory>, logging=<factory>, external_potential=None, integrator=<factory>)

Bases: object

Collects force, unit, logging, and integration configuration.

Parameters:

• force (FMMConfig | DirectSummationConfig | None) – Force configuration, or None to disable self-gravity.

• units (UnitConfig) – Simulation unit configuration.

• logging (LoggingConfig) – Logging configuration.

• external_potential (PotentialField | None) – Optional external potential field.

• integrator (IntegratorConfig) – Time-integrator configuration.

Force evaluation

These functions evaluate interactions using either the fast multipole method or direct summation.

jzfmm.fmm.direct_summation(part, cfg_direct, G=1)

Evaluates all pair interactions by direct summation.

Compatibility: JIT Local only Autodiff

Helpers: .jit

This function does not implement multi-device communication and will run only on local shards when called inside of a shard_map.

Parameters:

• part (PosMass) – Particle positions and masses.

• cfg_direct (DirectSummationConfig) – Direct-summation configuration.

• G (float | Array) – Gravitational constant or multiplicative interaction strength.

Returns:

Local expansion containing the potential and force at each particle.

Return type:

LocalExpansion

jzfmm.fmm.fast_multipole_method(part, cfg_fmm, th=None, result='loc', G=1., pout=1)

Evaluates particle interactions with the fast multipole method.

Compatibility: JIT Shard map Autodiff

Helpers: .jit .smap

Parameters:

• part (PosMass) – Particle data following the jztree.data.PosMass interface.

• cfg_fmm (FMMConfig) – Fast-multipole configuration.

• th (TreeHierarchy | None) – Existing tree hierarchy. If provided, part must already be in z-order and "loc" cannot be requested.

• result (str) – Underscore-separated selection of "loc" (input-order local expansions), "locz" (z-order local expansions), "partz" (z-order particles), and "tree" (tree hierarchy).

• G (float) – Gravitational constant or multiplicative interaction strength.

• pout (int) – Output expansion order. Currently only 1 is supported, returning the potential and force.

Returns:

The requested result, or a tuple when multiple results are requested.

Return type:

LocalExpansion

Time integration

These functions calculate forces, advance particles, and run simulations.

jzfmm.time_integration.force_and_potential(p, cfg)

Evaluates particle self-interactions selected by the simulation config.

Compatibility: JIT Shard map Autodiff

Helpers: .jit

Shard-map execution produces a global result with jzfmm.config.FMMConfig; direct summation is local only.

Parameters:

• p (Particles) – Particles at which to evaluate the force and potential.

• cfg (SimConfig) – Simulation configuration selecting the force method and gravitational constant.

Returns:

Local expansion containing potential and force.

Return type:

LocalExpansion

jzfmm.time_integration.find_center(p, npot=50, nbind=None)

Estimates the position and velocity center of a particle system.

Compatibility: JIT Local only Autodiff untested

Helpers: .jit

Parameters:

• p (Particles) – Particles with potentials stored in p.loc.

• npot (int) – Number of lowest-potential particles used for the position and initial velocity center.

• nbind (int | None) – Number of most-bound particles used for the final velocity center. Defaults to npot.

Returns:

Tuple containing the center position and center velocity.

Return type:

tuple[Array, Array]

jzfmm.time_integration.timestep(p, dt, cfg, t=0.)

Advances particles by one integration step.

Compatibility: JIT Shard map Autodiff

Helpers: .jit

Shard-map execution produces a global result with jzfmm.config.FMMConfig; direct summation is local only. Autodiff is supported for floating-point DKD and KDK integration, but not through the integer operations of jzfmm.config.DKDLatticeConfig.

Parameters:

• p (Particles) – Particles at the start of the step.

• dt (float | Array) – Step size.

• cfg (SimConfig) – Simulation configuration.

• t (float | Array) – Time at the start of the step.

Returns:

Particles at time t + dt.

Return type:

Particles

jzfmm.time_integration.simulate(p, ts, cfg)

Evolves particles through a sequence of times.

Compatibility: JIT Shard map Autodiff

Helpers: .jit .smap

Shard-map execution produces a global result with jzfmm.config.FMMConfig; direct summation is local only. Reverse-mode differentiation is supported with jzfmm.config.DKDConfig and jzfmm.config.DKDLatticeConfig, but not with jzfmm.config.KDKConfig.

Parameters:

• p (Particles) – Particles at the first time in ts.

• ts (Array) – One-dimensional sequence of times, including the initial time.

• cfg (SimConfig) – Simulation configuration.

Returns:

Particles evolved to the final time in ts.

Return type:

Particles

jzfmm.time_integration.simulate_with_outputs(p, tend, nout, steps_per_output, cfg, tstart=0.)

Runs a simulation and yields particles at regular output times.

Compatibility: JIT unsupported Shard map untested Autodiff unsupported

This host-side generator compiles and calls simulate() internally; it should not itself be transformed with JAX.

Parameters:

• p (Particles) – Particles at tstart.

• tend (float) – Total simulated duration.

• nout (int) – Number of output intervals.

• steps_per_output (int) – Integration steps per output interval.

• cfg (SimConfig) – Simulation configuration.

• tstart (float) – Initial simulation time.

Yields:

(time, particles) pairs, including the initial state.

External potentials

External potential fields can be attached to a simulation through jzfmm.config.SimConfig.

class jzfmm.external_potential.NFWPotential(rs=1.0, rhoc=1.0)

Bases: PotentialField

Spherical Navarro-Frenk-White potential.

The additive constant is chosen so that the potential approaches zero at the center.

Parameters:

• rs (float) – Scale radius.

• rhoc (float) – Characteristic density.

phic(G=1.)

Returns the characteristic potential scale.

potential(x, t=0., cfg=None)

Evaluates the potential at x.

class jzfmm.external_potential.HernquistPotential(a=1.0, mass=1.0)

Bases: PotentialField

Spherical Hernquist potential.

Parameters:

• a (float) – Scale radius.

• mass (float) – Total mass.

potential(x, t=0., cfg=None)

Evaluates the potential at x.

class jzfmm.external_potential.UniformAcceleration(acc=(0.0, 0.0, 0.0))

Bases: PotentialField

Potential producing a spatially uniform acceleration.

Parameters:

acc (tuple[float, float, float]) – Acceleration vector.

potential(x, t=0., cfg=None)

Evaluates the potential at x.

class jzfmm.external_potential.DiskPotential(mass, scale_radius, height)

Bases: PotentialField

Disk potential approximated by three Miyamoto-Nagai potentials.

Uses the approximation from Smith et al. (2015).

Parameters:

• mass (float) – Total disk mass.

• scale_radius (float) – Exponential scale radius.

• height (float) – Disk scale height.

potential(x, t=0., cfg=None)

Evaluates the potential at x.

class jzfmm.external_potential.MilkyWayPotential(halo=<factory>, bulge=<factory>, star_disk=<factory>, gas_disk=<factory>)

Bases: PotentialField

Composite Milky-Way-like potential.

Parameters:

• halo (PotentialField | None) – Dark-matter halo potential, or None to omit it.

• bulge (PotentialField | None) – Stellar bulge potential, or None to omit it.

• star_disk (PotentialField | None) – Stellar disk potential, or None to omit it.

• gas_disk (PotentialField | None) – Gas disk potential, or None to omit it.

potential(x, t=0., cfg=None)

Evaluates the sum of all enabled component potentials.

Loss functions

jzfmm.loss.maximum_mean_discrepancy(part, part_target, *, cfg_fmm, normalize=True)

Maximum mean discrepancy between two weighted particle sets.

Compatibility: JIT Shard map Autodiff

Helpers: .jit .smap

The MMD is evaluated as a signed kernel energy, using positive masses for part and negative masses for part_target: -sum_i m_i phi_i. Self interactions must be enabled so that the kernel diagonal is included. Shard-map execution produces a global result with jzfmm.config.FMMConfig; direct summation is local only.

Parameters:

• part (PosMass | Pos | Array) – First particle set or position array.

• part_target (PosMass | Pos | Array) – Target particle set or position array.

• cfg_fmm (FMMConfig | DirectSummationConfig) – FMM or direct-summation kernel configuration, with remove_self_interaction=False.

• normalize (bool) – Whether to normalize each particle set to unit total mass.

Returns:

Scalar discrepancy between the particle sets.

Return type:

Array

Additional helpers

The jzfmm_utils package contains optional helpers for creating initial conditions and plotting particle distributions. These conveniences are not part of the documented core API.