API reference¶

The public API is re-exported from the top-level yggdrax package. The modules below are documented from their source docstrings.

Tree building¶

Public tree-building API for Yggdrax.

class yggdrax.tree.Tree[source]¶

Bases: object

Public base class for concrete tree containers.

property num_nodes: int¶: Return number of nodes in the concrete topology.

property num_particles: int¶: Return number of particles represented by this tree.

property num_leaves: int¶: Return number of leaf nodes represented by this tree.

property missing_fmm_topology_fields: tuple[str, ...]¶: Return missing topology fields required by FMM core APIs.

property supports_fmm_topology: bool¶: Whether this tree exposes topology needed by FMM core routines.

require_fmm_topology()[source]¶

Raise when this tree cannot satisfy FMM core topology requirements.

Return type:: None

classmethod from_particles(positions, masses, *, tree_type='radix', build_mode='adaptive', bounds=None, return_reordered=True, workspace=None, return_workspace=False, leaf_size=8, target_leaf_particles=32, max_depth=None, refine_local=True, max_refine_levels=2, aspect_threshold=8.0, min_refined_leaf_particles=2)[source]¶

Build a concrete tree, dispatching on tree_type.

Primary entry point that normalizes arguments into a TreeBuildRequest and dispatches to the registered backend builder (see available_tree_types() and register_tree_builder()).

Parameters:

positions (Array) – Particle positions of shape (n, 3).
masses (Array) – Particle masses of shape (n,).
tree_type (str) – Backend identifier: "radix", "octree", "kdtree", or any registered type.
build_mode (str) – Construction mode: "adaptive", "fixed_depth", or "static_radix".
bounds (tuple[Array, Array] | None) – Optional (min_corner, max_corner) box; inferred when omitted.
return_reordered (bool) – If True (default), populate the reordered particle buffers on the returned tree.
workspace (RadixTreeWorkspace | None) – Optional reusable RadixTreeWorkspace.
return_workspace (bool) – If True, retain the workspace on the returned tree.
leaf_size (int) – Maximum particles per leaf (adaptive/static modes).
target_leaf_particles (int) – Target per-leaf occupancy (fixed-depth mode).
max_depth (int | None) – Optional hard cap on the fixed-depth Morton depth.
refine_local (bool) – Whether to locally refine elongated buckets (fixed-depth mode).
max_refine_levels (int) – Maximum additional local refinement depth.
aspect_threshold (float) – Axis aspect-ratio above which a bucket is refined.
min_refined_leaf_particles (int) – Smallest occupancy a locally refined leaf may have.

Returns:

A concrete tree container of the requested backend type.

Return type:

Tree

Raises:

ValueError – If tree_type is not registered.

class yggdrax.tree.TreeBuildRequest(positions, masses, build_mode, bounds, return_reordered, workspace, return_workspace, leaf_size, target_leaf_particles, max_depth, refine_local, max_refine_levels, aspect_threshold, min_refined_leaf_particles)[source]¶

Bases: object

Common request object passed to registered tree builders.

Registered builders receive a fully normalized request so wrapper code can share one dispatch path across radix and backend-provided tree families.

Parameters:

positions (Array)
masses (Array)
build_mode (str)
bounds (tuple[Array, Array] | None)
return_reordered (bool)
workspace (RadixTreeWorkspace | None)
return_workspace (bool)
leaf_size (int)
target_leaf_particles (int)
max_depth (int | None)
refine_local (bool)
max_refine_levels (int)
aspect_threshold (float)
min_refined_leaf_particles (int)

positions: Array¶

masses: Array¶

build_mode: str¶

bounds: tuple[Array, Array] | None¶

return_reordered: bool¶

workspace: RadixTreeWorkspace | None¶

return_workspace: bool¶

leaf_size: int¶

target_leaf_particles: int¶

max_depth: int | None¶

refine_local: bool¶

max_refine_levels: int¶

aspect_threshold: float¶

min_refined_leaf_particles: int¶

class yggdrax.tree.RadixTree(topology, build_mode='adaptive', positions_sorted=None, masses_sorted=None, inverse_permutation=None, workspace=None)[source]¶

Bases: Tree

Concrete radix-tree container implementing the generic Tree contract.

Parameters:

topology (RadixTree)
build_mode (Literal['adaptive', 'fixed_depth', 'static_radix'])
positions_sorted (Array | None)
masses_sorted (Array | None)
inverse_permutation (Array | None)
workspace (RadixTreeWorkspace | None)

topology: RadixTree¶

build_mode: Literal['adaptive', 'fixed_depth', 'static_radix'] = 'adaptive'¶

positions_sorted: Array | None = None¶

masses_sorted: Array | None = None¶

inverse_permutation: Array | None = None¶

workspace: RadixTreeWorkspace | None = None¶

property tree_type: Literal['radix', 'octree', 'kdtree']¶: Tree-family identifier for this concrete tree.

classmethod from_particles(positions, masses, *, build_mode='adaptive', bounds=None, return_reordered=True, workspace=None, return_workspace=False, leaf_size=8, target_leaf_particles=32, max_depth=None, refine_local=True, max_refine_levels=2, aspect_threshold=8.0, min_refined_leaf_particles=2)[source]¶

Build a radix tree from particles using a selected build mode.

Parameters:

positions (Array)
masses (Array)
build_mode (str)
bounds (tuple[Array, Array] | None)
return_reordered (bool)
workspace (RadixTreeWorkspace | None)
return_workspace (bool)
leaf_size (int)
target_leaf_particles (int)
max_depth (int | None)
refine_local (bool)
max_refine_levels (int)
aspect_threshold (float)
min_refined_leaf_particles (int)

Return type:

RadixTree

class yggdrax.tree.OctreeTree(topology, build_mode='adaptive', positions_sorted=None, masses_sorted=None, inverse_permutation=None, workspace=None)[source]¶

Bases: RadixTree

Oct-tree container built from an octree-specific Morton partition path.

Parameters:

topology (RadixTree)
build_mode (Literal['adaptive', 'fixed_depth', 'static_radix'])
positions_sorted (Array | None)
masses_sorted (Array | None)
inverse_permutation (Array | None)
workspace (RadixTreeWorkspace | None)

property tree_type: Literal['radix', 'octree', 'kdtree']¶: Tree-family identifier for this concrete tree.

property oct_num_nodes: int¶: Return the number of explicit octree cells carried by the topology.

property oct_num_leaf_nodes: int¶: Return the number of valid explicit octree leaves.

Build an octree from particles using the octree-specific build path.

Parameters:

positions (Array)
masses (Array)
build_mode (str)
bounds (tuple[Array, Array] | None)
return_reordered (bool)
workspace (RadixTreeWorkspace | None)
return_workspace (bool)
leaf_size (int)
target_leaf_particles (int)
max_depth (int | None)
refine_local (bool)
max_refine_levels (int)
aspect_threshold (float)
min_refined_leaf_particles (int)

Return type:

OctreeTree

class yggdrax.tree.OctreeTopology(parent, left_child, right_child, left_is_leaf, right_is_leaf, particle_indices, morton_codes, node_ranges, num_particles, num_internal_nodes, node_level, level_offsets, nodes_by_level, num_levels, bounds_min, bounds_max, leaf_codes, leaf_depths, use_morton_geometry, leaf_size, oct_parent, oct_children, oct_child_counts, oct_child_mask, oct_valid_mask, oct_node_codes, oct_node_depths, oct_node_ranges, oct_nodes_by_level, oct_level_offsets, oct_num_levels, oct_leaf_mask, oct_leaf_nodes, radix_node_to_oct, radix_leaf_to_oct)[source]¶

Bases: NamedTuple

Octree-native metadata plus shared compatibility topology fields.

Parameters:

parent (Array)
left_child (Array)
right_child (Array)
left_is_leaf (Array)
right_is_leaf (Array)
particle_indices (Array)
morton_codes (Array)
node_ranges (Array)
num_particles (int)
num_internal_nodes (int)
node_level (Array)
level_offsets (Array)
nodes_by_level (Array)
num_levels (Array)
bounds_min (Array)
bounds_max (Array)
leaf_codes (Array)
leaf_depths (Array)
use_morton_geometry (Array)
leaf_size (int | None)
oct_parent (Array)
oct_children (Array)
oct_child_counts (Array)
oct_child_mask (Array)
oct_valid_mask (Array)
oct_node_codes (Array)
oct_node_depths (Array)
oct_node_ranges (Array)
oct_nodes_by_level (Array)
oct_level_offsets (Array)
oct_num_levels (Array)
oct_leaf_mask (Array)
oct_leaf_nodes (Array)
radix_node_to_oct (Array)
radix_leaf_to_oct (Array)

parent: Array¶: Alias for field number 0

left_child: Array¶: Alias for field number 1

right_child: Array¶: Alias for field number 2

left_is_leaf: Array¶: Alias for field number 3

right_is_leaf: Array¶: Alias for field number 4

particle_indices: Array¶: Alias for field number 5

morton_codes: Array¶: Alias for field number 6

node_ranges: Array¶: Alias for field number 7

num_particles: int¶: Alias for field number 8

num_internal_nodes: int¶: Alias for field number 9

node_level: Array¶: Alias for field number 10

level_offsets: Array¶: Alias for field number 11

nodes_by_level: Array¶: Alias for field number 12

num_levels: Array¶: Alias for field number 13

bounds_min: Array¶: Alias for field number 14

bounds_max: Array¶: Alias for field number 15

leaf_codes: Array¶: Alias for field number 16

leaf_depths: Array¶: Alias for field number 17

use_morton_geometry: Array¶: Alias for field number 18

leaf_size: int | None¶: Alias for field number 19

oct_parent: Array¶: Alias for field number 20

oct_children: Array¶: Alias for field number 21

oct_child_counts: Array¶: Alias for field number 22

oct_child_mask: Array¶: Alias for field number 23

oct_valid_mask: Array¶: Alias for field number 24

oct_node_codes: Array¶: Alias for field number 25

oct_node_depths: Array¶: Alias for field number 26

oct_node_ranges: Array¶: Alias for field number 27

oct_nodes_by_level: Array¶: Alias for field number 28

oct_level_offsets: Array¶: Alias for field number 29

oct_num_levels: Array¶: Alias for field number 30

oct_leaf_mask: Array¶: Alias for field number 31

oct_leaf_nodes: Array¶: Alias for field number 32

radix_node_to_oct: Array¶: Alias for field number 33

radix_leaf_to_oct: Array¶: Alias for field number 34

class yggdrax.tree.KDParticleTree(topology, build_mode='adaptive', positions_sorted=None, masses_sorted=None, inverse_permutation=None, workspace=None)[source]¶

Bases: Tree

Concrete KD-tree container implementing the generic Tree contract.

Parameters:

topology (LeafKDTree)
build_mode (Literal['adaptive'])
positions_sorted (Array | None)
masses_sorted (Array | None)
inverse_permutation (Array | None)
workspace (RadixTreeWorkspace | None)

topology: LeafKDTree¶

build_mode: Literal['adaptive'] = 'adaptive'¶

positions_sorted: Array | None = None¶

masses_sorted: Array | None = None¶

inverse_permutation: Array | None = None¶

workspace: RadixTreeWorkspace | None = None¶

property tree_type: Literal['radix', 'octree', 'kdtree']¶: Tree-family identifier for this concrete tree.

Build a concrete tree, dispatching on tree_type.

Primary entry point that normalizes arguments into a TreeBuildRequest and dispatches to the registered backend builder (see available_tree_types() and register_tree_builder()).

Parameters:

