Skip to content

Numba Packed Trees

packingcubes.packed_tree.packed_tree_numba

Jitclass version of PackedTree and supporting functions

Underlying jitclass of a PackedTree that contains the actual implementation of most of the functions along with several supporting functions. While not strictly a private module, the functions contained here are not intended for use generally, but for extensions.

PackedTreeNumba 

PackedTreeNumba(box, tree, particle_threshold=_DEFAULT_PARTICLE_THRESHOLD)

Private jitted octree interface

This interface defines the methods for manipulating and traversing a packingcubes octree.

box instance-attribute 

box = box

Bounding box for the tree

Effectively the root level bounding box

particle_threshold instance-attribute 

particle_threshold = particle_threshold

Maximum number of particles in a leaf

tree instance-attribute 

tree = tree

Tree representation in memory

Note that this does not include any of the metadata. tree[0] should the start of the root node data in packed tree format.

__len__ 

__len__()

Return the number of particles in the tree

count_neighbors 

count_neighbors(other, r)

Compute the number of particles within r of particles in other

get_closest_particles 

get_closest_particles(data, xyz, distance_function, distance_upper_bound, k=1, use_shuffle=False, return_sorted=False)

Get kth nearest particle distances and indices to point

Parameters:

  • data (DataContainer) –

    Source of particle position data

  • xyz (NDArray) –

    Coordinates of point to check

  • distance_function (Callable[[float, float, float, float, float, float], float]) –

    Function to compute distance. Must be njit compatible! Needs to accept 6 coordinates: x, y, z and px, py, pz and return the distance as float64

  • distance_upper_bound (float) –

    Return only neighbors from other nodes within this distance. This is used for tree pruning, so if you are doing a series of nearest-neighbor queries, it may help to supply the distance to the nearest neighbor of the most recent point.

  • k (int, default: 1 ) –

    Number of closest particles to return. Default 1

  • use_shuffle (bool, default: False ) –

    Flag to return shuffle indices instead of sorted data indices. Default False.

  • return_sorted (bool, default: False ) –

    Return output in order by distance. Default False

Returns:

  • distances ( NDArray[float] ) –

    Distances to the kth nearest neighbors. Has shape (min(N,k),), where N is the number of particles in the sphere bounded by distance_upper_bound

  • indices ( NDArray[int] ) –

    Indices in data of the kth nearest neighbors. Has same shape as distances

Raises:

  • OctreeError –

    if something goes wrong with the search process

get_leaves 

get_leaves()

Return the list of PackedNodeNumba leaves

Note that the node objects are generated on the fly and are not backed by the tree; any modifications will not propagate.

get_node 

get_node(tag)

Get a PackedNodeNumba object represented by tag or None if non-existent

Note that this object is effectively a copy; any modifications will not propagate.

Parameters:

  • tag (str) –

    The name of the node to search for. Nodes are named by child order, in z-order so e.g. "012" is the front-mid-left-bottom grandchild.

Returns:

  • PackedNodeNumba | None –

    A copy of the node in the tree or None if it DNE

get_particle_indices_in_box 

get_particle_indices_in_box(box)

Return all particles contained within the box

Parameters:

  • box (BoxLike) –

    Box to check

Returns:

  • indices ( list[tuple[int, int, int]] ) –

    List of particle start-stop indices contained within box. Each tuple describes a chunk/slice of data in the form (start, stop, partial), where partial is a flag:

    • 1 if the data chunk is entirely contained within box,
    • 0 otherwise.

get_particle_indices_in_sphere 

get_particle_indices_in_sphere(center, radius)

Return all particles contained within the sphere defined by center and radius

Parameters:

  • center (NDArray) –

    Center point of the sphere

  • radius (float) –

    Radius of the sphere

Returns:

  • indices ( list[tuple[int, int, int]] ) –

    List of particle start-stop indices contained within sphere. Each tuple describes a chunk/slice of data in the form (start, stop, partial), where partial is a flag:

    • 1 if the data chunk is entirely contained within the sphere,
    • 0 otherwise.

euclidean_distance_particle 

euclidean_distance_particle(x, y, z, px, py, pz)

Compute the Euclidean distance between two points

euclidean_d2 

euclidean_d2(x, y, z, px, py, pz)

Compute the squared Euclidean distance between two points

_construct_tree 

_construct_tree(data, box=None, particle_threshold=_DEFAULT_PARTICLE_THRESHOLD)

Construct a packed tree, generating a node at a time recursively

Functions as the initial entry point into _construct_node_recursive.

Tip

If you know that all the particles are in a subregion of the data's bounding box, specifying a smaller box can significantly speed up tree construction and especially search.

Parameters:

  • data (DataContainer) –

    The backing data of this tree

  • box (BoundingBox | None, default: None ) –

    Bounding box of the tree. Defaults to the data's bounding box

  • particle_threshold (int, default: _DEFAULT_PARTICLE_THRESHOLD ) –

    The maximum number of particles a leaf node can hold before splitting Defaults to octree._DEFAULT_PARTICLE_THRESHOLD

