balltree.BallTree

class balltree.BallTree(xyz: ArrayLike, weight: ArrayLike | None = None, leafsize: int = default_leafsize)

A wrapper for the C implementation of BallTree.

The tree is implemented for 3-dim coordinates and an Euclidean distance metric. It provides different alorigthms to count neighbours within a single or multiple query points.

The data point(s) xyz can be a numpy array of shape (3,) or (N, 3), or an equivalent python object. The optional weights can be a float or a 1-dim sequence of matching length, the optional leafsize determines when the tree query algorithms switch from traversal to brute force.

classmethod from_random(cls, low: float, high: float, size: int, leafsize: int = default_leafsize) → BallTree

Build a new BallTree instance from randomly generated points.

The (x, y, z) coordinates are generated uniformly in the interval [low, high), size controlls the number of points generated. The optional leafsize determines when the tree query algorithms switch from traversal to brute force.

classmethod from_file(cls, path: str) → BallTree

Load back a class instance from a file created by the to_file() method.

property data: NDArray

Get a view of the data points stored in the tree.

Note that the points are reordered when the tree is constructed.

num_points: int

Get the number of data points stored in the tree.

leafsize: int

Get the leaf size of the tree.

sum_weight: float

Get the sum over the data point weights stored in the tree.

center: tuple(float, float, float)

Get the coordinates of the root node of the tree.

radius: float

Get the radius of the root node of the tree.

to_file(self, path: str) → None

Store a representation of the tree instance in a binary file.

count_nodes(self) → int

Get a count of all nodes of the tree, including the root node.

get_node_data(self) → NDArray

Collect the meta data of all tree nodes in a numpy array.

The array fields record depth (starting from the root node), num_points, sum_weight, x, y, z (node center) and node radius.

nearest_neighbours(self, xyz: ArrayLike, k: int, max_dist: float = -1.0) → NDArray

Query a fixed number of nearest neighbours.

The query point(s) xyz can be a numpy array of shape (3,) or (N, 3), or an equivalent python object. The number of neighbours k must be a positive integer and the optional max_dist parameter puts an upper bound on the separation to the neighbours.

Returns an array with fields index, holding the index to the neighbour in the array from which the tree was constructed, and distance, the separation. The result is sorted by separation, missing neighbours (e.g. if distance > max_dist) are indicated by an index of -1 and infinite separation.

brute_radius(self, xyz: ArrayLike, radius: float, weight: ArrayLike | None = None) → float

Count neighbours within a given radius using brute force.

The query point(s) xyz can be a numpy array of shape (3,) or (N, 3), or an equivalent python object. The optional weights can be a float or a 1-dim sequence of matching length.

count_radius(self, xyz: ArrayLike, radius: float, weight: ArrayLike | None = None) → float

Count neighbours within a given radius using tree traversal.

The query point(s) xyz can be a numpy array of shape (3,) or (N, 3), or an equivalent python object. The optional weights can be a float or a 1-dim sequence of matching length.

dualcount_radius(self, tree: BallTree, radius: float) → NDArray

Count all pairs within a given radius from points in another tree.

The pairs between the two trees are computed with the efficient dualtree algorithm.

brute_range(self, xyz: ArrayLike, radii: ArrayLike, weight: ArrayLike | None = None) → NDArray

Count neighbours within a sequence of radii using brute force.

The query point(s) xyz can be a numpy array of shape (3,) or (N, 3), or an equivalent python object. The radii must either be a float or monotonic sequence. The optional weights can be a float or a 1-dim sequence of matching length.

Returns an array of counts. The first element contains the count of all neighbours 0 <= r <= r_1, the following values contain the incremental counts r_i-1 <= r <= r_i.

count_range(self, xyz: ArrayLike, radii: ArrayLike, weight: ArrayLike | None = None) → NDArray

Count neighbours within a sequence of radii using tree traversal.

The query point(s) xyz can be a numpy array of shape (3,) or (N, 3), or an equivalent python object. The radii must either be a float or monotonic sequence. The optional weights can be a float or a 1-dim sequence of matching length.

Returns an array of counts. The first element contains the count of all neighbours 0 <= r <= r_1, the following values contain the incremental counts r_i-1 <= r <= r_i.

dualcount_range(self, other: BallTree, radii: ArrayLike) → NDArray

Count all pairs within a sequence of radii from points in another tree.

The pairs between the two trees are computed with the efficient dualtree algorithm. The radii must either be a float or monotonic sequence. The optional weights can be a float or a 1-dim.

Returns an array of counts. The first element contains the count of all neighbours 0 <= r <= r_1, the following values contain the incremental counts r_i-1 <= r <= r_i.