positions (Array) – Particle positions of shape (n, 3).
masses (Array) – Particle masses of shape (n,).
tree_type – Backend identifier: "radix", "octree", "kdtree", or any registered type.
build_mode (str) – Construction mode: "adaptive", "fixed_depth", or "static_radix".
bounds (tuple[Array, Array] | None) – Optional (min_corner, max_corner) box; inferred when omitted.
return_reordered (bool) – If True (default), populate the reordered particle buffers on the returned tree.
workspace (RadixTreeWorkspace | None) – Optional reusable RadixTreeWorkspace.
return_workspace (bool) – If True, retain the workspace on the returned tree.
leaf_size (int) – Maximum particles per leaf (adaptive/static modes).
target_leaf_particles (int) – Target per-leaf occupancy (fixed-depth mode).
max_depth (int | None) – Optional hard cap on the fixed-depth Morton depth.
refine_local (bool) – Whether to locally refine elongated buckets (fixed-depth mode).
max_refine_levels (int) – Maximum additional local refinement depth.
aspect_threshold (float) – Axis aspect-ratio above which a bucket is refined.
min_refined_leaf_particles (int) – Smallest occupancy a locally refined leaf may have.

Returns:

A concrete tree container of the requested backend type.

Return type:

Tree

Raises:

ValueError – If tree_type is not registered.

class yggdrax.tree.RadixTreeWorkspace(parent, left_child, right_child, left_is_leaf, right_is_leaf, node_ranges)[source]¶

Bases: NamedTuple

Reusable buffers for radix tree construction.

Parameters:

parent (Array)
left_child (Array)
right_child (Array)
left_is_leaf (Array)
right_is_leaf (Array)
node_ranges (Array)

parent: Array¶: Alias for field number 0

left_child: Array¶: Alias for field number 1

right_child: Array¶: Alias for field number 2

left_is_leaf: Array¶: Alias for field number 3

right_is_leaf: Array¶: Alias for field number 4

node_ranges: Array¶: Alias for field number 5

class yggdrax.tree.TreeBuildConfig(leaf_size=8, return_reordered=False, workspace=None, return_workspace=False)[source]¶

Bases: object

Resolved options for standard LBVH tree construction.

Attributes:: leaf_size: Maximum particles per Morton leaf. return_reordered: Whether to return Morton-sorted particle arrays. workspace: Optional reusable radix workspace. return_workspace: Whether to return the workspace alongside the tree.

Parameters:

leaf_size (int)
return_reordered (bool)
workspace (RadixTreeWorkspace | None)
return_workspace (bool)

leaf_size: int = 8¶

return_reordered: bool = False¶

workspace: RadixTreeWorkspace | None = None¶

return_workspace: bool = False¶

class yggdrax.tree.FixedDepthTreeBuildConfig(target_leaf_particles=32, return_reordered=False, workspace=None, return_workspace=False, max_depth=None, refine_local=True, max_refine_levels=2, aspect_threshold=8.0, min_refined_leaf_particles=2)[source]¶

Bases: object

Resolved options for fixed-depth tree construction.

Attributes:: target_leaf_particles: Target occupancy used to resolve Morton depth. return_reordered: Whether to return Morton-sorted particle arrays. workspace: Optional reusable radix workspace. return_workspace: Whether to return the workspace alongside the tree. max_depth: Optional upper bound on Morton depth. refine_local: Whether to locally refine elongated Morton buckets. max_refine_levels: Maximum extra local refinement depth. aspect_threshold: Axis-aligned aspect ratio threshold for refinement. min_refined_leaf_particles: Smallest locally refined leaf occupancy.

Parameters:

target_leaf_particles (int)
return_reordered (bool)
workspace (RadixTreeWorkspace | None)
return_workspace (bool)
max_depth (int | None)
refine_local (bool)
max_refine_levels (int)
aspect_threshold (float)
min_refined_leaf_particles (int)

target_leaf_particles: int = 32¶

return_reordered: bool = False¶

workspace: RadixTreeWorkspace | None = None¶

return_workspace: bool = False¶

max_depth: int | None = None¶

refine_local: bool = True¶

max_refine_levels: int = 2¶

aspect_threshold: float = 8.0¶

min_refined_leaf_particles: int = 2¶

yggdrax.tree.build_static_radix_tree(positions, masses, bounds=None, *, leaf_size=8, return_reordered=False, return_workspace=False)[source]¶

Build a fixed-shape radix tree from Morton-sorted count buckets.

The tree structure (node count, parent/child topology) is static for a fixed particle count and leaf_size, so it can be rebuilt cheaply for new particle values via rebuild_static_radix_tree_from_template(). Leaves are not fixed spatial cells; each leaf owns a contiguous chunk of the current Morton-sorted particle order.

Parameters:

positions (Array) – Particle positions of shape (n, 3).
masses (Array) – Particle masses of shape (n,).
bounds (tuple[Array, Array] | None) – Optional (min_corner, max_corner) box; inferred when omitted.
leaf_size (int) – Fixed number of particles per bucket that determines the static shape.
return_reordered (bool) – If True, also return the reordered particle buffers.
return_workspace (bool) – If True, also return the reusable workspace/template.

Returns:

The tree, or a tuple additionally containing the reordered buffers and/or workspace when the corresponding flags are set.

Return type:

RadixTree or tuple

yggdrax.tree.build_fixed_depth_tree(positions, masses, bounds=None, *, target_leaf_particles=32, return_reordered=False, workspace=None, return_workspace=False, max_depth=None, refine_local=True, max_refine_levels=2, aspect_threshold=8.0, min_refined_leaf_particles=2, config=None)[source]¶

Build a fixed-depth Morton tree, inferring bounds when not provided.

Resolves a uniform Morton depth from target_leaf_particles and, when refine_local is set, locally refines elongated leaf buckets by axis aspect ratio.

Parameters:

positions (Array) – Particle positions of shape (n, 3).
masses (Array) – Particle masses of shape (n,).
bounds (tuple[Array, Array] | None) – Optional (min_corner, max_corner) box; inferred when omitted.
target_leaf_particles (int) – Target per-leaf occupancy used to resolve the Morton depth.
return_reordered (bool) – If True, also return the reordered particle buffers.
workspace (RadixTreeWorkspace | None) – Optional reusable RadixTreeWorkspace.
return_workspace (bool) – If True, also return the workspace.
max_depth (int | None) – Optional hard cap on the Morton depth.
refine_local (bool) – Whether to locally refine elongated Morton buckets.
max_refine_levels (int) – Maximum additional local refinement depth.
aspect_threshold (float) – Axis aspect-ratio above which a bucket is refined.
min_refined_leaf_particles (int) – Smallest occupancy a locally refined leaf may have.
config (FixedDepthTreeBuildConfig | None) – Optional FixedDepthTreeBuildConfig overriding the individual keyword arguments.

Returns:

The tree, or a tuple additionally containing the reordered buffers and/or workspace when the corresponding flags are set.

Return type:

RadixTree or tuple

yggdrax.tree.build_fixed_depth_octree(positions, masses, bounds=None, *, target_leaf_particles=32, return_reordered=False, workspace=None, return_workspace=False, max_depth=None, refine_local=True, max_refine_levels=2, aspect_threshold=8.0, min_refined_leaf_particles=2, config=None)[source]¶

Build a fixed-depth octree through the octree-specific build path.

Parameters:

positions (Array)
masses (Array)
bounds (tuple[Array, Array] | None)
target_leaf_particles (int)
return_reordered (bool)
workspace (RadixTreeWorkspace | None)
return_workspace (bool)
max_depth (int | None)
refine_local (bool)
max_refine_levels (int)
aspect_threshold (float)
min_refined_leaf_particles (int)
config (FixedDepthTreeBuildConfig | None)

yggdrax.tree.build_fixed_depth_tree_jit(positions, masses, bounds=None, *, target_leaf_particles=32, return_reordered=False, workspace=None, return_workspace=False, max_depth=None, refine_local=True, max_refine_levels=2, aspect_threshold=8.0, min_refined_leaf_particles=2, config=None)[source]¶

JIT build for a fixed-depth tree, inferring bounds when not provided.

Parameters:

positions (Array)
masses (Array)
bounds (tuple[Array, Array] | None)
target_leaf_particles (int)
return_reordered (bool)
workspace (RadixTreeWorkspace | None)
return_workspace (bool)
max_depth (int | None)
refine_local (bool)
max_refine_levels (int)
aspect_threshold (float)
min_refined_leaf_particles (int)
config (FixedDepthTreeBuildConfig | None)

yggdrax.tree.build_fixed_depth_octree_jit(positions, masses, bounds=None, *, target_leaf_particles=32, return_reordered=False, workspace=None, return_workspace=False, max_depth=None, refine_local=True, max_refine_levels=2, aspect_threshold=8.0, min_refined_leaf_particles=2, config=None)[source]¶

JIT build for a fixed-depth octree through the octree-native path.

Parameters:

positions (Array)
masses (Array)
bounds (tuple[Array, Array] | None)
target_leaf_particles (int)
return_reordered (bool)
workspace (RadixTreeWorkspace | None)
return_workspace (bool)
max_depth (int | None)
refine_local (bool)
max_refine_levels (int)
aspect_threshold (float)
min_refined_leaf_particles (int)
config (FixedDepthTreeBuildConfig | None)

yggdrax.tree.build_octree(positions, masses, bounds=None, *, return_reordered=False, leaf_size=8, workspace=None, return_workspace=False, config=None)[source]¶

Build an octree through the octree-specific Morton partition pipeline.

Produces an OctreeTree that carries the same compatibility fields as build_tree() plus explicit octree buffers (oct_children, oct_node_depths, radix_node_to_oct, …) for level-wise FMM scheduling.

Parameters:

positions (Array) – Particle positions of shape (n, 3).
masses (Array) – Particle masses of shape (n,).
bounds (tuple[Array, Array] | None) – Optional (min_corner, max_corner) box; inferred when omitted.
return_reordered (bool) – If True, also return the reordered particle buffers and inverse permutation.
leaf_size (int) – Maximum number of particles per leaf.
workspace (RadixTreeWorkspace | None) – Optional reusable RadixTreeWorkspace.
return_workspace (bool) – If True, also return the workspace.
config (TreeBuildConfig | None) – Optional TreeBuildConfig overriding the individual keyword arguments.

Returns:

The octree, or a tuple additionally containing the reordered buffers and/or workspace when the corresponding flags are set.

Return type:

OctreeTree or tuple

yggdrax.tree.build_octree_jit(positions, masses, bounds=None, *, return_reordered=False, leaf_size=8, workspace=None, return_workspace=False, config=None)[source]¶

JIT-compiled variant of build_octree() (see it for parameters/returns).

Parameters:

positions (Array)
masses (Array)
bounds (tuple[Array, Array] | None)
return_reordered (bool)
leaf_size (int)
workspace (RadixTreeWorkspace | None)
return_workspace (bool)
config (TreeBuildConfig | None)

yggdrax.tree.build_tree(positions, masses, bounds=None, *, return_reordered=False, leaf_size=8, workspace=None, return_workspace=False, config=None)[source]¶

Build an adaptive LBVH radix tree, inferring bounds when not provided.

Parameters:

positions (Array) – Particle positions of shape (n, 3).
masses (Array) – Particle masses of shape (n,).
bounds (tuple[Array, Array] | None) – Optional (min_corner, max_corner) box; inferred from positions when omitted.
return_reordered (bool) – If True, also return the Morton-sorted positions/masses and the inverse permutation.
leaf_size (int) – Maximum number of particles per leaf.
workspace (RadixTreeWorkspace | None) – Optional reusable RadixTreeWorkspace to avoid reallocating scratch buffers across repeated builds.
return_workspace (bool) – If True, also return the (possibly newly allocated) workspace.
config (TreeBuildConfig | None) – Optional TreeBuildConfig; when given it overrides the equivalent individual keyword arguments.

Returns:

The tree, or a tuple additionally containing the reordered particle buffers and/or workspace when return_reordered / return_workspace are set.

Return type:

RadixTree or tuple

yggdrax.tree.build_tree_jit(positions, masses, bounds=None, *, return_reordered=False, leaf_size=8, workspace=None, return_workspace=False, config=None)[source]¶

JIT-compiled variant of build_tree() (see it for parameters/returns).

Parameters:

positions (Array)
masses (Array)
bounds (tuple[Array, Array] | None)
return_reordered (bool)
leaf_size (int)
workspace (RadixTreeWorkspace | None)
return_workspace (bool)
config (TreeBuildConfig | None)