Returns:

  • tree ( NDArray[uint32] ) –

    Array of uint32s representing the packed tree.

_construct_node_recursive 

_construct_node_recursive(data, tree, node, parent_index, max_depth, particle_threshold=_DEFAULT_PARTICLE_THRESHOLD)

Construct a packed tree, generating a node at a time recursively

Parameters:

  • data (DataContainer) –

    The backing data of this tree

  • tree (list[uint32]) –

    List of uint32s representing the in-progress packed tree. Initial call should pass an empty list.

  • node (CurrentNode) –

    Pointer to the node this invocation is constructing. Initial call should pass a node corresponding to the root

  • parent_index (int) –

    Index of the current node's parent in the in-progress tree. Initial call should pass 0

  • max_depth (int) –

    The maximum depth supported by the bounding box of the data before floating point errors will occur. A node with level >=max_depth is required/guaranteed to be a leaf node

  • particle_threshold (int, default: _DEFAULT_PARTICLE_THRESHOLD ) –

    The maximum number of particles a leaf node can hold before splitting. Defaults to octree._DEFAULT_PARTICLE_THRESHOLD

Returns:

  • int –

    The index of the next node in the tree at the same level or higher

Additional semi-private methods for PackedTreeNumba:

packingcubes.packed_tree.packed_tree_numba.PackedTreeNumba._get_particle_indices_in_shape 

_get_particle_indices_in_shape(containment_obj)

Return all particles contained within a shape

Parameters:

Returns:

  • indices ( list[tuple[int, int, int]] ) –

    List of particle start-stop indices contained within shape. Each tuple describes a chunk/slice of data in the form (start, stop, partial), where partial is a flag:

    • 1 if the data chunk is entirely contained within shape,
    • 0 otherwise.

packingcubes.packed_tree.packed_tree_numba.PackedTreeNumba._get_particle_index_list_in_shape 

_get_particle_index_list_in_shape(data, containment_obj)

Return all particles contained within a shape

If the data argument is specified, will do additional containment-checks at the particle level

Parameters:

  • data (DataContainer | None) –

    Particle positions information. If None, only check whether node is within shape, which is equivalent to expanding the node start/stops from _get_particle_indices_in_shape

  • containment_obj (BoundingVolume) –

    Provides an exact containment test (e.g. contained within a sphere).

Returns:

  • indices ( NDArray[int]] ) –

    List of particle indices contained within shape. Will contain any additional particles that can be found in the same nodes if data is not provided

packingcubes.packed_tree.packed_node

PackedTree classes representing tree nodes and related functions

CurrentNode 

CurrentNode(box, index=0, node_start=0, node_end=0, tag=None, child_flag=0, my_index=0, level=0, empty=0)

Jitclass implementation of a node in a PackedTree

Essentially, this class represents a pointer to a PackedTree node, effectively acting as the current struct in the Structured Tree.

box instance-attribute 

box = box

The bounding box of this node. This is computed from the root node

child_flag instance-attribute 

child_flag = child_flag

Bitflag where values signify which of the possible 8 children are present e.g. 00010100 corresponds to the 3rd and 5th children. A leaf node has child_flag==0

empty instance-attribute 

empty = empty

Currently unused metadata. Only values within [0,255] would be meaningful

index instance-attribute 

index = index

Index in the tree data array of this node

level instance-attribute 

level = level

This node's level in the tree. Root is 0

my_index instance-attribute 

my_index = my_index

The index order of this node among it's 8 siblings. Root is 0

node_end instance-attribute 

node_end = node_end

(Inclusive) Ending index of the particles in this node

node_start instance-attribute 

node_start = node_start

Starting index of the particles in this node

tag instance-attribute 

tag = zeros((64,), dtype=uint8) if tag is None else tag

tag is a (64,) uint8 array where tag[63] is a pointer to the current level. So node 0264 would look like [0, 2, 6, 4, 0, 0, ..., 0, 3]

__len__ 

__len__()

Return the number of particles in this node

PackedNodeNumba 

PackedNodeNumba(node_start, node_end, box, tag=None)

Private representation of a node in a packed tree

This class is what should be emitted by PackedTreeNumba methods that return nodes, rather than CurrentNode instances.

Note that this only represents the node at time of creation. Changes do not propagate in either direction.

Parameters:

  • node_start (int) –

    Starting index of node in dataset

  • node_end (int) –

    End index of node in dataset

  • box (BoundingBox) –

    Bounding box of node

  • tag (str | None, default: None ) –

    Tag describing node. Default is "0" (root)

copy 

copy()

Return a (deep) copy

get_name 

get_name(node)

Get the name (tag) of this CurrentNode

get_children 

get_children(node)

Return a list of 0-based children indices for this CurrentNode

is_leaf 

is_leaf(node)

Return True if node is a leaf node

is_root 

is_root(node)

Return True if node is the root node

expand_range 

expand_range(node)

Return the array of indices contained within this node