Tree Structures

Tree Structures#

UXarry supports two different tree structures, BallTree and KDTree. These trees are tailored for neighbor searches, making them useful for spatial data queries.

import uxarray as ux

For this example we will be using a UGRID meshfile

grid_path = "../../test/meshfiles/ugrid/quad-hexagon/grid.nc"
uxgrid = ux.open_grid(grid_path)

BallTree#

UXarray BallTree is built off of sklearn.neighbors.BallTree. A BallTree data structure organizes points in a multi-dimensional space into a tree of spheres. It is highly efficient for higher-dimensional data and for fast queries.

Parameters#

The BallTree class can be accessed through the Grid.get_ball_tree(coordinates, coordinate_system, distance_metric, reconstruct) method, which takes in the following parameters:

  • coordinates allows us to specify what nodes to build the tree on. We can choose from nodes, face_centers, or edge_centers. Each one will change how the tree is built. nodes builds the tree from the corner nodes of the grid, face_centers builds the tree from the face centers of each face, and edge_centers builds from the center of each edge in the grid. The default is to build from nodes.

  • coordinate_system specifies the tree’s coordinate type when constructing. We can use either cartesian, which uses the (x, y, z) coordinate system, or spherical, which uses the (lat, lon) coordinate system. The default parameter is spherical.

  • distance_metric relates to the distance computation, typically returned as a distance when querying for neighbors. There are a large number of options for us to use. A list of these options can be found here. An important note here is that some distance metrics aren’t compatible with some coordinate systems. BallTree uses the haversine distance as default, which will only work with spherical coordinates and not with cartesian. The default parameter is haversine.

  • reconstruct is a bool variable that allows the user to reconstruct the tree. As default for performance, if a user calls get_ball_tree and a tree has already been created, it will simply use that one. If reconstruct is set to True, it will override this and reconstruct the tree. The default parameter is False.

Constructing a BallTree#

We can store the BallTree data structure in a variable, which allows us to access the tree in a simple way for queries.

ball_tree = uxgrid.get_ball_tree(
    coordinates="nodes",
    coordinate_system="spherical",
    distance_metric="haversine",
    reconstruct="False",
)

Query#

Now we can use that variable to query for the distance and indexes of the nearest neigbhors. The first parameter is the point from which to do the search. return_distance allows us to choose to return the distance of the neighbors, and k controls how many neighbors to find.

d, ind = ball_tree.query([0.0, 0.0], return_distance=True, k=1)

If we don’t plan on using the tree for other things in the future we can skip the extra step and query right away

d, ind = uxgrid.get_ball_tree(
    coordinates="nodes",
    coordinate_system="spherical",
    distance_metric="haversine",
    reconstruct="True",
).query([0.0, 0.0], return_distance=True, k=1)

Query Radius#

We can also query the tree using a radius search, instead of a nearest neighbor search. This allows us to get all points within a certain radius of a specific point. For spherical coordinates, the radius is in units of degrees, and for cartesian coordinates, the radius is in meters.

d, ind = ball_tree.query_radius([0.0, 0.0], r=5, return_distance=True)

KDTree#

The KDTree structure is a binary search tree useful for low-dimensional data. Its implementation is almost identical to BallTree, and the parameters are identical. An important note is that the different allowed inputs for distance_metric change between the trees. For KDTree allowed distance_metrics can be found here. We can call it using get_kd_tree(). Generally, KDTree is going to be slower than BallTree, and it is recommended to use BallTree for most im

kd_tree = uxgrid.get_kd_tree()

query() and query_radius() work identically to the BallTree.

kd_tree.query([1.0, 0.0, 0.0], k=1)

(array(0.00104681), array(0))

kd_tree.query_radius([1.0, 0.0, 0.0], r=5, return_distance=True)

(array([0.00104681, 0.00335129, 0.00454262, 0.00477313, 0.00355618,
        0.00179975, 0.00519071, 0.00539414, 0.00562861, 0.00646756,
        0.0066315 , 0.00730387, 0.00334007, 0.00563384, 0.00612996,
        0.00436176]),
 array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15]))