yggdrax.tree.available_tree_types()[source]¶

Return registered public tree-type identifiers.

Return type:: tuple[str, …]

yggdrax.tree.get_level_offsets(tree, *, node_levels=None)[source]¶

Return level offsets, deriving compact level partitions when absent.

Parameters:

tree (object)
node_levels (Array | None)

Return type:

Array

yggdrax.tree.get_leaf_nodes(tree)[source]¶

Return leaf-node indices, deriving a stable default when needed.

Parameters:: tree (object)
Return type:: Array

yggdrax.tree.get_node_levels(tree)[source]¶

Return per-node depth levels, deriving from parent links when missing.

Parameters:: tree (object)
Return type:: Array

yggdrax.tree.get_nodes_by_level(tree, *, node_levels=None)[source]¶

Return nodes sorted by level (stable by node index within each level).

Parameters:

tree (object)
node_levels (Array | None)

Return type:

Array

yggdrax.tree.get_num_internal_nodes(tree)[source]¶

Return number of internal nodes, deriving it from child buffers when needed.

Parameters:: tree (object)
Return type:: int

yggdrax.tree.get_num_levels(tree, *, node_levels=None)[source]¶

Return tree depth count, deriving from node levels when needed.

Parameters:

tree (object)
node_levels (Array | None)

Return type:

int

yggdrax.tree.has_fmm_core_topology(tree_or_topology)[source]¶

Return True when all FMM-core fields are available.

Parameters:: tree_or_topology (object)
Return type:: bool

yggdrax.tree.has_fmm_topology(tree_or_topology)[source]¶

Alias of has_fmm_core_topology for compatibility.

Parameters:: tree_or_topology (object)
Return type:: bool

yggdrax.tree.has_leaf_topology(tree_or_topology)[source]¶

Return True when leaf-node metadata can be resolved.

Parameters:: tree_or_topology (object)
Return type:: bool

yggdrax.tree.has_morton_topology(tree_or_topology)[source]¶

Return True when all Morton-geometry fields are available.

Parameters:: tree_or_topology (object)
Return type:: bool

yggdrax.tree.missing_fmm_core_topology_fields(tree_or_topology)[source]¶

Return FMM-core required topology fields missing on the provided object.

Parameters:: tree_or_topology (object)
Return type:: tuple[str, …]

yggdrax.tree.missing_fmm_topology_fields(tree_or_topology)[source]¶

Alias of missing_fmm_core_topology_fields for compatibility.

Parameters:: tree_or_topology (object)
Return type:: tuple[str, …]

yggdrax.tree.missing_leaf_topology_fields(tree_or_topology)[source]¶

Return fields needed to derive or expose leaf-node indices.

Parameters:: tree_or_topology (object)
Return type:: tuple[str, …]

yggdrax.tree.missing_morton_topology_fields(tree_or_topology)[source]¶

Return Morton-geometry required fields missing on the provided object.

Parameters:: tree_or_topology (object)
Return type:: tuple[str, …]

yggdrax.tree.resolve_tree_topology(tree_or_topology)[source]¶

Return a topology payload from a tree container or topology object.

Parameters:: tree_or_topology (object) – A tree container (with a .topology attribute) or a topology object.
Returns:: The concrete topology payload. The return is typed to the broad FMM-core/Morton contract so downstream field accesses type-check; a given backend may expose additional fields beyond it.
Return type:: MortonLeafBoundsProtocol

yggdrax.tree.require_fmm_core_topology(tree_or_topology)[source]¶

Raise ValueError when FMM-core topology fields are missing.

Parameters:: tree_or_topology (object)
Return type:: None

yggdrax.tree.require_fmm_topology(tree_or_topology)[source]¶

Alias of require_fmm_core_topology for compatibility.

Parameters:: tree_or_topology (object)
Return type:: None

yggdrax.tree.require_leaf_topology(tree_or_topology)[source]¶

Raise ValueError when leaf-node metadata cannot be resolved.

Parameters:: tree_or_topology (object)
Return type:: None

yggdrax.tree.require_morton_topology(tree_or_topology)[source]¶

Raise ValueError when Morton-geometry fields are missing.

Parameters:: tree_or_topology (object)
Return type:: None

yggdrax.tree.register_tree_builder(tree_type, builder, *, overwrite=False)[source]¶

Parameters:

tree_type (str)
builder (Callable[[TreeBuildRequest], Tree])
overwrite (bool)

Return type:

None

yggdrax.tree.rebuild_static_radix_tree_from_template(positions, masses, template, *, bounds=None, return_reordered=False)[source]¶

Refresh particles using an existing static-radix data structure.

Parameters:

positions (Array)
masses (Array)
template (RadixTree | RadixTree)
bounds (tuple[Array, Array] | None)
return_reordered (bool)

yggdrax.tree.reorder_particles_by_indices(positions, masses, sorted_indices)[source]¶

Reorder particle arrays to Morton order using sorted_indices and return (positions_sorted, masses_sorted, inverse_perm).

Parameters:

positions (Array)
masses (Array)
sorted_indices (Array)

Return type:

tuple[Array, Array, Array]

Explicit octree metadata derived from Morton/radix topology.

class yggdrax.octree.ExplicitOctreeBoxGeometry(centers, half_extents, radii, max_extents)[source]¶

Bases: NamedTuple

Axis-aligned box geometry for explicit octree cells.

Parameters:

centers (Array)
half_extents (Array)
radii (Array)
max_extents (Array)

centers: Array¶: Alias for field number 0

half_extents: Array¶: Alias for field number 1

radii: Array¶: Alias for field number 2

max_extents: Array¶: Alias for field number 3

class yggdrax.octree.ExplicitOctreeMetadata(oct_parent, oct_children, oct_child_counts, oct_child_mask, oct_valid_mask, oct_node_codes, oct_node_depths, oct_node_ranges, oct_nodes_by_level, oct_level_offsets, oct_num_levels, oct_leaf_mask, oct_leaf_nodes, radix_node_to_oct, radix_leaf_to_oct)[source]¶

Bases: NamedTuple

Explicit octree buffers derived from octree leaf-cell partitions.

Parameters:

oct_parent (Array)
oct_children (Array)
oct_child_counts (Array)
oct_child_mask (Array)
oct_valid_mask (Array)
oct_node_codes (Array)
oct_node_depths (Array)
oct_node_ranges (Array)
oct_nodes_by_level (Array)
oct_level_offsets (Array)
oct_num_levels (Array)
oct_leaf_mask (Array)
oct_leaf_nodes (Array)
radix_node_to_oct (Array)
radix_leaf_to_oct (Array)

oct_parent: Array¶: Alias for field number 0

oct_children: Array¶: Alias for field number 1

oct_child_counts: Array¶: Alias for field number 2

oct_child_mask: Array¶: Alias for field number 3

oct_valid_mask: Array¶: Alias for field number 4

oct_node_codes: Array¶: Alias for field number 5

oct_node_depths: Array¶: Alias for field number 6

oct_node_ranges: Array¶: Alias for field number 7

oct_nodes_by_level: Array¶: Alias for field number 8

oct_level_offsets: Array¶: Alias for field number 9

oct_num_levels: Array¶: Alias for field number 10

oct_leaf_mask: Array¶: Alias for field number 11

oct_leaf_nodes: Array¶: Alias for field number 12

radix_node_to_oct: Array¶: Alias for field number 13

radix_leaf_to_oct: Array¶: Alias for field number 14

class yggdrax.octree.ExplicitOctreeTraversalView(valid_mask, parent, children, child_counts, node_codes, node_depths, node_ranges, nodes_by_level, level_offsets, num_levels, leaf_mask, leaf_nodes, radix_node_to_oct, radix_leaf_to_oct, oct_to_radix_node, oct_to_radix_leaf, num_valid_nodes, num_leaf_nodes, box_centers, box_half_extents, box_radii, box_max_extents)[source]¶

Bases: NamedTuple

Source-owned octree traversal/geometry package for native consumers.

Parameters:

valid_mask (Array)
parent (Array)
children (Array)
child_counts (Array)
node_codes (Array)
node_depths (Array)
node_ranges (Array)
nodes_by_level (Array)
level_offsets (Array)
num_levels (Array)
leaf_mask (Array)
leaf_nodes (Array)
radix_node_to_oct (Array)
radix_leaf_to_oct (Array)
oct_to_radix_node (Array)
oct_to_radix_leaf (Array)
num_valid_nodes (Array)
num_leaf_nodes (Array)
box_centers (Array)
box_half_extents (Array)
box_radii (Array)
box_max_extents (Array)

valid_mask: Array¶: Alias for field number 0

parent: Array¶: Alias for field number 1

children: Array¶: Alias for field number 2

child_counts: Array¶: Alias for field number 3

node_codes: Array¶: Alias for field number 4

node_depths: Array¶: Alias for field number 5

node_ranges: Array¶: Alias for field number 6

nodes_by_level: Array¶: Alias for field number 7

level_offsets: Array¶: Alias for field number 8

num_levels: Array¶: Alias for field number 9

leaf_mask: Array¶: Alias for field number 10

leaf_nodes: Array¶: Alias for field number 11

radix_node_to_oct: Array¶: Alias for field number 12

radix_leaf_to_oct: Array¶: Alias for field number 13

oct_to_radix_node: Array¶: Alias for field number 14

oct_to_radix_leaf: Array¶: Alias for field number 15

num_valid_nodes: Array¶: Alias for field number 16

num_leaf_nodes: Array¶: Alias for field number 17

box_centers: Array¶: Alias for field number 18

box_half_extents: Array¶: Alias for field number 19

box_radii: Array¶: Alias for field number 20

box_max_extents: Array¶: Alias for field number 21

class yggdrax.octree.OctreeTopology(parent, left_child, right_child, left_is_leaf, right_is_leaf, particle_indices, morton_codes, node_ranges, num_particles, num_internal_nodes, node_level, level_offsets, nodes_by_level, num_levels, bounds_min, bounds_max, leaf_codes, leaf_depths, use_morton_geometry, leaf_size, oct_parent, oct_children, oct_child_counts, oct_child_mask, oct_valid_mask, oct_node_codes, oct_node_depths, oct_node_ranges, oct_nodes_by_level, oct_level_offsets, oct_num_levels, oct_leaf_mask, oct_leaf_nodes, radix_node_to_oct, radix_leaf_to_oct)[source]¶

Bases: NamedTuple

Octree-native metadata plus shared compatibility topology fields.

Parameters:

parent (Array)
left_child (Array)
right_child (Array)
left_is_leaf (Array)
right_is_leaf (Array)
particle_indices (Array)
morton_codes (Array)
node_ranges (Array)
num_particles (int)
num_internal_nodes (int)
node_level (Array)
level_offsets (Array)
nodes_by_level (Array)
num_levels (Array)
bounds_min (Array)
bounds_max (Array)
leaf_codes (Array)
leaf_depths (Array)
use_morton_geometry (Array)
leaf_size (int | None)
oct_parent (Array)
oct_children (Array)
oct_child_counts (Array)
oct_child_mask (Array)
oct_valid_mask (Array)
oct_node_codes (Array)
oct_node_depths (Array)
oct_node_ranges (Array)
oct_nodes_by_level (Array)
oct_level_offsets (Array)
oct_num_levels (Array)
oct_leaf_mask (Array)
oct_leaf_nodes (Array)
radix_node_to_oct (Array)
radix_leaf_to_oct (Array)

parent: Array¶: Alias for field number 0

left_child: Array¶: Alias for field number 1

right_child: Array¶: Alias for field number 2

left_is_leaf: Array¶: Alias for field number 3

right_is_leaf: Array¶: Alias for field number 4

particle_indices: Array¶: Alias for field number 5

morton_codes: Array¶: Alias for field number 6

node_ranges: Array¶: Alias for field number 7

num_particles: int¶: Alias for field number 8

num_internal_nodes: int¶: Alias for field number 9

node_level: Array¶: Alias for field number 10

level_offsets: Array¶: Alias for field number 11

nodes_by_level: Array¶: Alias for field number 12

