diff --git a/parcels/_index_search.py b/parcels/_index_search.py index 397d372164..7c1f11ef2d 100644 --- a/parcels/_index_search.py +++ b/parcels/_index_search.py @@ -275,11 +275,9 @@ def _search_indices_curvilinear_2d( grid: XGrid, y: float, x: float, yi_guess: int | None = None, xi_guess: int | None = None ): yi, xi = yi_guess, xi_guess - if yi is None: - yi = int(grid.ydim / 2) - 1 - - if xi is None: - xi = int(grid.xdim / 2) - 1 + if yi is None or xi is None: + faces = grid.get_spatial_hash().query(np.column_stack((y, x))) + yi, xi = faces[0] xsi = eta = -1.0 invA = np.array( @@ -350,7 +348,6 @@ def _search_indices_curvilinear_2d( if not ((0 <= xsi <= 1) and (0 <= eta <= 1)): _raise_field_sampling_error(y, x) - return (yi, eta, xi, xsi) diff --git a/parcels/spatialhash.py b/parcels/spatialhash.py new file mode 100644 index 0000000000..57942610e5 --- /dev/null +++ b/parcels/spatialhash.py @@ -0,0 +1,310 @@ +import numpy as np + + +class SpatialHash: + """Custom data structure that is used for performing grid searches using Spatial Hashing. This class constructs an overlying + uniformly spaced rectilinear grid, called the "hash grid" on top parcels.xgrid.XGrid. It is particularly useful for grid searching + on curvilinear grids. Faces in the Xgrid 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 : parcels.xgrid.XGrid + Source grid used to construct the hash grid and hash table + reconstruct : bool, default=False + If true, reconstructs the spatial hash + + Note + ---- + Does not currently support queries on periodic elements. + """ + + def __init__( + self, + grid, + reconstruct=False, + ): + # TODO : Enforce grid to be an instance of parcels.xgrid.XGrid + # Currently, this is not done due to circular import with parcels.xgrid + + self._source_grid = grid + self.reconstruct = reconstruct + + # Hash grid size + self._dh = self._hash_cell_size() + + # Lower left corner of the hash grid + lon_min = self._source_grid.lon.min() + lat_min = self._source_grid.lat.min() + lon_max = self._source_grid.lon.max() + lat_max = self._source_grid.lat.max() + + # Get corner vertices of each face + self._xbound = np.stack( + ( + self._source_grid.lon[:-1, :-1], + self._source_grid.lon[:-1, 1:], + self._source_grid.lon[1:, 1:], + self._source_grid.lon[1:, :-1], + ), + axis=-1, + ) + self._ybound = np.stack( + ( + self._source_grid.lat[:-1, :-1], + self._source_grid.lat[:-1, 1:], + self._source_grid.lat[1:, 1:], + self._source_grid.lat[1:, :-1], + ), + axis=-1, + ) + + self._xmin = lon_min - self._dh + self._ymin = lat_min - self._dh + self._xmax = lon_max + self._dh + self._ymax = lat_max + 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 square root of the median cell area + """ + return np.sqrt(np.median(_planar_quad_area(self._source_grid.lat, self._source_grid.lon))) * 0.5 + + def _hash_index2d(self, coords): + """Computes the 2-d hash index (i,j) for the location (x,y), where x and y is the same units + as the source grid coordinates + """ + # Wrap longitude to [-180, 180] + if self._source_grid.mesh == "spherical": + lon = (coords[:, 1] + 180.0) % (360.0) - 180.0 + else: + lon = coords[:, 1] + i = ((lon - self._xmin) / self._dh).astype(np.int32) + j = ((coords[:, 0] - self._ymin) / self._dh).astype(np.int32) + return j, i + + 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. + """ + j, i = self._hash_index2d(coords) + return i + self._nx * j + + def _grid_ji_for_eid(self, eid): + """Returns the (i,j) grid coordinates for the given element id (eid)""" + j = eid // (self._source_grid.xdim) + i = eid - j * (self._source_grid.xdim) + return j, i + + 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)] + # Get the bounds of each curvilinear faces + lat_bounds, lon_bounds = _curvilinear_grid_facebounds( + self._source_grid.lat, + self._source_grid.lon, + ) + coords = np.stack( + ( + lat_bounds[:, :, 0].flatten(), + lon_bounds[:, :, 0].flatten(), + ), + axis=-1, + ) + yi1, xi1 = self._hash_index2d(coords) + coords = np.stack( + ( + lat_bounds[:, :, 1].flatten(), + lon_bounds[:, :, 1].flatten(), + ), + axis=-1, + ) + yi2, xi2 = self._hash_index2d(coords) + nface = (self._source_grid.xdim) * (self._source_grid.ydim) + for eid in range(nface): + for j in range(yi1[eid], yi2[eid] + 1): + if xi1[eid] <= xi2[eid]: + # Normal case, no wrap + for i in range(xi1[eid], xi2[eid] + 1): + index_to_face[(i % self._nx) + self._nx * j].append(eid) + else: + # Wrap-around case + for i in range(xi1[eid], self._nx): + index_to_face[(i % self._nx) + self._nx * j].append(eid) + for i in range(0, xi2[eid] + 1): + index_to_face[(i % self._nx) + self._nx * j].append(eid) + return index_to_face + + def query( + self, + coords, + tol=1e-6, + ): + """Queries the hash table. + + Parameters + ---------- + coords : array_like + coordinate pairs in degrees (lat, lon) to query. + + + Returns + ------- + faces : ndarray of shape (coords.shape[0]), dtype=np.int32 + 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. + """ + num_coords = coords.shape[0] + + # Preallocate results + faces = np.full((num_coords, 2), -1, dtype=np.int32) + + # 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, strict=False)): + for face_id in candidates: + yi, xi = self._grid_ji_for_eid(face_id) + nodes = np.stack( + ( + self._ybound[yi, xi, :], + self._xbound[yi, xi, :], + ), + axis=-1, + ) + + 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, :] = [yi, xi] + break + + return faces + + +def _triangle_area(A, B, C): + """Compute the area of a triangle given by three points.""" + d1 = B - A + d2 = C - A + d3 = np.cross(d1, d2) + return 0.5 * np.linalg.norm(d3) + + +def _barycentric_coordinates(nodes, point, min_area=1e-8): + """ + 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 (lat,lon) of each corner node of a face + point : numpy.ndarray + Spherical coordinates (lat,lon) 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 = max(_triangle_area(point, vim1, vi), min_area) + a2 = max(_triangle_area(point, vi, vi1), min_area) + 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 _planar_quad_area(lat, lon): + """Computes the area of each quadrilateral face in a curvilinear grid. + The lon and lat arrays are assumed to be 2D arrays of points with dimensions (n_y, n_x). + The area is computed using the Shoelace formula. + This method is only used during hashgrid construction to determine the size of the hash cells. + + Parameters + ---------- + lon : np.ndarray + 2D array of shape (n_y, n_x) containing the longitude of each corner node of the curvilinear grid. + lat : np.ndarray + 2D array of shape (n_y, n_x) containing the latitude of each corner node of the curvilinear grid. + + Returns + ------- + area : np.ndarray + 2D array of shape (n_y-1, n_x-1) containing the area of each quadrilateral face in the curvilinear grid. + """ + x0 = lon[:-1, :-1] + x1 = lon[:-1, 1:] + x2 = lon[1:, 1:] + x3 = lon[1:, :-1] + + y0 = lat[:-1, :-1] + y1 = lat[:-1, 1:] + y2 = lat[1:, 1:] + y3 = lat[1:, :-1] + + # Shoelace formula: 0.5 * |sum(x_i*y_{i+1} - x_{i+1}*y_i)| + area = 0.5 * np.abs(x0 * y1 + x1 * y2 + x2 * y3 + x3 * y0 - y0 * x1 - y1 * x2 - y2 * x3 - y3 * x0) + return area + + +def _curvilinear_grid_facebounds(lat, lon): + """Computes the bounds of each curvilinear face in the grid. + The lon and lat arrays are assumed to be 2D arrays of points with dimensions (n_y, n_x). + The bounds are for faces whose corner node vertices are defined by lat,lon. + Face(yi,xi) is surrounding by points (yi,xi), (yi,xi+1), (yi+1,xi+1), (yi+1,xi). + This method is only used during hashgrid construction to determine which curvilinear + faces overlap with which hash cells. + + Parameters + ---------- + lon : np.ndarray + 2D array of shape (n_y, n_x) containing the longitude of each corner node of the curvilinear grid. + lat : np.ndarray + 2D array of shape (n_y, n_x) containing the latitude of each corner node of the curvilinear grid. + + Returns + ------- + xbounds : np.ndarray + Array of shape (n_y-1, n_x-1, 2) containing the bounds of each face in the x-direction. + ybounds : np.ndarray + Array of shape (n_y-1, n_x-1, 2) containing the bounds of each face in the y-direction. + """ + xf = np.stack((lon[:-1, :-1], lon[:-1, 1:], lon[1:, 1:], lon[1:, :-1]), axis=-1) + xf_low = xf.min(axis=-1) + xf_high = xf.max(axis=-1) + xbounds = np.stack([xf_low, xf_high], axis=-1) + + yf = np.stack((lat[:-1, :-1], lat[:-1, 1:], lat[1:, 1:], lat[1:, :-1]), axis=-1) + yf_low = yf.min(axis=-1) + yf_high = yf.max(axis=-1) + ybounds = np.stack([yf_low, yf_high], axis=-1) + + return ybounds, xbounds diff --git a/parcels/uxgrid.py b/parcels/uxgrid.py index 0074309032..cc622feaf7 100644 --- a/parcels/uxgrid.py +++ b/parcels/uxgrid.py @@ -4,10 +4,9 @@ import numpy as np import uxarray as ux -from uxarray.grid.coordinates import _lonlat_rad_to_xyz -from uxarray.grid.neighbors import _barycentric_coordinates -from parcels.field import FieldOutOfBoundError # Adjust import as necessary +from parcels.field import FieldOutOfBoundError +from parcels.spatialhash import _barycentric_coordinates from parcels.xgrid import _search_1d_array from .basegrid import BaseGrid @@ -141,52 +140,20 @@ def _get_barycentric_coordinates_cartesian(self, y, x, fi): axis=-1, ) - bcoord = np.asarray(_barycentric_coordinates_cartesian(nodes, cart_coord)) + bcoord = np.asarray(_barycentric_coordinates(nodes, cart_coord)) return bcoord -def _barycentric_coordinates_cartesian(nodes, point, min_area=1e-8): +def _lonlat_rad_to_xyz( + lon, + lat, +): + """Converts Spherical latitude and longitude coordinates into Cartesian x, + y, z coordinates. """ - 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 - Cartesian coordinates (x,y,z) of each corner node of a face - point : numpy.ndarray - Cartesian coordinates (x,y,z) of the point - - Returns - ------- - numpy.ndarray - Barycentric coordinates corresponding to each vertex. + x = np.cos(lon) * np.cos(lat) + y = np.sin(lon) * np.cos(lat) + z = np.sin(lat) - """ - 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_cartesian(vim1, vi, vi1) - a1 = max(_triangle_area_cartesian(point, vim1, vi), min_area) - a2 = max(_triangle_area_cartesian(point, vi, vi1), min_area) - 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 _triangle_area_cartesian(A, B, C): - """Compute the area of a triangle given by three points.""" - d1 = B - A - d2 = C - A - d3 = np.cross(d1, d2) - return 0.5 * np.linalg.norm(d3) + return x, y, z diff --git a/parcels/xgrid.py b/parcels/xgrid.py index 95d72eb2eb..77a49d7406 100644 --- a/parcels/xgrid.py +++ b/parcels/xgrid.py @@ -9,6 +9,7 @@ from parcels import xgcm from parcels._index_search import _search_indices_curvilinear_2d from parcels.basegrid import BaseGrid +from parcels.spatialhash import SpatialHash from parcels.tools.statuscodes import FieldOutOfBoundError, FieldOutOfBoundSurfaceError _XGRID_AXES = Literal["X", "Y", "Z"] @@ -95,6 +96,7 @@ class XGrid(BaseGrid): def __init__(self, grid: xgcm.Grid, mesh="flat"): self.xgcm_grid = grid self.mesh = mesh + self._spatialhash = None ds = grid._ds if len(set(grid.axes) & {"X", "Y", "Z"}) > 0: # Only if spatial grid is >0D (see #2054 for further development) @@ -317,6 +319,32 @@ def _fpoint_info(self): return axis_position_mapping + def get_spatial_hash( + self, + reconstruct=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 + ---------- + global_grid : bool, default=False + If true, the hash grid is constructed using the domain [-pi,pi] x [-pi,pi] + reconstruct : bool, default=False + If true, reconstructs the spatial hash + + Returns + ------- + self._spatialhash : parcels.spatialhash.SpatialHash + SpatialHash instance + + """ + if self._spatialhash is None or reconstruct: + self._spatialhash = SpatialHash(self, reconstruct) + + return self._spatialhash + def get_axis_from_dim_name(axes: _XGCM_AXES, dim: str) -> _XGCM_AXIS_DIRECTION | None: """For a given dimension name in a grid, returns the direction axis it is on.""" diff --git a/tests/v4/test_spatialhash.py b/tests/v4/test_spatialhash.py new file mode 100644 index 0000000000..9d958507f7 --- /dev/null +++ b/tests/v4/test_spatialhash.py @@ -0,0 +1,9 @@ +from parcels._datasets.structured.generic import datasets +from parcels.xgrid import XGrid + + +def test_spatialhash_init(): + ds = datasets["2d_left_rotated"] + grid = XGrid.from_dataset(ds) + spatialhash = grid.get_spatial_hash() + assert spatialhash is not None