diff --git a/docs/api.rst b/docs/api.rst index 2c9170447..ca024f4d1 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -160,6 +160,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 diff --git a/docs/user-guide/spatial-hashing.ipynb b/docs/user-guide/spatial-hashing.ipynb new file mode 100644 index 000000000..ea5f6a02b --- /dev/null +++ b/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 +} diff --git a/docs/userguide.rst b/docs/userguide.rst index 5f20cf4a5..feb137525 100644 --- a/docs/userguide.rst +++ b/docs/userguide.rst @@ -49,6 +49,9 @@ These user guides provide detailed explanations of the core functionality in UXa `Subsetting `_ Select specific regions of a grid +`Spatial Hashing `_ + Use spatial hashing to locate the faces a list of points reside in. + `Cross-Sections `_ Select cross-sections of a grid diff --git a/test/test_spatial_hashing.py b/test/test_spatial_hashing.py new file mode 100644 index 000000000..da64f4d49 --- /dev/null +++ b/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 diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index 44bdc37fa..cfd882e97 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -80,6 +80,7 @@ from uxarray.grid.neighbors import ( BallTree, KDTree, + SpatialHash, _populate_edge_face_distances, _populate_edge_node_distances, ) @@ -252,6 +253,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.""" diff --git a/uxarray/grid/neighbors.py b/uxarray/grid/neighbors.py index 2dcb70602..1983d3148 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -781,6 +781,207 @@ def coordinates(self, value): ) +class SpatialHash: + """Custom data structure that is used for performing grid searches using Spatial Hashing. This class constructs an overlying + uniformly spaced structured grid, called the "hash grid" on top an unstructured grid. Faces in the unstructured grid are related + to the cells in the hash grid by determining the hash cells the bounding box of the unstructured face cells overlap with. + + Parameters + ---------- + grid : ux.Grid + Source grid used to construct the hash grid and hash table + reconstruct : bool, default=False + If true, reconstructs the spatial hash + """ + + def __init__( + self, + grid, + reconstruct: bool = False, + ): + self._source_grid = grid + self._nelements = self._source_grid.n_face + + self.reconstruct = reconstruct + + # Hash grid size + self._dh = self._hash_cell_size() + # Lower left corner of the hash grid + self._xmin = np.deg2rad(self._source_grid.node_lon.min().to_numpy()) - self._dh + self._ymin = np.deg2rad(self._source_grid.node_lat.min().to_numpy()) - self._dh + self._xmax = np.deg2rad(self._source_grid.node_lon.max().to_numpy()) + self._dh + self._ymax = np.deg2rad(self._source_grid.node_lat.max().to_numpy()) + self._dh + # Number of x points in the hash grid; used for + # array flattening + Lx = self._xmax - self._xmin + Ly = self._ymax - self._ymin + self._nx = int(np.ceil(Lx / self._dh)) + self._ny = int(np.ceil(Ly / self._dh)) + + # Generate the mapping from the hash indices to unstructured grid elements + self._face_hash_table = None + self._face_hash_table = self._initialize_face_hash_table() + + def _hash_cell_size(self): + """Computes the size of the hash cells from the source grid. + The hash cell size is set to 1/2 of the median edge length in the grid (in radians)""" + return self._source_grid.edge_node_distances.median().to_numpy() * 0.5 + + def _hash_index2d(self, coords): + """Computes the 2-d hash index (i,j) for the location (x,y), where x and y are given in spherical + coordinates (in degrees)""" + + i = ((coords[:, 0] - self._xmin) / self._dh).astype(INT_DTYPE) + j = ((coords[:, 1] - self._ymin) / self._dh).astype(INT_DTYPE) + return i, j + + def _hash_index(self, coords): + """Computes the flattened hash index for the location (x,y), where x and y are given in spherical + coordinates (in degrees). The single dimensioned hash index orders the flat index with all of the + i-points first and then all the j-points.""" + i, j = self._hash_index2d(coords) + return i + self._nx * j + + def _initialize_face_hash_table(self): + """Create a mapping that relates unstructured grid faces to hash indices by determining + which faces overlap with which hash cells""" + + if self._face_hash_table is None or self.reconstruct: + index_to_face = [[] for i in range(self._nx * self._ny)] + lon_bounds = np.sort(self._source_grid.face_bounds_lon.to_numpy(), 1) + lat_bounds = self._source_grid.face_bounds_lat.to_numpy() + + coords = np.column_stack( + ( + np.deg2rad(lon_bounds[:, 0].flatten()), + np.deg2rad(lat_bounds[:, 0].flatten()), + ) + ) + i1, j1 = self._hash_index2d(coords) + coords = np.column_stack( + ( + np.deg2rad(lon_bounds[:, 1].flatten()), + np.deg2rad(lat_bounds[:, 1].flatten()), + ) + ) + i2, j2 = self._hash_index2d(coords) + + for eid in range(self._source_grid.n_face): + for j in range(j1[eid], j2[eid] + 1): + for i in range(i1[eid], i2[eid] + 1): + index_to_face[i + self._nx * j].append(eid) + + return index_to_face + + def query( + self, + coords: Union[np.ndarray, list, tuple], + in_radians: Optional[bool] = False, + tol: Optional[float] = 1e-6, + ): + """Queries the hash table. + + Parameters + ---------- + coords : array_like + coordinate pairs in degrees (lon, lat) to query + in_radians : bool, optional + if True, queries assuming coords are inputted in radians, not degrees. Only applies to spherical coords + + + Returns + ------- + faces : ndarray of shape (coords.shape[0]), dtype=INT_DTYPE + Face id's in the self._source_grid where each coords element is found. When a coords element is not found, the + corresponding array entry in faces is set to -1. + bcoords : ndarray of shape (coords.shape[0], self._source_grid.n_max_face_nodes), dtype=double + Barycentric coordinates of each coords element + """ + + coords = _prepare_xy_for_query(coords, in_radians, distance_metric=None) + num_coords = coords.shape[0] + max_nodes = self._source_grid.n_max_face_nodes + + # Preallocate results + bcoords = np.zeros((num_coords, max_nodes), dtype=np.double) + faces = np.full(num_coords, -1, dtype=INT_DTYPE) + + # Get grid variables + n_nodes_per_face = self._source_grid.n_nodes_per_face.to_numpy() + face_node_connectivity = self._source_grid.face_node_connectivity.to_numpy() + + # Precompute radian values for node coordinates: + node_lon = np.deg2rad(self._source_grid.node_lon.to_numpy()) + node_lat = np.deg2rad(self._source_grid.node_lat.to_numpy()) + + # Get the list of candidate faces for each coordinate + candidate_faces = [ + self._face_hash_table[pid] for pid in self._hash_index(coords) + ] + + for i, (coord, candidates) in enumerate(zip(coords, candidate_faces)): + for face_id in candidates: + n_nodes = n_nodes_per_face[face_id] + node_ids = face_node_connectivity[face_id, :n_nodes] + nodes = np.column_stack((node_lon[node_ids], node_lat[node_ids])) + bcoord = np.asarray(_barycentric_coordinates(nodes, coord)) + err = abs(np.dot(bcoord, nodes[:, 0]) - coord[0]) + abs( + np.dot(bcoord, nodes[:, 1]) - coord[1] + ) + if (bcoord >= 0).all() and err < tol: + faces[i] = face_id + bcoords[i, :n_nodes] = bcoord[:n_nodes] + break + + return faces, bcoords + + +@njit(cache=True) +def _triangle_area(A, B, C): + """ + Compute the area of a triangle given by three points. + """ + return 0.5 * abs(A[0] * (B[1] - C[1]) + B[0] * (C[1] - A[1]) + C[0] * (A[1] - B[1])) + + +@njit(cache=True) +def _barycentric_coordinates(nodes, point): + """ + Compute the barycentric coordinates of a point P inside a convex polygon using area-based weights. + So that this method generalizes to n-sided polygons, we use the Waschpress points as the generalized + barycentric coordinates, which is only valid for convex polygons. + + Parameters + ---------- + nodes : numpy.ndarray + Spherical coordinates (lon,lat) of each corner node of a face + point : numpy.ndarray + Spherical coordinates (lon,lat) of the point + Returns + ------- + numpy.ndarray + Barycentric coordinates corresponding to each vertex. + + """ + n = len(nodes) + sum_wi = 0 + w = [] + + for i in range(0, n): + vim1 = nodes[i - 1] + vi = nodes[i] + vi1 = nodes[(i + 1) % n] + a0 = _triangle_area(vim1, vi, vi1) + a1 = _triangle_area(point, vim1, vi) + a2 = _triangle_area(point, vi, vi1) + sum_wi += a0 / (a1 * a2) + w.append(a0 / (a1 * a2)) + + barycentric_coords = [w_i / sum_wi for w_i in w] + + return barycentric_coords + + def _prepare_xy_for_query(xy, use_radians, distance_metric): """Prepares xy coordinates for query with the sklearn BallTree or KDTree."""