KDTree Data Structure for Center and Corner Nodes

[aaronzedwick]

Closes #496

Overview

Adds a KDTree data structure for center nodes (Mesh2_face_cart_x, Mesh2_face_cart_y, Mesh2_face_cart_z) and corner nodes (Mesh2_node_cart_x, Mesh2_node_cart_y, Mesh2_node_cart_z) for grid files. This can be used for nearest neighbor queries as well as querying in a specific radius. This implementation mirrors the BallTree data structure #406 and uses sklearn.neighbors.KDTree with a minkowski distance for querying.

Expected Usage

import uxarray as ux

grid_path = "/path/to/grid.nc"

uxgrid = ux.open_grid(grid_path)

# find the nearest neighbor to (0.0, 0.0, 1.0) using corner nodes
d, ind = uxgrid.get_kd_tree(tree_type="nodes").query([0.0, 0.0, 1.0])

# find the 5 nearest neighbors to (0.0, 0.0, 1.0) using face centers
d, ind = uxgrid.get_kd_tree(tree_type="face centers").query([1.0, 0.0, 0.0], k=5)

# find all face centers within a radius of 80.00 from (0.0, 0.0, 1.0)
d, ind = uxgrid.get_kd_tree(tree_type="face centers").query_radius([0.0, 0.0, 1.0], r=80.00)

PR Checklist

General

• An issue is linked created and linked
• Add appropriate labels
• Filled out Overview and Expected Usage (if applicable) sections

Testing

• Adequate tests are created if there is new functionality
• Tests cover all possible logical paths in your function
• Tests are not too basic (such as simply calling a function and nothing else)

Documentation

• Docstrings have been added to all new functions
• Docstrings have updated with any function changes
• Internal functions have a preceding underscore (_) and have been added to docs/internal_api/index.rst
• User functions have been added to docs/user_api/index.rst

Examples

• Any new notebook examples added to docs/examples/ folder
• Clear the output of all cells before committing
• New notebook files added to docs/examples.rst toctree
• New notebook files added to new entry in docs/gallery.yml with appropriate thumbnail photo in docs/_static/thumbnails/

[aaronzedwick]

@philipc2 I think I updated the user api incorrectly causing the docs to fail, anything you see that would be causing the problem? I just added what you did from the balltree, but using the kdtree instead.

[philipc2]

Nice work! Please run the pre-commit though!