num_levels: Array¶: Alias for field number 13

bounds_min: Array¶: Alias for field number 14

bounds_max: Array¶: Alias for field number 15

leaf_codes: Array¶: Alias for field number 16

leaf_depths: Array¶: Alias for field number 17

use_morton_geometry: Array¶: Alias for field number 18

leaf_size: int | None¶: Alias for field number 19

oct_parent: Array¶: Alias for field number 20

oct_children: Array¶: Alias for field number 21

oct_child_counts: Array¶: Alias for field number 22

oct_child_mask: Array¶: Alias for field number 23

oct_valid_mask: Array¶: Alias for field number 24

oct_node_codes: Array¶: Alias for field number 25

oct_node_depths: Array¶: Alias for field number 26

oct_node_ranges: Array¶: Alias for field number 27

oct_nodes_by_level: Array¶: Alias for field number 28

oct_level_offsets: Array¶: Alias for field number 29

oct_num_levels: Array¶: Alias for field number 30

oct_leaf_mask: Array¶: Alias for field number 31

oct_leaf_nodes: Array¶: Alias for field number 32

radix_node_to_oct: Array¶: Alias for field number 33

radix_leaf_to_oct: Array¶: Alias for field number 34

yggdrax.octree.augment_radix_topology_with_octree(topology)[source]¶

Return an octree-augmented topology that preserves radix compatibility.

Attaches the explicit octree buffers (see build_explicit_octree_metadata()) to a radix topology without disturbing its FMM-core fields, so the result works with both the radix traversal core and octree-style level scheduling.

Parameters:: topology (object) – Radix topology exposing the Morton leaf contract.
Returns:: The topology extended with oct_* and radix_*_to_oct fields.
Return type:: OctreeTopology

yggdrax.octree.build_explicit_octree_traversal_view(topology)[source]¶

Package explicit octree structure, mappings, and box geometry.

Convenience view that bundles the octree cell tables, the radix<->octree mappings, and per-cell box geometry for downstream octree-native traversal.

Parameters:: topology (object) – Octree-augmented topology (see augment_radix_topology_with_octree()) exposing the oct_* fields.
Returns:: Bundled octree structure, mappings, and box geometry.
Return type:: ExplicitOctreeTraversalView
Raises:: ValueError – If topology is missing any required oct_* field.

yggdrax.octree.compute_explicit_octree_box_geometry(*, valid_mask, node_codes, node_depths, bounds_min, bounds_max)[source]¶

Compute explicit octree box geometry directly from code/depth pairs.

Parameters:

valid_mask (Array)
node_codes (Array)
node_depths (Array)
bounds_min (Array)
bounds_max (Array)

Return type:

ExplicitOctreeBoxGeometry

yggdrax.octree.build_explicit_octree_metadata(topology)[source]¶

Derive explicit octree cells and map the radix nodes onto them.

Builds the compressed set of octree cells (Morton leaves plus their consecutive-pair least-common-ancestors) and the mappings that link the radix/compat nodes to octree cells.

Parameters:: topology (object) – Radix topology exposing the Morton leaf contract.
Returns:: Octree cell tables (children, depths, ranges, level offsets) and the radix_node_to_oct / radix_leaf_to_oct mappings.
Return type:: ExplicitOctreeMetadata

Experimental KD-tree API with JAX-friendly exact query kernels.

class yggdrax.kdtree.KDTree(points, indices, particle_indices, node_start, node_end, node_ranges, parent, left_child, right_child, num_internal_nodes, num_particles, use_morton_geometry, split_dim, split_value, bbox_min, bbox_max, leaf_nodes, leaf_start, leaf_end, leaf_point_ids, leaf_valid_mask, node_to_leaf, leaf_size)[source]¶

Bases: object

KD-tree container plus precomputed topology metadata.

Parameters:

points (Array)
indices (Array)
particle_indices (Array)
node_start (Array)
node_end (Array)
node_ranges (Array)
parent (Array)
left_child (Array)
right_child (Array)
num_internal_nodes (int)
num_particles (Array)
use_morton_geometry (Array)
split_dim (Array)
split_value (Array)
bbox_min (Array)
bbox_max (Array)
leaf_nodes (Array)
leaf_start (Array)
leaf_end (Array)
leaf_point_ids (Array)
leaf_valid_mask (Array)
node_to_leaf (Array)
leaf_size (int)

points: Array¶

indices: Array¶

particle_indices: Array¶

node_start: Array¶

node_end: Array¶

node_ranges: Array¶

parent: Array¶

left_child: Array¶

right_child: Array¶

num_internal_nodes: int¶

num_particles: Array¶

use_morton_geometry: Array¶

split_dim: Array¶

split_value: Array¶

bbox_min: Array¶

bbox_max: Array¶

leaf_nodes: Array¶

leaf_start: Array¶

leaf_end: Array¶

leaf_point_ids: Array¶

leaf_valid_mask: Array¶

node_to_leaf: Array¶

leaf_size: int¶

property num_points: int¶: Return the number of reference points in the tree.

property dimension: int¶: Return spatial dimensionality of reference points.

yggdrax.kdtree.build_and_query(points, queries, *, k=1, leaf_size=32, backend='tiled', return_squared=False)[source]¶

Build a KD-tree and run nearest-neighbor queries in one call.

Convenience wrapper around build_kdtree() followed by query_neighbors().

Parameters:

points (Array) – Reference points of shape (n_points, dim).
queries (Array) – Query points of shape (n_queries, dim).
k (int) – Number of nearest neighbors to return per query.
leaf_size (int) – Maximum number of points per KD leaf bucket.
backend (Literal['dense', 'tiled', 'tree']) – Query backend passed to query_neighbors().
return_squared (bool) – If True, return squared Euclidean distances instead of distances.

Returns:

(indices, distances), each of shape (n_queries, k).

Return type:

tuple of Array

yggdrax.kdtree.build_kdtree(points, *, leaf_size=32)[source]¶

Build a median-split KD-tree container over reference points.

The returned container exposes both KD-specific split metadata and the FMM-core topology contract (parent/left_child/right_child/ node_ranges), so it can feed the same geometry and dual-tree traversal APIs as the radix and octree backends.

Parameters:

points (Array) – Reference points of shape (n_points, dim).
leaf_size (int) – Maximum number of points stored per leaf bucket; larger values yield a shallower tree.

Returns:

Immutable KD-tree container with topology and per-node bounding boxes.

Return type:

KDTree

Raises:

ValueError – If leaf_size is less than 1.

yggdrax.kdtree.count_neighbors(tree, queries, *, radius, include_self=True, backend='tiled', point_block_size=2048)[source]¶

Count reference points within a radius (or radii) of each query point.

Parameters:

tree (KDTree) – Reference-point container built by build_kdtree().
queries (Array) – Query points with shape (n_queries, dim).
radius (float | Array) – Search radius. A scalar counts within a single radius; an array counts within several radii per query and returns one column per radius.
include_self (bool) – If True, a query coinciding with a reference point counts itself.
backend (Literal['tiled', 'tree', 'auto']) – Counting backend: "tiled" (blockwise, lower memory), "tree" (KD pruning), or "auto" (chosen by problem size and device).
point_block_size (int) – Block size used by the tiled backend.

Returns:

Neighbor counts with shape (n_queries,) for a scalar radius or (n_queries, n_radii) when radius is an array.

Return type:

Array

Raises:

ValueError – If backend is not one of the supported values.

yggdrax.kdtree.query_neighbors(tree, queries, *, k=1, exclude_self=False, backend='tiled', point_block_size=4096, return_squared=False)[source]¶

Return nearest-neighbor indices and distances for each query point.

Parameters:

tree (KDTree) – Reference-point container built by build_kdtree().
queries (Array) – Query points with shape (n_queries, dim).
k (int) – Number of nearest neighbors to return per query.
exclude_self (bool) – If True and queries.shape[0] == tree.num_points, masks diagonal elements so a point is not returned as its own neighbor.
backend (Literal['dense', 'tiled', 'tree', 'auto']) – Query backend. "dense" builds a full pairwise matrix; "tiled" computes exact neighbors blockwise with lower memory; "tree" uses KD leaf bounding boxes to prune the exact search; "auto" selects a backend by problem size and device.
point_block_size (int) – Block size used by the tiled backend.
return_squared (bool) – If True, return squared Euclidean distances instead of distances.

Returns:

(indices, distances), each of shape (n_queries, k).

Return type:

tuple of Array

Raises:

ValueError – If k is out of range or backend is not a supported value.

Geometry¶

Public geometry API for Yggdrax.

class yggdrax.geometry.LevelMajorTreeGeometry(centers, half_extents, radii, max_extents, level_counts, node_indices)[source]¶

Bases: NamedTuple

Tree geometry reshaped into dense level-major buffers.

Parameters:

centers (Array)
half_extents (Array)
radii (Array)
max_extents (Array)
level_counts (Array)
node_indices (Array)

centers: Array¶: Alias for field number 0

half_extents: Array¶: Alias for field number 1

radii: Array¶: Alias for field number 2

max_extents: Array¶: Alias for field number 3

level_counts: Array¶: Alias for field number 4

node_indices: Array¶: Alias for field number 5

class yggdrax.geometry.TreeGeometry(center, half_extent, radius, max_extent)[source]¶

Bases: NamedTuple

Bounding information for every node in a radix tree.

Parameters:

center (Array)
half_extent (Array)
radius (Array)
max_extent (Array)

center: Array¶: Alias for field number 0

half_extent: Array¶: Alias for field number 1

radius: Array¶: Alias for field number 2

max_extent: Array¶: Alias for field number 3

yggdrax.geometry.compute_tree_geometry(tree, positions_sorted, *, max_leaf_size=None)[source]¶

Compute per-node geometric bounds and helper radii.

Produces, for every node, its bounding-box center and half-extent, the box max half-extent (L-infinity radius), and the bounding-sphere radius, from Morton-sorted particle positions.

Parameters:

tree (object) – Tree container or topology exposing the FMM-core contract.
positions_sorted (Array) – Particle positions in the tree’s Morton order (i.e. indexed by tree.particle_indices), shape (n_particles, 3).
max_leaf_size (int | None) – Optional cap on the temporary leaf-gather buffer. Optional for correctness, but important for large JIT-compiled radix trees: it bounds the staging shape used during construction and avoids falling back to a num_particles-sized buffer under tracing. Defaults to the tree’s leaf_size when available.

Returns:

Per-node center, half_extent, max_extent, and radius.

Return type:

TreeGeometry

yggdrax.geometry.geometry_to_level_major(tree, geometry)[source]¶

Convert per-node geometry to padded level-major buffers.

Regroups the flat per-node geometry into (num_levels, max_nodes_per_level) padded arrays, so batched kernels can process one tree level at a time.

Parameters:

tree (object) – Tree container or topology exposing level-order metadata.
geometry (TreeGeometry) – Per-node geometry from compute_tree_geometry() for tree.

Returns:

Level-major padded geometry buffers and per-level node counts.

Return type:

LevelMajorTreeGeometry

Interactions and traversal¶

Public interaction/traversal API for Yggdrax.

class yggdrax.interactions.DualTreeRetryEvent(attempt, queue_capacity, interaction_capacity, status, far_pair_count, near_pair_count)[source]¶

Bases: NamedTuple

Metadata describing a single dual-tree retry attempt.

Parameters:

attempt (int)
queue_capacity (int)
interaction_capacity (int)
status (str)
far_pair_count (int)
near_pair_count (int)

attempt: int¶: Alias for field number 0

queue_capacity: int¶: Alias for field number 1

interaction_capacity: int¶: Alias for field number 2

status: str¶: Alias for field number 3

far_pair_count: int¶: Alias for field number 4

near_pair_count: int¶: Alias for field number 5

class yggdrax.interactions.DualTreeTraversalConfig(max_pair_queue, process_block, max_interactions_per_node, max_neighbors_per_leaf=2048)[source]¶

Bases: object

Fixed traversal parameters for the dual-tree walk.

Parameters:

