Add spatial hashing

docs/api.rst

@@ -159,6 +159,8 @@ Methods
    Grid.calculate_total_face_area
    Grid.normalize_cartesian_coordinates
    Grid.construct_face_centers
+   Grid.get_spatial_hash
+
 
 Inheritance of Xarray Functionality
 -----------------------------------

docs/userguide.rst

@@ -49,6 +49,9 @@ These user guides provide detailed explanations of the core functionality in UXa
 `Subsetting <user-guide/subset.ipynb>`_
  Select specific regions of a grid
 
+`Spatial Hashing <user-guide/spatial-hashing.ipynb>`_
+ Use spatial hashing to locate the faces a list of points reside in.
+
 `Cross-Sections <user-guide/cross-sections.ipynb>`_
  Select cross-sections of a grid
 

test/test_spatial_hashing.py

@@ -0,0 +1,115 @@
+import os
+import numpy as np
+import pytest
+import xarray as xr
+from pathlib import Path
+import uxarray as ux
+
+current_path = Path(os.path.dirname(os.path.realpath(__file__)))
+
+# Sample grid file paths
+gridfile_CSne8 = current_path / "meshfiles" / "scrip" / "outCSne8" / "outCSne8.nc"
+gridfile_RLL1deg = current_path / "meshfiles" / "ugrid" / "outRLL1deg" / "outRLL1deg.ug"
+gridfile_RLL10deg_CSne4 = current_path / "meshfiles" / "ugrid" / "ov_RLL10deg_CSne4" / "ov_RLL10deg_CSne4.ug"
+gridfile_CSne30 = current_path / "meshfiles" / "ugrid" / "outCSne30" / "outCSne30.ug"
+gridfile_fesom = current_path / "meshfiles" / "ugrid" / "fesom" / "fesom.mesh.diag.nc"
+gridfile_geoflow = current_path / "meshfiles" / "ugrid" / "geoflow-small" / "grid.nc"
+gridfile_mpas = current_path / 'meshfiles' / "mpas" / "QU" / 'mesh.QU.1920km.151026.nc'
+
+grid_files = [gridfile_CSne8,
+              gridfile_RLL1deg,
+              gridfile_RLL10deg_CSne4,
+              gridfile_CSne30,
+              gridfile_fesom,
+              gridfile_geoflow,
+              gridfile_mpas]
+
+def test_construction():
+    """Tests the construction of the SpatialHash object"""
+    for grid_file in grid_files:
+        uxgrid = ux.open_grid(grid_file)
+        face_ids, bcoords = uxgrid.get_spatial_hash().query([0.9, 1.8])
+        assert face_ids.shape[0] == bcoords.shape[0]
+
+
+def test_is_inside():
+    """Verifies simple test for points inside and outside an element."""
+    verts = [(0.0, 90.0), (-180, 0.0), (0.0, -90)]
+    uxgrid = ux.open_grid(verts, latlon=True)
+    # Verify that a point outside the element returns a face id of -1
+    face_ids, bcoords = uxgrid.get_spatial_hash().query([90.0, 0.0])
+    assert face_ids[0] == -1
+    # Verify that a point inside the element returns a face id of 0
+    face_ids, bcoords = uxgrid.get_spatial_hash().query([-90.0, 0.0])
+
+    assert face_ids[0] == 0
+    assert np.allclose(bcoords[0], [0.25, 0.5, 0.25], atol=1e-06)
+
+def test_list_of_coords_simple():
+    """Verifies test using list of points inside and outside an element"""
+    verts = [(0.0, 90.0), (-180, 0.0), (0.0, -90)]
+    uxgrid = ux.open_grid(verts, latlon=True)
+
+    coords = [[90.0, 0.0], [-90.0, 0.0]]
+    face_ids, bcoords = uxgrid.get_spatial_hash().query(coords)
+    assert face_ids[0] == -1
+    assert face_ids[1] == 0
+    assert np.allclose(bcoords[1], [0.25, 0.5, 0.25], atol=1e-06)
+
+def test_list_of_coords_fesom():
+    """Verifies test using list of points on the fesom grid"""
+    uxgrid = ux.open_grid(gridfile_fesom)
+
+    num_particles = 20
+    coords = np.zeros((num_particles,2))
+    x_min = 1.0
+    x_max = 3.0
+    y_min = 2.0
+    y_max = 10.0
+    for k in range(num_particles):
+        coords[k,0] = np.deg2rad(np.random.uniform(x_min, x_max))
+        coords[k,1] = np.deg2rad(np.random.uniform(y_min, y_max))
+    face_ids, bcoords = uxgrid.get_spatial_hash().query(coords)
+    assert len(face_ids) == num_particles
+    assert bcoords.shape[0] == num_particles
+    assert bcoords.shape[1] == 3
+    assert np.all(face_ids >= 0) # All particles should be inside an element
+
+def test_list_of_coords_mpas_dual():
+    """Verifies test using list of points on the dual MPAS grid"""
+    uxgrid = ux.open_grid(gridfile_mpas, use_dual=True)
+
+    num_particles = 20
+    coords = np.zeros((num_particles,2))
+    x_min = -40.0
+    x_max = 40.0
+    y_min = -20.0
+    y_max = 20.0
+    for k in range(num_particles):
+        coords[k,0] = np.deg2rad(np.random.uniform(x_min, x_max))
+        coords[k,1] = np.deg2rad(np.random.uniform(y_min, y_max))
+    face_ids, bcoords = uxgrid.get_spatial_hash().query(coords)
+    assert len(face_ids) == num_particles
+    assert bcoords.shape[0] == num_particles
+    assert bcoords.shape[1] == 3 # max sides of an element
+    assert np.all(face_ids >= 0) # All particles should be inside an element
+
+
+def test_list_of_coords_mpas_primal():
+    """Verifies test using list of points on the primal MPAS grid"""
+    uxgrid = ux.open_grid(gridfile_mpas, use_dual=False)
+
+    num_particles = 20
+    coords = np.zeros((num_particles,2))
+    x_min = -40.0
+    x_max = 40.0
+    y_min = -20.0
+    y_max = 20.0
+    for k in range(num_particles):
+        coords[k,0] = np.deg2rad(np.random.uniform(x_min, x_max))
+        coords[k,1] = np.deg2rad(np.random.uniform(y_min, y_max))
+    face_ids, bcoords = uxgrid.get_spatial_hash().query(coords)
+    assert len(face_ids) == num_particles
+    assert bcoords.shape[0] == num_particles
+    assert bcoords.shape[1] == 6 # max sides of an element
+    assert np.all(face_ids >= 0) # All particles should be inside an element

