Add spatial hashing (#1169)

I attempted to stay close to the KD-Tree and BallTree API. However, the
purpose of the SpatialHash class is to help locate the faces that a
list of coordinates lie within. Because of this, the intent of the
`query` function for the SpatialHash is a bit different than that for
the KD and Ball trees.

Additional support routines are provided for assessing whether a point
is inside or outside a polygon. This is done by calculating the
barycentric coordinates of a point in a convex polygon. The method is
defined so that the sum of the barycentric coordinates is exactly one.
Because of this, if any of the barycentric coordinates are negative, we
immediately know that a point is outside the polygon.

All coordinate calculations are done using (lon,lat) in radians for the
SpatialHash class. Cartesian coordinates are not supported in this
commit.

* Add `get_spatial_hash` to grid api

* Add example usage in spatial hashing docs

* Add spatialhashing to userguide

* Add basic tests

* Fix call to get_spatialhash

* Fix formatting

* Prepend attributes intended for internal use with _

* Add a few patches for hashtable setup and query check

* Ensure that the lon bounds are sorted from low to high. This is needed
to make sure that the index looping for filling the hashtable actually
executes.

* All barycentric coordinates are calculated directly. Rather than
calculating the last coordinate to ensure they sum to one, we now
compute all coordinates. The coordinates are calculated now to align
with the indexing of the node ID's so that np.sum(bcoords*nodes) returns
the barycentric interpolation estimate of point. If the point is within
the element, then the barycentric interpolation estimate matches the
input coordinate (to machine precision). I've added this as a check.

* Fix docstrings for tests; add fesom test with particle list

* Add docstring for get_spatialhash

* remove empty notes

* Update uxarray/grid/grid.py

Co-authored-by: Philip Chmielowiec <[email protected]>

* Fix naming in test

* Update uxarray/grid/neighbors.py

Co-authored-by: Philip Chmielowiec <[email protected]>

* Update uxarray/grid/neighbors.py

Co-authored-by: Philip Chmielowiec <[email protected]>

* Update test/test_spatial_hashing.py

Co-authored-by: Philip Chmielowiec <[email protected]>

* Fix docstring for reconstruct field

Co-authored-by: Philip Chmielowiec <[email protected]>

* Update uxarray/grid/grid.py

Co-authored-by: Philip Chmielowiec <[email protected]>

* Add example to docstring

* Use waschpress points for generalized barycentric coordinates

Added dual and primal grids for mpas. The generalized waschpress points
give us the correct generalized barycentric coordinates that satisfy the
Lagrange property for barycentric interpolation in convex polygons.

* Switch to mpas example in notebook

* Update uxarray/grid/grid.py

Co-authored-by: Philip Chmielowiec <[email protected]>

* Retain output in notebook. Make requested formatting changes

* Fix formatting

* Change to smaller/simpler 4xhex grid

* Add get_spatial_hash to the Grid/Methods section

* Clear outputs from notebook.

* Fix docstring; we're not computing signed area any longer.

Instead, we're favoring a check of the predicted coordinates using
barycentric interpolation

* Update uxarray/grid/neighbors.py

Co-authored-by: Philip Chmielowiec <[email protected]>

* Update uxarray/grid/neighbors.py

Co-authored-by: Philip Chmielowiec <[email protected]>

* Add query optimization suggested by @philipc2

* Remove extra return statement

* Clip denominator in barycentric weights calculation

The denominator in the barycentric weights calculation (Weischell
weights) is clipped to be no smaller than `ERROR_TOLERANCE` from the
uxarray.constants module.

This change is made to fix an issue, reported anectdotally by @Philip2c
in #1169

* Clip triangle areas rather than denominator

When the denominator is clipped, rather than the individual triangle
areas, the Weischell weights are incorrect for the case when a query
point is on a face vertex.

I've added tests for query points on edges and vertices

---------

Co-authored-by: Philip Chmielowiec <[email protected]>

Author: fluidnumerics-joe

docs/api.rst

@@ -161,6 +161,7 @@ Methods
    Grid.calculate_total_face_area
    Grid.normalize_cartesian_coordinates
    Grid.construct_face_centers
+   Grid.get_spatial_hash
    Grid.get_faces_containing_point
 
 Inheritance of Xarray Functionality

docs/user-guide/spatial-hashing.ipynb

@@ -0,0 +1,106 @@
+{
+ "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 simple hexagonal grid."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "grid_path = \"../../test/meshfiles/ugrid/quad-hexagon/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 = uxgrid.get_spatial_hash().query([-0.1, -0.1])\n",
+    "print(f\"face ids : {face_ids}\")\n",
+    "print(f\"Barycentric Coordinates: {bcoords}\")"
+   ]
+  }
+ ],
+ "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,143 @@
+import os
+import numpy as np
+import pytest
+import xarray as xr
+from pathlib import Path
+import uxarray as ux
+from uxarray.constants import ERROR_TOLERANCE
+
+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_query_on_vertex():
+    """Verifies correct values when a query is made exactly on a vertex"""
+    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([0.0, 90.0])
+    assert face_ids[0] == 0
+    assert np.isclose(bcoords[0,0],1.0,atol=ERROR_TOLERANCE)
+    assert np.isclose(bcoords[0,1],0.0,atol=ERROR_TOLERANCE)
+    assert np.isclose(bcoords[0,2],0.0,atol=ERROR_TOLERANCE)
+
+
+def test_query_on_edge():
+    """Verifies correct values when a query is made exactly on an edge of a face"""
+    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([0.0, 0.0])
+    assert face_ids[0] == 0
+    assert np.isclose(bcoords[0,0],0.5,atol=ERROR_TOLERANCE)
+    assert np.isclose(bcoords[0,1],0.0,atol=ERROR_TOLERANCE)
+    assert np.isclose(bcoords[0,2],0.5,atol=ERROR_TOLERANCE)
+
+
+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

@@ -80,6 +80,7 @@
 from uxarray.grid.neighbors import (
     BallTree,
     KDTree,
+    SpatialHash,
     _populate_edge_face_distances,
     _populate_edge_node_distances,
 )
@@ -255,6 +256,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
@@ -1762,6 +1764,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."""