max_pair_queue (int)
process_block (int)
max_interactions_per_node (int)
max_neighbors_per_leaf (int)

max_pair_queue: int¶

process_block: int¶

max_interactions_per_node: int¶

max_neighbors_per_leaf: int = 2048¶

class yggdrax.interactions.DualTreeWalkResult(interaction_offsets, interaction_sources, interaction_targets, interaction_tags, interaction_counts, neighbor_offsets, neighbor_indices, neighbor_counts, leaf_indices, far_pair_count, near_pair_count, queue_overflow, far_overflow, near_overflow, accept_decisions, near_decisions, refine_decisions)[source]¶

Bases: NamedTuple

Container for far-field and near-field results from a dual walk.

Parameters:

interaction_offsets (Array)
interaction_sources (Array)
interaction_targets (Array)
interaction_tags (Array)
interaction_counts (Array)
neighbor_offsets (Array)
neighbor_indices (Array)
neighbor_counts (Array)
leaf_indices (Array)
far_pair_count (Array)
near_pair_count (Array)
queue_overflow (Array)
far_overflow (Array)
near_overflow (Array)
accept_decisions (Array)
near_decisions (Array)
refine_decisions (Array)

interaction_offsets: Array¶: Alias for field number 0

interaction_sources: Array¶: Alias for field number 1

interaction_targets: Array¶: Alias for field number 2

interaction_tags: Array¶: Alias for field number 3

interaction_counts: Array¶: Alias for field number 4

neighbor_offsets: Array¶: Alias for field number 5

neighbor_indices: Array¶: Alias for field number 6

neighbor_counts: Array¶: Alias for field number 7

leaf_indices: Array¶: Alias for field number 8

far_pair_count: Array¶: Alias for field number 9

near_pair_count: Array¶: Alias for field number 10

queue_overflow: Array¶: Alias for field number 11

far_overflow: Array¶: Alias for field number 12

near_overflow: Array¶: Alias for field number 13

accept_decisions: Array¶: Alias for field number 14

near_decisions: Array¶: Alias for field number 15

refine_decisions: Array¶: Alias for field number 16

class yggdrax.interactions.CompactTaggedFarPairs(sources, targets, tags, far_pair_count)[source]¶

Bases: NamedTuple

Compact far-pair payload plus the number of active entries.

When built under tracing, sources/targets/tags may be fixed-capacity padded arrays. far_pair_count records how many leading entries are valid.

Parameters:

sources (Array)
targets (Array)
tags (Array)
far_pair_count (Array)

sources: Array¶: Alias for field number 0

targets: Array¶: Alias for field number 1

tags: Array¶: Alias for field number 2

far_pair_count: Array¶: Alias for field number 3

class yggdrax.interactions.CompactTaggedOctreeFarPairs(sources, targets, tags)[source]¶

Bases: NamedTuple

Exact-length far-pair payload stored in explicit octree node space.

Parameters:

sources (Array)
targets (Array)
tags (Array)

sources: Array¶: Alias for field number 0

targets: Array¶: Alias for field number 1

tags: Array¶: Alias for field number 2

class yggdrax.interactions.OctreeNativeNeighborList(offsets, neighbors, leaf_indices, counts, particle_order_leaf_indices, particle_order_to_native_leaf, neighbor_leaf_positions, target_block_leaf_ids, target_block_source_leaf_ids, target_block_valid_mask, target_block_offsets, target_block_size)[source]¶

Bases: NamedTuple

Compressed near-field neighbor list stored in explicit octree leaf space.

Parameters:

offsets (Array)
neighbors (Array)
leaf_indices (Array)
counts (Array)
particle_order_leaf_indices (Array)
particle_order_to_native_leaf (Array)
neighbor_leaf_positions (Array)
target_block_leaf_ids (Array)
target_block_source_leaf_ids (Array)
target_block_valid_mask (Array)
target_block_offsets (Array)
target_block_size (int)

offsets: Array¶: Alias for field number 0

neighbors: Array¶: Alias for field number 1

leaf_indices: Array¶: Alias for field number 2

counts: Array¶: Alias for field number 3

particle_order_leaf_indices: Array¶: Alias for field number 4

particle_order_to_native_leaf: Array¶: Alias for field number 5

neighbor_leaf_positions: Array¶: Alias for field number 6

target_block_leaf_ids: Array¶: Alias for field number 7

target_block_source_leaf_ids: Array¶: Alias for field number 8

target_block_valid_mask: Array¶: Alias for field number 9

target_block_offsets: Array¶: Alias for field number 10

target_block_size: int¶: Alias for field number 11

class yggdrax.interactions.NodeInteractionList(offsets, sources, targets, counts, level_offsets, target_levels)[source]¶

Bases: NamedTuple

Compressed far-field interaction list for all tree nodes.

Variables:

offsets (jax.Array) – Start index for each node within sources/targets; combine with counts to recover the per-node slice length.
sources (jax.Array) – Source node indices for every far-field pair, in level-major order.
targets (jax.Array) – Target node indices matching sources.
counts (jax.Array) – Interaction counts per node (entries per target).
level_offsets (jax.Array) – Prefix offsets delimiting the interaction ranges for each tree level (length num_levels + 1).
target_levels (jax.Array) – Tree level for each pair (monotonically non-decreasing).

Parameters:

offsets (Array)
sources (Array)
targets (Array)
counts (Array)
level_offsets (Array)
target_levels (Array)

offsets: Array¶: Alias for field number 0

sources: Array¶: Alias for field number 1

targets: Array¶: Alias for field number 2

counts: Array¶: Alias for field number 3

level_offsets: Array¶: Alias for field number 4

target_levels: Array¶: Alias for field number 5

class yggdrax.interactions.NodeNeighborList(offsets, neighbors, leaf_indices, counts, particle_order_leaf_indices, particle_order_to_native_leaf, neighbor_leaf_positions, target_block_leaf_ids, target_block_source_leaf_ids, target_block_valid_mask, target_block_offsets, target_block_size)[source]¶

Bases: NamedTuple

Compressed near-field neighbor list for leaf nodes.

Parameters:

offsets (Array)
neighbors (Array)
leaf_indices (Array)
counts (Array)
particle_order_leaf_indices (Array)
particle_order_to_native_leaf (Array)
neighbor_leaf_positions (Array)
target_block_leaf_ids (Array)
target_block_source_leaf_ids (Array)
target_block_valid_mask (Array)
target_block_offsets (Array)
target_block_size (int)

offsets: Array¶: Alias for field number 0

neighbors: Array¶: Alias for field number 1

leaf_indices: Array¶: Alias for field number 2

counts: Array¶: Alias for field number 3

particle_order_leaf_indices: Array¶: Alias for field number 4

particle_order_to_native_leaf: Array¶: Alias for field number 5

neighbor_leaf_positions: Array¶: Alias for field number 6

target_block_leaf_ids: Array¶: Alias for field number 7

target_block_source_leaf_ids: Array¶: Alias for field number 8

target_block_valid_mask: Array¶: Alias for field number 9

target_block_offsets: Array¶: Alias for field number 10

target_block_size: int¶: Alias for field number 11

yggdrax.interactions.build_interactions_and_neighbors(tree, geometry, theta=0.5, max_interactions_per_node=None, max_neighbors_per_leaf=2048, max_pair_queue=None, process_block=None, traversal_config=None, retry_logger=None, mac_type='bh', dehnen_radius_scale=1.0, pair_policy=None, policy_state=None, *, return_result=False, return_compact_far_pairs=False, return_interactions=True, return_grouped=False)[source]¶

Construct both far-field interactions and near-field neighbors.

Runs a single dual-tree traversal that classifies node/leaf pairs into well-separated far-field interactions (M2L) and near-field neighbor pairs (P2P), using the multipole acceptance criterion selected by mac_type (or a custom pair_policy).

When pair_policy is provided it overrides the built-in MAC decision for each candidate pair and may attach integer tags to accepted far pairs. Policies are evaluated in both directions; a pair is accepted only when both directed decisions accept, and the directed tags are exposed on DualTreeWalkResult.interaction_tags when return_result=True.

Fixed-capacity buffers are grown automatically (subject to internal caps) when the traversal reports overflow, so the capacity arguments are hints rather than hard limits unless they are pinned via traversal_config.

Parameters:

tree (object) – Tree container exposing FMM-core topology (radix, octree, or kd-tree).
geometry (TreeGeometry) – Per-node geometry from compute_tree_geometry() for tree.
theta (float) – Opening-angle parameter of the multipole acceptance criterion. Smaller values accept fewer far pairs (more accurate, more work).
max_interactions_per_node (int | None) – Capacity of the per-node far-interaction buffer. None auto-sizes it and grows on overflow.
max_neighbors_per_leaf (int) – Capacity of the per-leaf near-neighbor buffer.
max_pair_queue (int | None) – Capacity of the traversal wavefront queue. None auto-sizes it.
process_block (int | None) – Number of pairs processed per traversal iteration. None auto-sizes.
traversal_config (DualTreeTraversalConfig | None) – Bundled capacity/queue/block settings. When given, it takes precedence over the equivalent individual keyword arguments.
retry_logger (Callable[[DualTreeRetryEvent], object] | None) – Optional callable invoked with a DualTreeRetryEvent on each capacity-driven retry (for diagnostics/tuning).
mac_type (Literal['bh', 'engblom', 'dehnen']) – Multipole acceptance criterion: "bh" (Barnes-Hut opening angle), "dehnen", or "engblom".
dehnen_radius_scale (float) – Effective-radius scale applied for the Dehnen MAC.
pair_policy (Callable[[...], tuple[Array, Array]] | None) – Optional JAX-traceable callable overriding the built-in MAC decision.
policy_state (object) – Opaque state threaded to pair_policy on every evaluation.
return_result (bool) – If True, also return the raw DualTreeWalkResult (including interaction_tags).
return_compact_far_pairs (bool) – If True, also return exact-length CompactTaggedFarPairs.
return_interactions (bool) – If True (default), include the sparse NodeInteractionList in the result.
return_grouped (bool) – If True, also return displacement-grouped interaction buffers.

Returns:

By default (interactions, neighbors) where interactions is a NodeInteractionList and neighbors is a NodeNeighborList. Additional elements are appended, in declaration order, for each enabled return_* flag (return_compact_far_pairs, return_grouped, return_result).

Return type:

tuple

yggdrax.interactions.build_interactions_and_neighbors_split(tree, geometry, theta=0.5, max_interactions_per_node=None, max_neighbors_per_leaf=2048, max_pair_queue=None, process_block=None, traversal_config=None, retry_logger=None, mac_type='bh', dehnen_radius_scale=1.0, pair_policy=None, policy_state=None)[source]¶

Construct far and near products using two separate dual-tree walks.

Same outputs as build_interactions_and_neighbors() but runs the far-field and near-field collection as independent walks instead of one shared traversal.

Parameters:

tree (object) – Tree container or topology exposing the FMM-core contract.
geometry (TreeGeometry) – Per-node geometry from compute_tree_geometry() for tree.
theta (float) – Opening-angle parameter of the multipole acceptance criterion.
max_interactions_per_node (int | None) – Per-node far-interaction capacity; auto-sized when None.
max_neighbors_per_leaf (int) – Per-leaf near-neighbor capacity.
max_pair_queue (int | None) – Traversal wavefront capacity; auto-sized when None.
process_block (int | None) – Pairs processed per traversal iteration; auto-sized when None.
traversal_config (DualTreeTraversalConfig | None) – Bundled traversal settings; overrides the individual arguments.
retry_logger (Callable[[DualTreeRetryEvent], object] | None) – Optional callable invoked on each capacity-driven retry.
mac_type (Literal['bh', 'engblom', 'dehnen']) – MAC variant: "bh", "dehnen", or "engblom".
dehnen_radius_scale (float) – Effective-radius scale applied for the Dehnen MAC.
pair_policy (Callable[[...], tuple[Array, Array]] | None) – Optional JAX-traceable callable overriding the built-in MAC.
policy_state (object) – Opaque state threaded to pair_policy.

Returns:

(interactions, neighbors) as (NodeInteractionList, NodeNeighborList).

Return type:

tuple

yggdrax.interactions.build_compact_far_pairs(tree, geometry, theta=0.5, max_interactions_per_node=None, mac_type='bh', *, pair_policy=None, policy_state=None, max_pair_queue=None, process_block=None, traversal_config=None, retry_logger=None, dehnen_radius_scale=1.0)[source]¶

Construct exact-length tagged far pairs from the dual-tree walk.

Returns far-field pairs packed to their exact count (rather than the padded per-node layout of build_well_separated_interactions()).

Parameters:

tree (object) – Tree container or topology exposing the FMM-core contract.
geometry (TreeGeometry) – Per-node geometry from compute_tree_geometry() for tree.
theta (float) – Opening-angle parameter of the multipole acceptance criterion.
max_interactions_per_node (int | None) – Per-node far-interaction capacity; auto-sized when None.
mac_type (Literal['bh', 'engblom', 'dehnen']) – MAC variant: "bh", "dehnen", or "engblom".
pair_policy (Callable[[...], tuple[Array, Array]] | None) – Optional JAX-traceable callable overriding the built-in MAC.
policy_state (object) – Opaque state threaded to pair_policy.
max_pair_queue (int | None) – Traversal wavefront capacity; auto-sized when None.
process_block (int | None) – Pairs processed per traversal iteration; auto-sized when None.
traversal_config (DualTreeTraversalConfig | None) – Bundled traversal settings; overrides the individual arguments.
retry_logger (Callable[[DualTreeRetryEvent], object] | None) – Optional callable invoked on each capacity-driven retry.
dehnen_radius_scale (float) – Effective-radius scale applied for the Dehnen MAC.

Returns:

Exact-length far pairs with per-pair tags.

Return type:

CompactTaggedFarPairs

yggdrax.interactions.build_compact_far_pairs_and_leaf_neighbor_lists(tree, geometry, theta=0.5, max_neighbors_per_leaf=2048, max_interactions_per_node=None, mac_type='bh', *, pair_policy=None, policy_state=None, max_pair_queue=None, process_block=None, traversal_config=None, retry_logger=None, dehnen_radius_scale=1.0, timing_callback=None, compact_far_pair_capacity=None)[source]¶

Construct compact far pairs and near neighbors from one count walk.

Produces both exact-length far pairs and the near-field neighbor list while sharing a single bounded count pass (cheaper than two separate walks).

Parameters:

tree (object) – Tree container or topology exposing the FMM-core contract.
geometry (TreeGeometry) – Per-node geometry from compute_tree_geometry() for tree.
theta (float) – Opening-angle parameter of the multipole acceptance criterion.
max_neighbors_per_leaf (int) – Per-leaf near-neighbor capacity.
max_interactions_per_node (int | None) – Per-node far-interaction capacity; auto-sized when None.
mac_type (Literal['bh', 'engblom', 'dehnen']) – MAC variant: "bh", "dehnen", or "engblom".
pair_policy (Callable[[...], tuple[Array, Array]] | None) – Optional JAX-traceable callable overriding the built-in MAC.
policy_state (object) – Opaque state threaded to pair_policy.
max_pair_queue (int | None) – Traversal wavefront capacity; auto-sized when None.
process_block (int | None) – Pairs processed per traversal iteration; auto-sized when None.
traversal_config (DualTreeTraversalConfig | None) – Bundled traversal settings; overrides the individual arguments.
retry_logger (Callable[[DualTreeRetryEvent], object] | None) – Optional callable invoked on each capacity-driven retry.
dehnen_radius_scale (float) – Effective-radius scale applied for the Dehnen MAC.
timing_callback (Callable[[...], object] | None) – Optional callable invoked with per-stage timing diagnostics.
compact_far_pair_capacity (int | None) – Optional explicit capacity for the compact far-pair buffer.

Returns:

(compact_far_pairs, neighbors) as (CompactTaggedFarPairs, NodeNeighborList).

Return type:

tuple

yggdrax.interactions.build_octree_native_far_pairs(tree, geometry, theta=0.5, max_interactions_per_node=None, mac_type='bh', *, pair_policy=None, policy_state=None, max_pair_queue=None, process_block=None, traversal_config=None, retry_logger=None, dehnen_radius_scale=1.0)[source]¶

Construct exact-length far-field pairs in explicit octree node space.

Like build_compact_far_pairs() but emits pairs indexed in the explicit octree cell space (requires an octree-augmented tree).

Parameters:

tree (object) – Octree-augmented tree exposing the explicit oct_* buffers.
geometry (TreeGeometry) – Per-node geometry from compute_tree_geometry() for tree.
theta (float) – Opening-angle parameter of the multipole acceptance criterion.
max_interactions_per_node (int | None) – Per-node far-interaction capacity; auto-sized when None.
mac_type (Literal['bh', 'engblom', 'dehnen']) – MAC variant: "bh", "dehnen", or "engblom".
pair_policy (Callable[[...], tuple[Array, Array]] | None) – Optional JAX-traceable callable overriding the built-in MAC.
policy_state (object) – Opaque state threaded to pair_policy.
max_pair_queue (int | None) – Traversal wavefront capacity; auto-sized when None.
process_block (int | None) – Pairs processed per traversal iteration; auto-sized when None.
traversal_config (DualTreeTraversalConfig | None) – Bundled traversal settings; overrides the individual arguments.
retry_logger (Callable[[DualTreeRetryEvent], object] | None) – Optional callable invoked on each capacity-driven retry.
dehnen_radius_scale (float) – Effective-radius scale applied for the Dehnen MAC.

Returns:

Exact-length far pairs in octree node space.

Return type:

CompactTaggedOctreeFarPairs

yggdrax.interactions.build_octree_native_neighbor_lists(tree, geometry, theta=0.5, max_neighbors_per_leaf=2048, max_interactions_per_node=None, mac_type='bh', *, pair_policy=None, policy_state=None, max_pair_queue=None, process_block=None, traversal_config=None, retry_logger=None, dehnen_radius_scale=1.0)[source]¶

Construct exact-length near neighbors in explicit octree leaf space.

Like build_leaf_neighbor_lists() but emits neighbor pairs indexed in the explicit octree leaf space (requires an octree-augmented tree).

Parameters:

tree (object) – Octree-augmented tree exposing the explicit oct_* buffers.
geometry (TreeGeometry) – Per-node geometry from compute_tree_geometry() for tree.
theta (float) – Opening-angle parameter of the multipole acceptance criterion.
max_neighbors_per_leaf (int) – Per-leaf near-neighbor capacity.
max_interactions_per_node (int | None) – Per-node far-interaction capacity used during the walk; auto-sized when None.
mac_type (Literal['bh', 'engblom', 'dehnen']) – MAC variant: "bh", "dehnen", or "engblom".
pair_policy (Callable[[...], tuple[Array, Array]] | None) – Optional JAX-traceable callable overriding the built-in MAC.
policy_state (object) – Opaque state threaded to pair_policy.
max_pair_queue (int | None) – Traversal wavefront capacity; auto-sized when None.
process_block (int | None) – Pairs processed per traversal iteration; auto-sized when None.
traversal_config (DualTreeTraversalConfig | None) – Bundled traversal settings; overrides the individual arguments.
retry_logger (Callable[[DualTreeRetryEvent], object] | None) – Optional callable invoked on each capacity-driven retry.
dehnen_radius_scale (float) – Effective-radius scale applied for the Dehnen MAC.

Returns:

Exact-length near neighbors in octree leaf space.

Return type:

OctreeNativeNeighborList

yggdrax.interactions.build_leaf_neighbor_lists(tree, geometry, theta=0.5, max_neighbors_per_leaf=2048, max_interactions_per_node=None, mac_type='bh', *, pair_policy=None, policy_state=None, max_pair_queue=None, process_block=None, traversal_config=None, retry_logger=None, dehnen_radius_scale=1.0)[source]¶

Construct only the near-field neighbor list from a dual-tree walk.

Like build_interactions_and_neighbors() but returns the leaf-leaf near-field (P2P) neighbor list alone, skipping far-field collection.

Parameters:

tree (object) – Tree container or topology exposing the FMM-core contract.
geometry (TreeGeometry) – Per-node geometry from compute_tree_geometry() for tree.
theta (float) – Opening-angle parameter of the multipole acceptance criterion.
max_neighbors_per_leaf (int) – Per-leaf near-neighbor capacity.
max_interactions_per_node (int | None) – Per-node far-interaction capacity used during the walk; auto-sized when None.
mac_type (Literal['bh', 'engblom', 'dehnen']) – MAC variant: "bh", "dehnen", or "engblom".
pair_policy (Callable[[...], tuple[Array, Array]] | None) – Optional JAX-traceable callable overriding the built-in MAC.
policy_state (object) – Opaque state threaded to pair_policy.
max_pair_queue (int | None) – Traversal wavefront capacity; auto-sized when None.
process_block (int | None) – Pairs processed per traversal iteration; auto-sized when None.
traversal_config (DualTreeTraversalConfig | None) – Bundled traversal settings; overrides the individual arguments.
retry_logger (Callable[[DualTreeRetryEvent], object] | None) – Optional callable invoked on each capacity-driven retry.
dehnen_radius_scale (float) – Effective-radius scale applied for the Dehnen MAC.

Returns:

Near-field leaf neighbor list.

Return type:

NodeNeighborList

yggdrax.interactions.build_grouped_interactions_from_pairs(tree, geometry, interaction_sources, interaction_targets, *, level_offsets=None)[source]¶

Group far-field pairs into displacement classes for class-major M2L.

Thin wrapper over yggdrax.grouped_interactions.build_grouped_interactions_from_pairs() that first resolves tree to its topology payload.

Parameters:

tree (object) – Tree container or topology exposing level-order metadata.
geometry (TreeGeometry) – Per-node geometry from compute_tree_geometry() for tree.
interaction_sources (Array) – Far-pair source node indices, shape (num_pairs,).
interaction_targets (Array) – Far-pair target node indices, shape (num_pairs,).
level_offsets (Array | None) – Optional precomputed level offsets passed through for diagnostics; derived from tree when omitted.

Returns:

Class keys, per-class displacements, and CSR-style class offsets/ids.

Return type:

GroupedInteractionBuffers

yggdrax.interactions.build_well_separated_interactions(tree, geometry, theta=0.5, max_interactions_per_node=None, mac_type='bh', *, pair_policy=None, policy_state=None, max_pair_queue=None, process_block=None, traversal_config=None, retry_logger=None, dehnen_radius_scale=1.0)[source]¶

Construct only the far-field interaction list from a dual-tree walk.

Like build_interactions_and_neighbors() but returns the sparse far-field (M2L) list alone, skipping near-field neighbor collection.

Parameters:

tree (object) – Tree container or topology exposing the FMM-core contract.
geometry (TreeGeometry) – Per-node geometry from compute_tree_geometry() for tree.
theta (float) – Opening-angle parameter of the multipole acceptance criterion.
max_interactions_per_node (int | None) – Capacity of the per-node far-interaction buffer; auto-sized when None.
mac_type (Literal['bh', 'engblom', 'dehnen']) – MAC variant: "bh", "dehnen", or "engblom".
pair_policy (Callable[[...], tuple[Array, Array]] | None) – Optional JAX-traceable callable overriding the built-in MAC.
policy_state (object) – Opaque state threaded to pair_policy.
max_pair_queue (int | None) – Traversal wavefront capacity; auto-sized when None.
process_block (int | None) – Pairs processed per traversal iteration; auto-sized when None.
traversal_config (DualTreeTraversalConfig | None) – Bundled traversal settings; overrides the individual arguments.
retry_logger (Callable[[DualTreeRetryEvent], object] | None) – Optional callable invoked on each capacity-driven retry.
dehnen_radius_scale (float) – Effective-radius scale applied for the Dehnen MAC.

Returns:

Sparse far-field interaction list.

Return type:

NodeInteractionList

