Numba Cubes¶

packingcubes.cubes.cubes_numba ¶

Jitted functions for handling cubes structures

All cubes functionality that can be numba jitted should be included in this module.

In general these functions assume that cubes structure can be decomposed into three parts similar to a structured array. The parts are

An array of offset indices, cube_indices, such that the particles in cube i correspond to the data indices cube_indices[i]:cube_indices[i+1]. For the case i==len(cubes)-1 (the last cube), the end index is len(data).
A list of cube bounding boxes, cube_boxes.
A list of PackedTrees (actually PackedTreeNumbas`)

cube ¶

cube(data, cubes_per_side, box)

Bin the loaded particles into the different cubes

Effectivaly perform a parallel radix sort of the particles into the CxCxC cube bins, where C=cubes_per_side.

Parameters:

data (DataContainer) –

The 3D particle data to sort
cubes_per_side (int) –

The number of cubes on a side
box (BoundingBox) –

The bounding box to use, sets the bounding boxes of the individual cubes. Does not need to be the same as data.bounding_box

Returns:

chopped ( NDArray of ints ) –

An array of data offsets such that cube 0 is the data from chopped[0]:chopped[1], cube 1 is the data from chopped[1]:chopped[2], etc. The data in the final cube is from chopped[number_cubes-1]:len(data)

get_array_in_shape ¶

get_array_in_shape(cubes, trees, cube_offsets, shape, data_positions, array, strict)

Get the array of particle indices in the specified shape

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

Parameters:

cubes (List[BoundingBox]) –

List of cube bounding boxes
trees (List[PackedTreeNumba]) –

List of cube PackedTreeNumbas
cube_offsets (NDArray) –

Array of cube offset indices into the data
shape (BoundingVolume) –

BoundingVolume to check
data_positions (NDArray) –

Particle positions information. If data is not available and you want non-strict containment, consider calling [_parallel_expand_all_array][packingcubes.cubes.cubes_numba._parallel_expand_all_array] or [_parallel_expand_all_matrix][packingcubes.cubes.cubes_numba._parallel_expand_all_matrix] directly.
array (NDArray) –

The actual array to expand.
strict (bool) –

If False, only check whether node is within shape, which is equivalent to expanding the node start/stops from _get_particle_indices_in_shape. Default True.

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

get_closest_particles ¶

get_closest_particles(cubes, trees, cube_indices, data, xyz, k, distance_function, distance_upper_bound, use_shuffle, return_sorted)

Return the k-closest particle distances and their indices

Parameters:

cubes (List[BoundingBox]) –

The cube boxes
trees (List[PackedTreeNumba]) –

The cube trees
cube_indices (NDArray) –

The cube index offsets
data (DataContainer) –

The container of the position data
xyz (NDArray) –

The 3 Cartesian coordinates
k (int) –

The number of particles to return. No verification of sign is performed
distance_function (Callable[[float, float, float, float, float, float], float]) –

The distance function between two Cartesian points, e.g. d = distance_function(x1, y1, z1, x2, y2, z2)
distance_upper_bound (float) –

The maximum distance to consider particles within. May result in fewer than k particles being returned if too stringent
use_shuffle (bool) –

Flag to return shuffle indices instead of sorted data indices

Returns:

dists ( NDArray[float] ) –

K-length vector of distances
inds ( NDArray[int] ) –

K-length vector of particle indices

get_packednodes_in_shape ¶

get_packednodes_in_shape(cubes, trees, cube_offsets, shape)

Get the PackedNodes contained within the shape

This is not done in parallel, unlike get_particle_indices_in_shape, and is not expected to be especially performant, though is otherwise implemented similarly.

Parameters:

cubes (List[BoundingBox]) –

List of cube bounding boxes
trees (List[PackedTreeNumba]) –

List of cube PackedTreeNumbas
cube_offsets (NDArray) –

Array of cube offset indices into the data
shape (BoundingVolume) –

BoundingVolume to check

Returns:

entire ( List[PackedNode] ) –

List of PackedNodes entirely contained within shape
partial ( List[PackedNode] ) –

List of PackedNodes partially contained within shape

get_particle_index_list_in_shape ¶

get_particle_index_list_in_shape(cubes, trees, cube_offsets, shape, data, strict, use_data_indices)

Get the array of particle indices in the specified shape

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

Parameters:

cubes (List[BoundingBox]) –

List of cube bounding boxes
trees (List[PackedTreeNumba]) –

List of cube PackedTreeNumbas
cube_offsets (NDArray) –

Array of cube offset indices into the data
shape (BoundingVolume) –

BoundingVolume to check
data (DataContainer) –

Particle positions information. If data is not available and you want non-strict data indices, consider calling [_parallel_expand_all_data_indices][packingcubes.cubes.cubes_numba._parallel_expand_all_data_indices] directly.
strict (bool) –

If False, only check whether node is within shape, which is equivalent to expanding the node start/stops from _get_particle_indices_in_shape. Default True.
use_data_indices (bool) –

Return shuffle indices if False. Default True.

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

get_particle_indices_in_shape ¶

get_particle_indices_in_shape(cubes, trees, cube_offsets, shape)

Get the particle start-stop tuples in the specified shape

Parameters:

cubes (List[BoundingBox]) –

List of cube bounding boxes
trees (List[PackedTreeNumba]) –

List of cube PackedTreeNumbas
cube_offsets (NDArray) –

Array of cube offset indices into the data
shape (BoundingVolume) –

BoundingVolume to check

Returns:

indices ( Xx3 NDArray[np.int_] ) –

Array of index information. Each row 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.

make_trees ¶

make_trees(data, cube_indices, cube_boxes, particle_threshold)

Create a packed tree for the data in each cube in parallel

Parameters:

data (DataContainer) –

The particle data to use
cube_indices (NDArray) –

The particle index offsets for each cube
cube_boxes (List[BoundingBox]) –

The list of bounding boxes of each cube
particle_threshold (int) –

The number of particles needed to split a node when creating the PackedTreeNumba

Returns:

list[NDArray] –

The packed tree for each cube in packed array format