Add spatial hashing

docs/user-guide/spatial-hashing.ipynb

@@ -0,0 +1,104 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Spatial Hashing\n",
+    "UXArray supports spatial hashing, which a fast search method for finding the face that a point lies within on an unstructured grid. This can be useful for building particle-in-cell interpolators for Lagrangian fluid flow analysis. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import uxarray as ux"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "For this example we will be using a FESOM2 grid"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "grid_path = \"../../test/meshfiles/fesom/soufflet-netcdf/grid.nc\"\n",
+    "uxgrid = ux.open_grid(grid_path)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "## SpatialHash\n",
+    "UXArray `SpatialHash` class provides an API for locating the unstructured grid faces that a list of $(x,y)$ points (in spherical coordinates) lie within. This search is conducted by relating unstructured grid faces and arbitrary physical positions to a uniformly spaced structured grid (the hash grid). The relationship between the unstructured grid elements and the hash grid is maintained in a hash table. \n",
+    "\n",
+    "### Parameters\n",
+    "The `SpatialHash` class can be accessed through the `Grid.get_spatial_hash(reconstruct)` method, which takes in the following parameters:\n",
+    "* `reconstruct` is a bool variable that allows the user to reconstruct the spatial hash structure. As default for performance, if a user calls `get_spatial_hash` and a spatial hash structure has already been created, it will simply use that one. If `reconstruct` is set to `True`, it will override this and reconstruct the spatial hash structure. The default parameter is `False`.\n",
+    "\n",
+    "\n",
+    "### Constructing a Spatial Hash\n",
+    "We can store the SpatialHash data structure in a variable, which allows us to access the spatial hash in a simple way for queries."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "spatial_hash = uxgrid.get_spatial_hash(\n",
+    "    reconstruct=\"False\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Query\n",
+    "Now we can use that variable to query for the face that a point lies within in addition to its barycentric coordinates. The first parameter is the point from which to do the search, which must be in spherical coordinates."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "face_ids, bcoords = spatial_hash.query([0.9, 1.8])"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "uxarray-dev",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.16"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

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,76 @@
+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

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,29 @@ 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 tree
+
+        Returns
+        -------
+        self._spatialhash : grid.Neighbors.SpatialHash
+            SpatialHash instance
+        """
+        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."""