yggdrax.interactions.diagnose_leaf_neighbor_growth(tree, geometry, theta=0.5, *, max_neighbors_per_leaf=2048, top_k=10, sample_neighbors=20)[source]¶

Return a compact report of the highest per-leaf neighbor counts.

Diagnostic helper for tuning max_neighbors_per_leaf: runs the near-field walk and summarizes which leaves accumulate the most neighbors.

Parameters:

tree (object) – Tree container or topology exposing the FMM-core contract.
geometry (TreeGeometry) – Per-node geometry from compute_tree_geometry() for tree.
theta (float) – Opening-angle parameter of the multipole acceptance criterion.
max_neighbors_per_leaf (int) – Per-leaf near-neighbor capacity used for the probe walk.
top_k (int) – Number of highest-count leaves to report.
sample_neighbors (int) – Number of sample neighbor ids to include per reported leaf.

Returns:

Report with the top per-leaf neighbor counts and sampled neighbor ids.

Return type:

dict

yggdrax.interactions.interactions_for_node(data, node)[source]¶

Return source node indices interacting with node.

Parameters:

data (NodeInteractionList)
node (int)

Return type:

Array

yggdrax.interactions.neighbors_for_leaf(data, leaf_node)[source]¶

Return neighbouring leaf nodes interacting with leaf_node.

Parameters:

data (NodeNeighborList)
leaf_node (int)

Return type:

Array

Prepared tree artifact construction for Yggdrax.

yggdrax.traversal.build_prepared_tree_artifacts(positions, masses, bounds=None, *, tree_type='radix', leaf_size=16, theta=0.6, mac_type='bh', traversal_config=None, dehnen_radius_scale=1.0)[source]¶

Build a tree, geometry, and interactions in one call.

End-to-end convenience: reorders particles into a tree, computes per-node geometry, and runs the dual-tree walk, returning everything a downstream solver needs in a single bundle with stable public field names. A KD-tree-tuned traversal configuration is chosen automatically when traversal_config is omitted and tree_type == "kdtree".

Parameters:

positions (Array) – Particle positions of shape (n, 3).
masses (Array) – Particle masses of shape (n,).
bounds (tuple[Array, Array] | None) – Optional (min_corner, max_corner) box; inferred from positions when omitted.
tree_type (Literal['radix', 'octree', 'kdtree']) – Backend to build: "radix", "octree", or "kdtree".
leaf_size (int) – Maximum particles per leaf.
theta (float) – Opening-angle parameter of the multipole acceptance criterion.
mac_type (Literal['bh', 'engblom', 'dehnen']) – MAC variant: "bh", "dehnen", or "engblom".
traversal_config (DualTreeTraversalConfig | None) – Optional traversal capacity/queue settings; a KD-tree-tuned default is used when omitted for KD-trees.
dehnen_radius_scale (float) – Effective-radius scale applied for the Dehnen MAC.

Returns:

Bundle of the tree, reordered particle buffers, inverse permutation, geometry, far-field interactions, near-field neighbors, and the raw traversal result.

Return type:

PreparedTreeArtifacts

Raises:

RuntimeError – If the tree build did not return reordered particle buffers.

Dense level-major interaction buffers for the downward sweep.

class yggdrax.dense_interactions.DenseInteractionBuffers(geometry, m2l_sources, m2l_displacements, m2l_mask, m2l_counts, sparse_interactions)[source]¶

Bases: NamedTuple

Dense level-major representation of far-field interaction pairs.

Parameters:

geometry (LevelMajorTreeGeometry)
m2l_sources (Array)
m2l_displacements (Array)
m2l_mask (Array)
m2l_counts (Array)
sparse_interactions (NodeInteractionList)

geometry: LevelMajorTreeGeometry¶: Alias for field number 0

m2l_sources: Array¶: Alias for field number 1

m2l_displacements: Array¶: Alias for field number 2

m2l_mask: Array¶: Alias for field number 3

m2l_counts: Array¶: Alias for field number 4

sparse_interactions: NodeInteractionList¶: Alias for field number 5

yggdrax.dense_interactions.build_dense_interactions(tree, geometry, theta=0.5, *, max_pair_queue=None, process_block=None, traversal_config=None, retry_logger=None)[source]¶

Build sparse far-field interactions and return their dense view.

Convenience wrapper: runs build_interactions_and_neighbors() and then densify_interactions().

Parameters:

tree (object) – Tree container or topology exposing the FMM-core contract.
geometry (TreeGeometry) – Per-node geometry from compute_tree_geometry() for tree.
theta (float) – Opening-angle parameter of the multipole acceptance criterion.
max_pair_queue (int | None) – Traversal wavefront capacity; auto-sized when None.
process_block (int | None) – Pairs processed per traversal iteration; auto-sized when None.
traversal_config (DualTreeTraversalConfig | None) – Bundled traversal settings; overrides the individual arguments.
retry_logger (Callable[[DualTreeRetryEvent], None] | None) – Optional callable invoked on each capacity-driven retry.

Returns:

Level-major M2L buffers plus the original sparse list.

Return type:

DenseInteractionBuffers

yggdrax.dense_interactions.densify_interactions(tree, geometry, interactions)[source]¶

Convert sparse far-field interactions into dense level-major tensors.

Regroups a NodeInteractionList into padded (num_levels, max_nodes, max_interactions) buffers (source ids, validity mask, counts, and target-minus-source displacements) so a batched M2L kernel can process one tree level at a time.

Parameters:

tree (object) – Tree container or topology exposing the FMM-core contract.
geometry (TreeGeometry) – Per-node geometry from compute_tree_geometry() for tree.
interactions (NodeInteractionList) – Sparse far-field list from build_interactions_and_neighbors().

Returns:

Level-major M2L buffers plus the original sparse list.

Return type:

DenseInteractionBuffers

Grouped interaction utilities for class-major far-field execution.

class yggdrax.grouped_interactions.GroupedInteractionBuffers(class_keys, class_displacements, class_offsets, class_sources, class_targets, class_ids, level_offsets, level_nodes)[source]¶

Bases: NamedTuple

Class-major grouped representation of far-field interactions.

Parameters:

class_keys (Array)
class_displacements (Array)
class_offsets (Array)
class_sources (Array)
class_targets (Array)
class_ids (Array)
level_offsets (Array)
level_nodes (Array)

class_keys: Array¶: Alias for field number 0

class_displacements: Array¶: Alias for field number 1

class_offsets: Array¶: Alias for field number 2

class_sources: Array¶: Alias for field number 3

class_targets: Array¶: Alias for field number 4

class_ids: Array¶: Alias for field number 5

level_offsets: Array¶: Alias for field number 6

level_nodes: Array¶: Alias for field number 7

yggdrax.grouped_interactions.build_grouped_interactions(tree, geometry, interactions)[source]¶

Group sparse far-field pairs into displacement classes.

Host-side (NumPy) grouping of a NodeInteractionList into classes keyed by (target_level, source_level, displacement xyz) so a solver can execute one translation stencil per class (“class-major” M2L).

Parameters:

tree (object) – Tree container or topology exposing level-order metadata.
geometry (TreeGeometry) – Per-node geometry from compute_tree_geometry() for tree.
interactions (NodeInteractionList) – Sparse far-field list with sources/targets arrays.

Returns:

Class keys, per-class displacements, and CSR-style class offsets/ids.

Return type:

GroupedInteractionBuffers

yggdrax.grouped_interactions.build_grouped_interactions_from_pairs(tree, geometry, sources, targets, *, level_offsets=None)[source]¶

Group far-field source/target node arrays into displacement classes.

Lower-level entry point for build_grouped_interactions() that takes the source/target node-index arrays directly. Runs eagerly on the host (NumPy lexsort + boundary detection); not JAX-traceable.

Parameters:

tree (object) – Tree container or topology exposing level-order metadata.
geometry (TreeGeometry) – Per-node geometry from compute_tree_geometry() for tree.
sources (Array) – Far-pair source node indices, shape (num_pairs,).
targets (Array) – Far-pair target node indices, shape (num_pairs,).
level_offsets (Array | None) – Optional precomputed level offsets passed through for diagnostics; derived from tree when omitted.

Returns:

Class keys, per-class displacements, and CSR-style class offsets/ids.

Return type:

GroupedInteractionBuffers

Moments¶

Mass moments for radix tree nodes.

This module accumulates per-node total masses and centres of mass from particles that were reordered into Morton order by yggdrax.tree. These summaries are useful when building multipole expansions during the upward FMM sweep.

class yggdrax.tree_moments.TreeMassMoments(mass, center_of_mass)[source]¶

Bases: NamedTuple

Aggregate mass information for every node in a radix tree.

Parameters:

mass (Array)
center_of_mass (Array)

mass: Array¶: Alias for field number 0

center_of_mass: Array¶: Alias for field number 1

class yggdrax.tree_moments.TreeMultipoleMoments(max_order, mass, center, dipole, second_moment, quadrupole, third_moment, octupole, fourth_moment, hexadecapole, raw_packed)[source]¶

Bases: NamedTuple

Multipole summaries (through hexadecapole) for tree nodes.

Parameters:

max_order (int)
mass (Array)
center (Array)
dipole (Array)
second_moment (Array)
quadrupole (Array)
third_moment (Array)
octupole (Array)
fourth_moment (Array)
hexadecapole (Array)
raw_packed (Array)

max_order: int¶: Alias for field number 0

mass: Array¶: Alias for field number 1

center: Array¶: Alias for field number 2

dipole: Array¶: Alias for field number 3

second_moment: Array¶: Alias for field number 4

quadrupole: Array¶: Alias for field number 5

third_moment: Array¶: Alias for field number 6

octupole: Array¶: Alias for field number 7

fourth_moment: Array¶: Alias for field number 8

hexadecapole: Array¶: Alias for field number 9

raw_packed: Array¶: Alias for field number 10

yggdrax.tree_moments.compute_tree_mass_moments(tree, positions_sorted, masses_sorted)[source]¶

Compute total mass and center of mass for every tree node.

Uses prefix sums over the Morton-sorted masses and mass-weighted positions for O(1)-per-node range queries.

Parameters:

tree (object) – Tree container or topology exposing the FMM-core contract.
positions_sorted (Array) – Particle positions in Morton order, shape (n_particles, 3).
masses_sorted (Array) – Particle masses reordered identically to positions_sorted.

Returns:

Per-node total mass and center of mass.

Return type:

TreeMassMoments

yggdrax.tree_moments.compute_tree_multipole_moments(tree, positions_sorted, masses_sorted, expansion_centers=None, max_order=2)[source]¶

Compute multipole information up to max_order for each node.

Parameters:

tree (object) – Tree/topology exposing the FMM-core topology contract.
positions_sorted (Array) – Particle positions in Morton order.
masses_sorted (Array) – Particle masses reordered identically to positions_sorted.
expansion_centers (Array | None) – Optional array (num_nodes, 3) specifying the expansion center for each node. When omitted, the node center-of-mass is used, which yields zero dipole moments.
max_order (int) – Highest multipole order to accumulate (0 ≤ max_order ≤ 4).

Returns:

Packed multipole moments for each node through max_order.

Return type:

TreeMultipoleMoments

Raises:

ValueError – If expansion_centers has an incompatible shape or max_order is outside the supported range.

yggdrax.tree_moments.multipole_from_packed(packed, centers, max_order)[source]¶

Reconstruct multipole tensors from packed triangular coefficients.

Inverse of the packing done by pack_multipole_expansions(): expands the packed per-node coefficient rows back into named Cartesian moment tensors (dipole, quadrupole, octupole, hexadecapole) up to max_order.

Parameters:

packed (Array) – Packed coefficients of shape (num_nodes, >= total_coefficients(max_order)).
centers (Array) – Per-node expansion centers of shape (num_nodes, 3).
max_order (int) – Highest multipole order to reconstruct (0 ≤ max_order ≤ 4).

Returns:

Per-node moment tensors through max_order.

Return type:

TreeMultipoleMoments

Raises:

ValueError – If packed is not rank-2, lacks enough coefficients, or its row count does not match centers.

yggdrax.tree_moments.tree_moments_from_raw(packed, centers, max_order)[source]¶