uxarray/grid/grid.py

@@ -74,6 +74,7 @@
 from uxarray.grid.neighbors import (
     BallTree,
     KDTree,
+    SpatialHash,
     _populate_edge_face_distances,
     _populate_edge_node_distances,
 )
@@ -246,6 +247,7 @@ def __init__(
         # initialize cached data structures (nearest neighbor operations)
         self._ball_tree = None
         self._kd_tree = None
+        self._spatialhash = None
 
         # flag to track if coordinates are normalized
         self._normalized = None
@@ -1749,6 +1751,44 @@ def get_kd_tree(
 
         return self._kd_tree
 
+    def get_spatial_hash(
+        self,
+        reconstruct: bool = False,
+    ):
+        """Get the SpatialHash data structure of this Grid that allows for
+        fast face search queries. Face searches are used to find the faces that
+        a list of points, in spherical coordinates, are contained within.
+
+        Parameters
+        ----------
+        reconstruct : bool, default=False
+            If true, reconstructs the spatial hash
+
+        Returns
+        -------
+        self._spatialhash : grid.Neighbors.SpatialHash
+            SpatialHash instance
+
+        Examples
+        --------
+        Open a grid from a file path:
+
+        >>> import uxarray as ux
+        >>> uxgrid = ux.open_grid("grid_filename.nc")
+
+        Obtain SpatialHash instance:
+
+        >>> spatial_hash = uxgrid.get_spatial_hash()
+
+        Query to find the face a point lies within in addition to its barycentric coordinates:
+
+        >>> face_ids, bcoords = spatial_hash.query([0.0, 0.0])
+        """
+        if self._spatialhash is None or reconstruct:
+            self._spatialhash = SpatialHash(self, reconstruct)
+
+        return self._spatialhash
+
     def copy(self):
         """Returns a deep copy of this grid."""