Build TreeMultipoleMoments from raw central packed coefficients.

Like multipole_from_packed(), but treats packed as coefficients already expressed about centers (the “central”/raw convention) rather than converting from mass moments.

Parameters:

packed (Array) – Raw packed coefficients of shape (num_nodes, >= total_coefficients(max_order)).
centers (Array) – Per-node expansion centers of shape (num_nodes, 3).
max_order (int) – Highest multipole order to reconstruct (0 ≤ max_order ≤ 4).

Returns:

Per-node moment tensors through max_order.

Return type:

TreeMultipoleMoments

Raises:

ValueError – If packed is not rank-2, lacks enough coefficients, or its row count does not match centers.

yggdrax.tree_moments.pack_multipole_expansions(moments, max_order)[source]¶

Pack multipole coefficients into the triangular layout up to max_order.

Flattens the named moment tensors on moments into the contiguous packed representation consumed by translate_packed_moments() and multipole_from_packed().

Parameters:

moments (TreeMultipoleMoments) – Per-node multipole moments to pack.
max_order (int) – Highest order to include (must not exceed moments.max_order).

Returns:

Packed coefficients of shape (num_nodes, total_coefficients(max_order)).

Return type:

Array

Raises:

ValueError – If max_order exceeds the order stored on moments.

Utilities for handling symmetric multipole tensors.

This module provides helpers for working with the packed triangular representation discussed in recent FMM optimisation work. The packed layout stores all Cartesian symmetric tensor components of order l contiguously using a triangular indexing scheme, which avoids redundant entries while remaining friendly to vectorised JAX code.

yggdrax.multipole_utils.multi_index_tuples(level)[source]¶

Return tuples (i, j, k) satisfying i + j + k = level.

Parameters:: level (int)
Return type:: tuple[tuple[int, int, int], …]

yggdrax.multipole_utils.multi_index_factorial(combo)[source]¶

Return i! * j! * k! for a multi-index tuple.

Parameters:: combo (tuple[int, int, int])
Return type:: int

yggdrax.multipole_utils.multi_power(vec, combo)[source]¶

Return vec[0]^i * vec[1]^j * vec[2]^k for combo = (i, j, k).

Parameters:

vec (Array)
combo (tuple[int, int, int])

Return type:

Array

yggdrax.multipole_utils.level_size(level)[source]¶

Return coefficient count for a symmetric tensor of order level.

Parameters:: level (int)
Return type:: int

yggdrax.multipole_utils.level_offset(level)[source]¶

Return the packed offset for order level.

Offsets accumulate contributions from lower orders using the closed form level(level+1)(level+2)/6.

Parameters:: level (int)
Return type:: int

yggdrax.multipole_utils.total_coefficients(max_order)[source]¶

Return total packed length for orders 0..max_order inclusive.

Parameters:: max_order (int)
Return type:: int

yggdrax.multipole_utils.triangular_index(level, i, j)[source]¶

Map Cartesian indices to the packed triangular index.

The mapping assumes i >= 0, j >= 0 and i + j <= level. The remaining index is k = level - i - j.

Parameters:

level (int)
i (int)
j (int)

Return type:

int

yggdrax.multipole_utils.triangular_indices(level)[source]¶

Enumerate all (i, j, k) tuples for a given tensor order.

Parameters:: level (int)
Return type:: Array

yggdrax.multipole_utils.pack_tensor(level, tensor)[source]¶

Pack a symmetric Cartesian tensor of order level.

Parameters:

level (int) – Tensor order l.
tensor (Array) – Array of shape (l + 1, l + 1, l + 1) containing Cartesian components. Only entries with i + j + k = l are read.

Returns:

1-D flattened packed representation with length level_size(level).

Return type:

Array

Raises:

ValueError – If tensor does not match shape (level + 1, level + 1, level + 1).

yggdrax.multipole_utils.unpack_tensor(level, data)[source]¶

Unpack a packed triangular buffer back into Cartesian components.

Inverse of pack_tensor().

Parameters:

level (int) – Tensor order l.
data (Array) – Packed representation whose last axis has length level_size(level).

Returns:

Dense tensor of shape (l + 1, l + 1, l + 1) with the symmetric components filled in (entries with i + j + k != l are zero).

Return type:

Array

Raises:

ValueError – If data’s last axis does not have length level_size(level).

Morton ordering and primitives¶

Morton code utilities for 3D spatial ordering.

yggdrax.morton.get_common_prefix_length(code1, code2)[source]¶

Return the number of common leading bits of two Morton codes.

This is the length of the shared binary prefix, which corresponds to the depth of the deepest tree node containing both codes.

Parameters:

code1 (Array) – First scalar uint64 Morton code.
code2 (Array) – Second scalar uint64 Morton code.

Returns:

Number of shared leading bits (64 when the codes are equal).

Return type:

int

yggdrax.morton.morton_decode(codes, bounds)[source]¶

Decode Morton codes back to approximate 3D positions.

Inverse of morton_encode() up to the 21-bit-per-axis quantization (decoded positions land at their cell centers within bounds).

Parameters:

codes (Array) – uint64 Morton codes of shape (n,).
bounds (tuple[Float[Array, '3'], Float[Array, '3']]) – (min_corner, max_corner) box used at encode time, each shape (3,).

Returns:

Decoded positions of shape (n, 3).

Return type:

Array

yggdrax.morton.morton_encode(positions, bounds)¶

Encode 3D positions as 21-bit-per-axis Morton (Z-order) codes.

Un-jitted implementation. Prefer morton_encode() for standalone use; this raw version is for calling inside another transform (e.g. shard_map), where nesting a separately-jit-compiled function would clash with the enclosing mesh’s sharding-in-types.

Parameters:

positions (Array) – Points of shape (n, 3).
bounds (tuple[Float[Array, '3'], Float[Array, '3']]) – (min_corner, max_corner) axis-aligned box, each of shape (3,). Positions are normalized into this box and clipped to stay in range.

Returns:

uint64 Morton codes of shape (n,).

Return type:

Array

yggdrax.morton.morton_encode_impl(positions, bounds)[source]¶

Encode 3D positions as 21-bit-per-axis Morton (Z-order) codes.

Parameters:

positions (Array) – Points of shape (n, 3).
bounds (tuple[Float[Array, '3'], Float[Array, '3']]) – (min_corner, max_corner) axis-aligned box, each of shape (3,). Positions are normalized into this box and clipped to stay in range.

Returns:

uint64 Morton codes of shape (n,).

Return type:

Array

yggdrax.morton.sort_by_morton(codes)[source]¶

Return indices that stably sort points by Morton code.

Parameters:: codes (Array) – Morton codes of shape (n,).
Returns:: Permutation indices of shape (n,) such that codes[result] is ascending.
Return type:: Array

Bounding-box helpers for tree construction.

yggdrax.bounds.infer_bounds(positions)[source]¶

Infer a padded axis-aligned bounding box from particle positions.

The box spans the particle extent plus 5% padding per axis (at least 1e-6), so points on the boundary do not clip to the edge during Morton encoding.

Parameters:: positions (Array) – Particle positions of shape (n, 3).
Returns:: (min_corner, max_corner), each of shape (3,).
Return type:: tuple of Array

Local dtype policy for Yggdrax contracts.

yggdrax.dtypes.INDEX_DTYPE¶: alias of int64

yggdrax.dtypes.as_index(x)[source]¶

Convert a scalar/array to yggdrax index dtype.

Parameters:: x (Array | ndarray | bool | number | bool | int | float | complex | TypedNdArray)
Return type:: Array

yggdrax.dtypes.complex_dtype_for_real(real_dtype)[source]¶

Return complex dtype paired with a real floating dtype.

Parameters:: real_dtype (str | type[Any] | dtype | SupportsDType)
Return type:: str | type[Any] | dtype | SupportsDType

Contracts and types¶

Structural protocols for tree/topology capabilities.

class yggdrax.protocols.MortonLeafBoundsProtocol(*args, **kwargs)[source]¶

Bases: TreeRangesProtocol, Protocol

Adds Morton-derived leaf box information for geometry construction.

use_morton_geometry: Array | bool¶

bounds_min: Array¶

bounds_max: Array¶

leaf_codes: Array¶

leaf_depths: Array¶

class yggdrax.protocols.TopologyContainerProtocol(*args, **kwargs)[source]¶

Bases: Protocol

Container exposing a concrete topology payload.

topology: object¶

class yggdrax.protocols.TreeLevelIndexProtocol(*args, **kwargs)[source]¶

Bases: TreeStructureProtocol, Protocol

Adds explicit level-order indexing metadata.

node_level: Array¶

nodes_by_level: Array¶

level_offsets: Array¶

num_levels: Array | int¶

class yggdrax.protocols.TreeRangesProtocol(*args, **kwargs)[source]¶

Bases: TreeStructureProtocol, Protocol

Adds node-to-particle range metadata.

node_ranges: Array¶

num_particles: Array | int¶

class yggdrax.protocols.TreeStructureProtocol(*args, **kwargs)[source]¶

Bases: Protocol

Minimal structure needed for parent/child traversal.

parent: Array¶

left_child: Array¶

right_child: Array¶

Shared prepared artifact contracts for Yggdrax consumers.

class yggdrax.types.PreparedTreeArtifacts(tree, positions_sorted, masses_sorted, inverse_permutation, geometry, interactions, neighbors, traversal_result=None)[source]¶

Bases: object

Canonical prepared tree/traversal bundle.

Parameters:

tree (Tree)
positions_sorted (Array)
masses_sorted (Array)
inverse_permutation (Array)
geometry (TreeGeometry)
interactions (NodeInteractionList)
neighbors (NodeNeighborList)
traversal_result (TraversalResult | None)

tree: Tree¶

positions_sorted: Array¶

masses_sorted: Array¶

inverse_permutation: Array¶

geometry: TreeGeometry¶

interactions: NodeInteractionList¶

neighbors: NodeNeighborList¶

traversal_result: TraversalResult | None = None¶

class yggdrax.types.TraversalResult(interaction_offsets, interaction_sources, interaction_targets, interaction_tags, interaction_counts, neighbor_offsets, neighbor_indices, neighbor_counts, leaf_indices, far_pair_count, near_pair_count, queue_overflow, far_overflow, near_overflow, accept_decisions, near_decisions, refine_decisions)[source]¶

Bases: object

Local yggdrax view of dual-tree traversal outputs.

Parameters:

interaction_offsets (Array)
interaction_sources (Array)
interaction_targets (Array)
interaction_tags (Array)
interaction_counts (Array)
neighbor_offsets (Array)
neighbor_indices (Array)
neighbor_counts (Array)
leaf_indices (Array)
far_pair_count (Array)
near_pair_count (Array)
queue_overflow (Array)
far_overflow (Array)
near_overflow (Array)
accept_decisions (Array)
near_decisions (Array)
refine_decisions (Array)

interaction_offsets: Array¶

interaction_sources: Array¶

interaction_targets: Array¶

interaction_tags: Array¶

interaction_counts: Array¶

neighbor_offsets: Array¶

neighbor_indices: Array¶

neighbor_counts: Array¶

leaf_indices: Array¶

far_pair_count: Array¶

near_pair_count: Array¶

queue_overflow: Array¶

far_overflow: Array¶

near_overflow: Array¶

accept_decisions: Array¶

near_decisions: Array¶

refine_decisions: Array¶

yggdrax.types.TraversalArtifacts¶: alias of PreparedTreeArtifacts

yggdrax.types.traversal_result_from_expanse(result)[source]¶

Convert expanse traversal result into yggdrax local contract.

Parameters:: result (DualTreeWalkResult)
Return type:: TraversalResult

Traversal policy helpers for Yggdrax.

class yggdrax.policies.TraversalPolicy(config, description='Explicit dual-tree traversal capacities.')[source]¶

Bases: object

Top-level traversal policy contract.

Parameters:

config (DualTreeTraversalConfig)
description (str)

config: DualTreeTraversalConfig¶

description: str = 'Explicit dual-tree traversal capacities.'¶