diff --git a/parcels/_datasets/unstructured/generic.py b/parcels/_datasets/unstructured/generic.py index edd432a96..1d5b456dc 100644 --- a/parcels/_datasets/unstructured/generic.py +++ b/parcels/_datasets/unstructured/generic.py @@ -218,10 +218,9 @@ def _fesom2_square_delaunay_antimeridian(): All fields are placed on location consistent with FESOM2 variable placement conventions """ lon, lat = np.meshgrid( - np.linspace(-210.0, -150.0, Nx, dtype=np.float32), np.linspace(0, 60.0, Nx, dtype=np.float32) + np.linspace(-210.0, -150.0, Nx, dtype=np.float32), np.linspace(-40.0, 40.0, Nx, dtype=np.float32) ) # wrap longitude from [-180,180] - lon = np.where(lon < -180, lon + 360, lon) lon_flat = lon.ravel() lat_flat = lat.ravel() zf = np.linspace(0.0, 1000.0, 10, endpoint=True, dtype=np.float32) # Vertical element faces @@ -231,7 +230,10 @@ def _fesom2_square_delaunay_antimeridian(): # mask any point on one of the boundaries mask = ( - np.isclose(lon_flat, 0.0) | np.isclose(lon_flat, 60.0) | np.isclose(lat_flat, 0.0) | np.isclose(lat_flat, 60.0) + np.isclose(lon_flat, -210.0) + | np.isclose(lon_flat, -150.0) + | np.isclose(lat_flat, -40.0) + | np.isclose(lat_flat, 40.0) ) boundary_points = np.flatnonzero(mask) diff --git a/parcels/_index_search.py b/parcels/_index_search.py index 574a31159..08d73e560 100644 --- a/parcels/_index_search.py +++ b/parcels/_index_search.py @@ -128,3 +128,150 @@ def _search_indices_curvilinear_2d( eta = coords[:, 1] return (yi, eta, xi, xsi) + + +def uxgrid_point_in_cell(grid, y: np.ndarray, x: np.ndarray, yi: np.ndarray, xi: np.ndarray): + """Check if points are inside the grid cells defined by the given face indices. + + Parameters + ---------- + grid : ux.grid.Grid + The uxarray grid object containing the unstructured grid data. + y : np.ndarray + Array of latitudes of the points to check. + x : np.ndarray + Array of longitudes of the points to check. + yi : np.ndarray + Array of face indices corresponding to the points. + xi : np.ndarray + Not used, but included for compatibility with other search functions. + + Returns + ------- + is_in_cell : np.ndarray + An array indicating whether each point is inside (1) or outside (0) the corresponding cell. + coords : np.ndarray + Barycentric coordinates of the points within their respective cells. + """ + if grid._mesh == "spherical": + lon_rad = np.deg2rad(x) + lat_rad = np.deg2rad(y) + x_cart, y_cart, z_cart = _latlon_rad_to_xyz(lat_rad, lon_rad) + points = np.column_stack((x_cart.flatten(), y_cart.flatten(), z_cart.flatten())) + + # Get the vertex indices for each face + nids = grid.uxgrid.face_node_connectivity[yi].values + face_vertices = np.stack( + ( + grid.uxgrid.node_x[nids.ravel()].values.reshape(nids.shape), + grid.uxgrid.node_y[nids.ravel()].values.reshape(nids.shape), + grid.uxgrid.node_z[nids.ravel()].values.reshape(nids.shape), + ), + axis=-1, + ) + else: + nids = grid.uxgrid.face_node_connectivity[yi].values + face_vertices = np.stack( + ( + grid.uxgrid.node_lon[nids.ravel()].values.reshape(nids.shape), + grid.uxgrid.node_lat[nids.ravel()].values.reshape(nids.shape), + ), + axis=-1, + ) + points = np.stack((x, y)) + + M = len(points) + + is_in_cell = np.zeros(M, dtype=np.int32) + + coords = _barycentric_coordinates(face_vertices, points) + is_in_cell = np.where(np.all((coords >= -1e-6) & (coords <= 1 + 1e-6), axis=1), 1, 0) + + return is_in_cell, coords + + +def _triangle_area(A, B, C): + """Compute the area of a triangle given by three points.""" + d1 = B - A + d2 = C - A + if A.shape[-1] == 2: + # 2D case: cross product reduces to scalar z-component + cross = d1[..., 0] * d2[..., 1] - d1[..., 1] * d2[..., 0] + area = 0.5 * np.abs(cross) + elif A.shape[-1] == 3: + # 3D case: full vector cross product + cross = np.cross(d1, d2) + area = 0.5 * np.linalg.norm(cross, axis=-1) + else: + raise ValueError(f"Expected last dim=2 or 3, got {A.shape[-1]}") + + return area + # d3 = np.cross(d1, d2, axis=-1) + # breakpoint() + # return 0.5 * np.linalg.norm(d3, axis=-1) + + +def _barycentric_coordinates(nodes, points, 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 + Polygon verties per query of shape (M, 3, 2/3) where M is the number of query points. The second dimension corresponds to the number + of vertices + The last dimension can be either 2 or 3, where 3 corresponds to the (z, y, x) coordinates of each vertex and 2 corresponds to the + (lat, lon) coordinates of each vertex. + + points : numpy.ndarray + Spherical coordinates of the point (M,2/3) where M is the number of query points. + + Returns + ------- + numpy.ndarray + Barycentric coordinates corresponding to each vertex. + + """ + M, K = nodes.shape[:2] + + # roll(-1) to get vi+1, roll(+1) to get vi-1 + vi = nodes # (M,K,2) + vi1 = np.roll(nodes, shift=-1, axis=1) # (M,K,2) + vim1 = np.roll(nodes, shift=+1, axis=1) # (M,K,2) + + # a0 = area(v_{i-1}, v_i, v_{i+1}) + a0 = _triangle_area(vim1, vi, vi1) # (M,K) + + # a1 = area(P, v_{i-1}, v_i); a2 = area(P, v_i, v_{i+1}) + P = points[:, None, :] # (M,1,2) -> (M,K,2) + a1 = _triangle_area(P, vim1, vi) + a2 = _triangle_area(P, vi, vi1) + + # clamp tiny denominators for stability + a1c = np.maximum(a1, min_area) + a2c = np.maximum(a2, min_area) + + wi = a0 / (a1c * a2c) # (M,K) + + sum_wi = wi.sum(axis=1, keepdims=True) # (M,1) + # Avoid 0/0: if sum_wi==0 (degenerate), keep zeros + with np.errstate(invalid="ignore", divide="ignore"): + bcoords = wi / sum_wi + + return bcoords + + +def _latlon_rad_to_xyz( + lat, + lon, +): + """Converts Spherical latitude and longitude coordinates into Cartesian x, + y, z coordinates. + """ + x = np.cos(lon) * np.cos(lat) + y = np.sin(lon) * np.cos(lat) + z = np.sin(lat) + + return x, y, z diff --git a/parcels/application_kernels/interpolation.py b/parcels/application_kernels/interpolation.py index f5fb35cfe..8caf5231a 100644 --- a/parcels/application_kernels/interpolation.py +++ b/parcels/application_kernels/interpolation.py @@ -657,8 +657,8 @@ def UXPiecewiseLinearNode( # The zi refers to the vertical layer index. The field in this routine are assumed to be defined at the vertical interface levels. # For interface zi, the interface indices are [zi, zi+1], so we need to use the values at zi and zi+1. # First, do barycentric interpolation in the lateral direction for each interface level - fzk = np.dot(field.data.values[ti, k, node_ids], bcoords) - fzkp1 = np.dot(field.data.values[ti, k + 1, node_ids], bcoords) + fzk = np.sum(field.data.values[ti, k, node_ids] * bcoords, axis=-1) + fzkp1 = np.sum(field.data.values[ti, k + 1, node_ids] * bcoords, axis=-1) # Then, do piecewise linear interpolation in the vertical direction zk = field.grid.z.values[k] diff --git a/parcels/basegrid.py b/parcels/basegrid.py index d758d5f17..37d56f612 100644 --- a/parcels/basegrid.py +++ b/parcels/basegrid.py @@ -6,6 +6,8 @@ import numpy as np +from parcels.spatialhash import SpatialHash + if TYPE_CHECKING: import numpy as np @@ -178,6 +180,32 @@ def get_axis_dim(self, axis: str) -> int: """ ... + 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) + + return self._spatialhash + def _unravel(dims, ei): """ diff --git a/parcels/spatialhash.py b/parcels/spatialhash.py index 8f5e3ffc6..5e4aaa4aa 100644 --- a/parcels/spatialhash.py +++ b/parcels/spatialhash.py @@ -1,7 +1,7 @@ import numpy as np import parcels -from parcels._index_search import GRID_SEARCH_ERROR, curvilinear_point_in_cell +from parcels._index_search import GRID_SEARCH_ERROR, _latlon_rad_to_xyz, curvilinear_point_in_cell, uxgrid_point_in_cell class SpatialHash: @@ -27,95 +27,144 @@ def __init__( ): if isinstance(grid, parcels.xgrid.XGrid): self._point_in_cell = curvilinear_point_in_cell + elif isinstance(grid, parcels.uxgrid.UxGrid): + self._point_in_cell = uxgrid_point_in_cell else: - raise NotImplementedError("SpatialHash only supports parcels.xgrid.XGrid grids at this time.") + raise ValueError("Expected `grid` to be a parcels.xgrid.XGrid or parcels.uxgrid.UxGrid") self._source_grid = grid self._bitwidth = bitwidth # Max integer to use per coordinate in quantization (10 bits = 0..1023) - if self._source_grid._mesh == "spherical": - # Boundaries of the hash grid are the unit cube - self._xmin = -1.0 - self._ymin = -1.0 - self._zmin = -1.0 - self._xmax = 1.0 - self._ymax = 1.0 - self._zmax = 1.0 # Compute the cell centers of the source grid (for now, assuming Xgrid) - lon = np.deg2rad(self._source_grid.lon) - lat = np.deg2rad(self._source_grid.lat) - x, y, z = _latlon_rad_to_xyz(lat, lon) - _xbound = np.stack( - ( - x[:-1, :-1], - x[:-1, 1:], - x[1:, 1:], - x[1:, :-1], - ), - axis=-1, - ) - _ybound = np.stack( - ( - y[:-1, :-1], - y[:-1, 1:], - y[1:, 1:], - y[1:, :-1], - ), - axis=-1, - ) - _zbound = np.stack( - ( - z[:-1, :-1], - z[:-1, 1:], - z[1:, 1:], - z[1:, :-1], - ), - axis=-1, - ) - # Compute centroid locations of each cells - self._xlow = np.min(_xbound, axis=-1) - self._xhigh = np.max(_xbound, axis=-1) - self._ylow = np.min(_ybound, axis=-1) - self._yhigh = np.max(_ybound, axis=-1) - self._zlow = np.min(_zbound, axis=-1) - self._zhigh = np.max(_zbound, axis=-1) - - else: - # Boundaries of the hash grid are the bounding box of the source grid - self._xmin = self._source_grid.lon.min() - self._xmax = self._source_grid.lon.max() - self._ymin = self._source_grid.lat.min() - self._ymax = self._source_grid.lat.max() - # setting min and max below is needed for mesh="flat" - self._zmin = 0.0 - self._zmax = 0.0 - x = self._source_grid.lon - y = self._source_grid.lat - - _xbound = np.stack( - ( - x[:-1, :-1], - x[:-1, 1:], - x[1:, 1:], - x[1:, :-1], - ), - axis=-1, - ) - _ybound = np.stack( - ( - y[:-1, :-1], - y[:-1, 1:], - y[1:, 1:], - y[1:, :-1], - ), - axis=-1, - ) - # Compute centroid locations of each cells - self._xlow = np.min(_xbound, axis=-1) - self._xhigh = np.max(_xbound, axis=-1) - self._ylow = np.min(_ybound, axis=-1) - self._yhigh = np.max(_ybound, axis=-1) - self._zlow = np.zeros_like(self._xlow) - self._zhigh = np.zeros_like(self._xlow) + if isinstance(grid, parcels.xgrid.XGrid): + if self._source_grid._mesh == "spherical": + # Boundaries of the hash grid are the unit cube + self._xmin = -1.0 + self._ymin = -1.0 + self._zmin = -1.0 + self._xmax = 1.0 + self._ymax = 1.0 + self._zmax = 1.0 # Compute the cell centers of the source grid (for now, assuming Xgrid) + lon = np.deg2rad(self._source_grid.lon) + lat = np.deg2rad(self._source_grid.lat) + x, y, z = _latlon_rad_to_xyz(lat, lon) + _xbound = np.stack( + ( + x[:-1, :-1], + x[:-1, 1:], + x[1:, 1:], + x[1:, :-1], + ), + axis=-1, + ) + _ybound = np.stack( + ( + y[:-1, :-1], + y[:-1, 1:], + y[1:, 1:], + y[1:, :-1], + ), + axis=-1, + ) + _zbound = np.stack( + ( + z[:-1, :-1], + z[:-1, 1:], + z[1:, 1:], + z[1:, :-1], + ), + axis=-1, + ) + # Compute centroid locations of each cells + self._xlow = np.min(_xbound, axis=-1) + self._xhigh = np.max(_xbound, axis=-1) + self._ylow = np.min(_ybound, axis=-1) + self._yhigh = np.max(_ybound, axis=-1) + self._zlow = np.min(_zbound, axis=-1) + self._zhigh = np.max(_zbound, axis=-1) + + else: + # Boundaries of the hash grid are the bounding box of the source grid + self._xmin = self._source_grid.lon.min() + self._xmax = self._source_grid.lon.max() + self._ymin = self._source_grid.lat.min() + self._ymax = self._source_grid.lat.max() + # setting min and max below is needed for mesh="flat" + self._zmin = 0.0 + self._zmax = 0.0 + x = self._source_grid.lon + y = self._source_grid.lat + + _xbound = np.stack( + ( + x[:-1, :-1], + x[:-1, 1:], + x[1:, 1:], + x[1:, :-1], + ), + axis=-1, + ) + _ybound = np.stack( + ( + y[:-1, :-1], + y[:-1, 1:], + y[1:, 1:], + y[1:, :-1], + ), + axis=-1, + ) + # Compute bounding box of each face + self._xlow = np.min(_xbound, axis=-1) + self._xhigh = np.max(_xbound, axis=-1) + self._ylow = np.min(_ybound, axis=-1) + self._yhigh = np.max(_ybound, axis=-1) + self._zlow = np.zeros_like(self._xlow) + self._zhigh = np.zeros_like(self._xlow) + + elif isinstance(grid, parcels.uxgrid.UxGrid): + if self._source_grid._mesh == "spherical": + # Boundaries of the hash grid are the unit cube + self._xmin = -1.0 + self._ymin = -1.0 + self._zmin = -1.0 + self._xmax = 1.0 + self._ymax = 1.0 + self._zmax = 1.0 # Compute the cell centers of the source grid (for now, assuming Xgrid) + # Reshape node coordinates to (nfaces, nnodes_per_face) + nids = self._source_grid.uxgrid.face_node_connectivity.values + lon = self._source_grid.uxgrid.node_lon.values[nids] + lat = self._source_grid.uxgrid.node_lat.values[nids] + x, y, z = _latlon_rad_to_xyz(np.deg2rad(lat), np.deg2rad(lon)) + _xbound, _ybound, _zbound = _latlon_rad_to_xyz(np.deg2rad(lat), np.deg2rad(lon)) + + # Compute bounding box of each face + self._xlow = np.atleast_2d(np.min(_xbound, axis=-1)) + self._xhigh = np.atleast_2d(np.max(_xbound, axis=-1)) + self._ylow = np.atleast_2d(np.min(_ybound, axis=-1)) + self._yhigh = np.atleast_2d(np.max(_ybound, axis=-1)) + self._zlow = np.atleast_2d(np.min(_zbound, axis=-1)) + self._zhigh = np.atleast_2d(np.max(_zbound, axis=-1)) + + else: + # Boundaries of the hash grid are the bounding box of the source grid + self._xmin = self._source_grid.uxgrid.node_lon.min().values + self._xmax = self._source_grid.uxgrid.node_lon.max().values + self._ymin = self._source_grid.uxgrid.node_lat.min().values + self._ymax = self._source_grid.uxgrid.node_lat.max().values + # setting min and max below is needed for mesh="flat" + self._zmin = 0.0 + self._zmax = 0.0 + # Reshape node coordinates to (nfaces, nnodes_per_face) + nids = self._source_grid.uxgrid.face_node_connectivity.values + lon = self._source_grid.uxgrid.node_lon.values[nids] + lat = self._source_grid.uxgrid.node_lat.values[nids] + + # Compute bounding box of each face + self._xlow = np.atleast_2d(np.min(lon, axis=-1)) + self._xhigh = np.atleast_2d(np.max(lon, axis=-1)) + self._ylow = np.atleast_2d(np.min(lat, axis=-1)) + self._yhigh = np.atleast_2d(np.max(lat, axis=-1)) + self._zlow = np.zeros_like(self._xlow) + self._zhigh = np.zeros_like(self._xlow) # Generate the mapping from the hash indices to unstructured grid elements self._hash_table = self._initialize_hash_table() @@ -150,17 +199,21 @@ def _initialize_hash_table(self): self._zmax, self._bitwidth, ) - xqlow = xqlow.ravel() - yqlow = yqlow.ravel() - zqlow = zqlow.ravel() - xqhigh = xqhigh.ravel() - yqhigh = yqhigh.ravel() - zqhigh = zqhigh.ravel() - nx = xqhigh - xqlow + 1 - ny = yqhigh - yqlow + 1 - nz = zqhigh - zqlow + 1 - num_hash_per_face = nx * ny * nz - total_hash_entries = np.sum(num_hash_per_face) + xqlow = xqlow.ravel().astype(np.int32, copy=False) + yqlow = yqlow.ravel().astype(np.int32, copy=False) + zqlow = zqlow.ravel().astype(np.int32, copy=False) + xqhigh = xqhigh.ravel().astype(np.int32, copy=False) + yqhigh = yqhigh.ravel().astype(np.int32, copy=False) + zqhigh = zqhigh.ravel().astype(np.int32, copy=False) + nx = (xqhigh - xqlow + 1).astype(np.int32, copy=False) + ny = (yqhigh - yqlow + 1).astype(np.int32, copy=False) + nz = (zqhigh - zqlow + 1).astype(np.int32, copy=False) + num_hash_per_face = (nx * ny * nz).astype( + np.int32, copy=False + ) # Since nx, ny, nz are in the 10-bit range, their product fits in int32 + total_hash_entries = int(num_hash_per_face.sum()) + + # Preallocate output arrays morton_codes = np.zeros(total_hash_entries, dtype=np.uint32) # Compute the j, i indices corresponding to each hash entry @@ -373,17 +426,6 @@ def query(self, y, x): ) -def _latlon_rad_to_xyz(lat, lon): - """Converts Spherical latitude and longitude coordinates into Cartesian x, - y, z coordinates. - """ - x = np.cos(lon) * np.cos(lat) - y = np.sin(lon) * np.cos(lat) - z = np.sin(lat) - - return x, y, z - - def _dilate_bits(n): """ Takes a 10-bit integer n, in range [0,1023], and "dilates" its bits so that diff --git a/parcels/uxgrid.py b/parcels/uxgrid.py index 4702c6f7c..3aa91e177 100644 --- a/parcels/uxgrid.py +++ b/parcels/uxgrid.py @@ -5,8 +5,8 @@ import numpy as np import uxarray as ux +from parcels._index_search import GRID_SEARCH_ERROR, uxgrid_point_in_cell from parcels._typing import assert_valid_mesh -from parcels.tools.statuscodes import FieldOutOfBoundError from parcels.xgrid import _search_1d_array from .basegrid import BaseGrid @@ -42,6 +42,7 @@ def __init__(self, grid: ux.grid.Grid, z: ux.UxDataArray, mesh="flat") -> UxGrid raise ValueError("z must be a 1D array of vertical coordinates") self.z = z self._mesh = mesh + self._spatialhash = None assert_valid_mesh(mesh) @@ -73,138 +74,50 @@ def get_axis_dim(self, axis: _UXGRID_AXES) -> int: return self.uxgrid.n_face def search(self, z, y, x, ei=None, tol=1e-6): - def try_face(fid): - bcoords, err = self._get_barycentric_coordinates_latlon(y, x, fid) - if (bcoords >= 0).all() and (bcoords <= 1).all() and err < tol: - return bcoords - else: - bcoords = self._get_barycentric_coordinates_cartesian(y, x, fid) - if (bcoords >= 0).all() and (bcoords <= 1).all(): - return bcoords + """ + Search for the grid cell (face) and vertical layer that contains the given points. - return None + Parameters + ---------- + z : float or np.ndarray + The vertical coordinate(s) (depth) of the point(s). + y : float or np.ndarray + The latitude(s) of the point(s). + x : float or np.ndarray + The longitude(s) of the point(s). + ei : np.ndarray, optional + Precomputed horizontal indices (face indices) for the points. + + TO BE IMPLEMENTED : If provided, we'll check + if the points are within the faces specified by these indices. For cells where the particles + are not found, a nearest neighbor search will be performed. As a last resort, the spatial hash will be used. + tol : float, optional + Tolerance for barycentric coordinate checks. Default is 1e-6. + """ + x = np.asarray(x, dtype=np.float32) + y = np.asarray(y, dtype=np.float32) + z = np.asarray(z, dtype=np.float32) zi, zeta = _search_1d_array(self.z.values, z) - if ei is not None: + if np.any(ei): indices = self.unravel_index(ei) - fi = indices["FACE"][0] - bcoords = try_face(fi) - if bcoords is not None: - return {"Z": (zi, zeta), "FACE": (np.asarray([fi]), bcoords)} - # Try neighbors of current face - for neighbor in self.uxgrid.face_face_connectivity[fi, :]: - if neighbor == -1: - continue - bcoords = try_face(neighbor) - if bcoords is not None: - return {"Z": (zi, zeta), "FACE": (np.asarray([neighbor]), bcoords)} - - # Global fallback as last ditch effort - points = np.column_stack((x, y)) - face_ids = self.uxgrid.get_faces_containing_point(points, return_counts=False)[0] - fi = face_ids[0] if len(face_ids) > 0 else -1 - if fi == -1: - raise FieldOutOfBoundError(z, y, x) - bcoords = try_face(fi) - if bcoords is None: - raise FieldOutOfBoundError(z, y, x) - return {"Z": (zi, zeta), "FACE": (fi, bcoords)} - - def _get_barycentric_coordinates_latlon(self, y, x, fi): - """Checks if a point is inside a given face id on a UxGrid.""" - # Check if particle is in the same face, otherwise search again. - n_nodes = self.uxgrid.n_nodes_per_face[fi].to_numpy() - node_ids = self.uxgrid.face_node_connectivity[fi, 0:n_nodes] - nodes = np.column_stack( - ( - np.deg2rad(self.uxgrid.node_lon[node_ids].to_numpy()), - np.deg2rad(self.uxgrid.node_lat[node_ids].to_numpy()), - ) - ) - - coord = np.deg2rad(np.column_stack((x, y))) - bcoord = np.asarray(_barycentric_coordinates(nodes, coord)) - proj_coord = np.matmul(np.transpose(nodes), bcoord) - err = np.linalg.norm(proj_coord - coord) - return bcoord, err - - def _get_barycentric_coordinates_cartesian(self, y, x, fi): - n_nodes = self.uxgrid.n_nodes_per_face[fi].to_numpy() - node_ids = self.uxgrid.face_node_connectivity[fi, 0:n_nodes] - - coord = np.deg2rad([x, y]) - x, y, z = _lonlat_rad_to_xyz(coord[0], coord[1]) - cart_coord = np.array([x, y, z]).T - # Second attempt to find barycentric coordinates using cartesian coordinates - nodes = np.stack( - ( - self.uxgrid.node_x[node_ids].values, - self.uxgrid.node_y[node_ids].values, - self.uxgrid.node_z[node_ids].values, - ), - axis=-1, - ) - - bcoord = np.asarray(_barycentric_coordinates(nodes, cart_coord)) - - return bcoord - - -def _lonlat_rad_to_xyz( - lon, - lat, -): - """Converts Spherical latitude and longitude coordinates into Cartesian x, - y, z coordinates. - """ - x = np.cos(lon) * np.cos(lat) - y = np.sin(lon) * np.cos(lat) - z = np.sin(lat) - - return x, y, z - - -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 + fi = indices.get("FACE") + is_in_cell, coords = uxgrid_point_in_cell(self.uxgrid, y, x, fi, fi) + y_check = y[is_in_cell == 0] + x_check = x[is_in_cell == 0] + zero_indices = np.where(is_in_cell == 0)[0] + else: + # Otherwise, we need to check all points + fi = np.full(len(y), GRID_SEARCH_ERROR, dtype=np.int32) + y_check = y + x_check = x + coords = -1.0 * np.ones((len(y), 3), dtype=np.float32) + zero_indices = np.arange(len(y)) + + if len(zero_indices) > 0: + _, face_ids_q, coords_q = self.get_spatial_hash().query(y_check, x_check) + coords[zero_indices, :] = coords_q + fi[zero_indices] = face_ids_q + + return {"Z": (zi, zeta), "FACE": (fi, coords)} diff --git a/parcels/xgrid.py b/parcels/xgrid.py index ebb90f35f..83bbb46f7 100644 --- a/parcels/xgrid.py +++ b/parcels/xgrid.py @@ -10,7 +10,6 @@ from parcels._index_search import _search_indices_curvilinear_2d from parcels._typing import assert_valid_mesh from parcels.basegrid import BaseGrid -from parcels.spatialhash import SpatialHash _XGRID_AXES = Literal["X", "Y", "Z"] _XGRID_AXES_ORDERING: Sequence[_XGRID_AXES] = "ZYX" @@ -354,32 +353,6 @@ def get_axis_dim_mapping(self, dims: list[str]) -> dict[_XGRID_AXES, str]: result[cast(_XGRID_AXES, axis)] = dim return result - 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) - - 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_field.py b/tests/v4/test_field.py index 0a6e10cd5..75c6f6dc9 100644 --- a/tests/v4/test_field.py +++ b/tests/v4/test_field.py @@ -162,16 +162,18 @@ def test_field_unstructured_z_linear(): P = Field(name="p", data=ds.p, grid=grid, interp_method=UXPiecewiseConstantFace) # Test above first cell center - for piecewise constant, should return the depth of the first cell center - assert np.isclose(P.eval(time=ds.time[0].values, z=10.0, y=30.0, x=30.0, applyConversion=False), 55.555557) + assert np.isclose(P.eval(time=ds.time[0].values, z=[10.0], y=[30.0], x=[30.0], applyConversion=False), 55.555557) # Test below first cell center, but in the first layer - for piecewise constant, should return the depth of the first cell center - assert np.isclose(P.eval(time=ds.time[0].values, z=65.0, y=30.0, x=30.0, applyConversion=False), 55.555557) + assert np.isclose(P.eval(time=ds.time[0].values, z=[65.0], y=[30.0], x=[30.0], applyConversion=False), 55.555557) # Test bottom layer - for piecewise constant, should return the depth of the of the bottom layer cell center - assert np.isclose(P.eval(time=ds.time[0].values, z=900.0, y=30.0, x=30.0, applyConversion=False), 944.44445801) + assert np.isclose( + P.eval(time=ds.time[0].values, z=[900.0], y=[30.0], x=[30.0], applyConversion=False), 944.44445801 + ) W = Field(name="W", data=ds.W, grid=grid, interp_method=UXPiecewiseLinearNode) - assert np.isclose(W.eval(time=ds.time[0].values, z=10.0, y=30.0, x=30.0, applyConversion=False), 10.0) - assert np.isclose(W.eval(time=ds.time[0].values, z=65.0, y=30.0, x=30.0, applyConversion=False), 65.0) - assert np.isclose(W.eval(time=ds.time[0].values, z=900.0, y=30.0, x=30.0, applyConversion=False), 900.0) + assert np.isclose(W.eval(time=ds.time[0].values, z=[10.0], y=[30.0], x=[30.0], applyConversion=False), 10.0) + assert np.isclose(W.eval(time=ds.time[0].values, z=[65.0], y=[30.0], x=[30.0], applyConversion=False), 65.0) + assert np.isclose(W.eval(time=ds.time[0].values, z=[900.0], y=[30.0], x=[30.0], applyConversion=False), 900.0) def test_field_constant_in_time(): @@ -183,8 +185,8 @@ def test_field_constant_in_time(): # Assert that the field can be evaluated at any time, and returns the same value time = np.datetime64("2000-01-01T00:00:00") - P1 = P.eval(time=time, z=10.0, y=30.0, x=30.0, applyConversion=False) - P2 = P.eval(time=time + np.timedelta64(1, "D"), z=10.0, y=30.0, x=30.0, applyConversion=False) + P1 = P.eval(time=time, z=[10.0], y=[30.0], x=[30.0], applyConversion=False) + P2 = P.eval(time=time + np.timedelta64(1, "D"), z=[10.0], y=[30.0], x=[30.0], applyConversion=False) assert np.isclose(P1, P2) diff --git a/tests/v4/test_particleset_execute.py b/tests/v4/test_particleset_execute.py index 0bfaca1a4..c5aef6905 100644 --- a/tests/v4/test_particleset_execute.py +++ b/tests/v4/test_particleset_execute.py @@ -400,8 +400,8 @@ def test_uxstommelgyre_pset_execute(): dt=np.timedelta64(60, "s"), pyfunc=AdvectionEE, ) - assert utils.round_and_hash_float_array([p.lon for p in pset]) == 1165397121 - assert utils.round_and_hash_float_array([p.lat for p in pset]) == 1142123780 + assert utils.round_and_hash_float_array([p.lon for p in pset]) == 1165396086 + assert utils.round_and_hash_float_array([p.lat for p in pset]) == 1142124776 @pytest.mark.xfail(reason="Output file not implemented yet") diff --git a/tests/v4/test_uxarray_fieldset.py b/tests/v4/test_uxarray_fieldset.py index d3c8b5de6..0ab13f8ca 100644 --- a/tests/v4/test_uxarray_fieldset.py +++ b/tests/v4/test_uxarray_fieldset.py @@ -119,10 +119,10 @@ def test_fesom2_square_delaunay_uniform_z_coordinate_eval(): P = Field(name="p", data=ds.p, grid=UxGrid(ds.uxgrid, z=ds.coords["nz"]), interp_method=UXPiecewiseLinearNode) fieldset = FieldSet([UVW, P, UVW.U, UVW.V, UVW.W]) - assert fieldset.U.eval(time=ds.time[0].values, z=1.0, y=30.0, x=30.0, applyConversion=False) == 1.0 - assert fieldset.V.eval(time=ds.time[0].values, z=1.0, y=30.0, x=30.0, applyConversion=False) == 1.0 - assert fieldset.W.eval(time=ds.time[0].values, z=1.0, y=30.0, x=30.0, applyConversion=False) == 0.0 - assert fieldset.p.eval(time=ds.time[0].values, z=1.0, y=30.0, x=30.0, applyConversion=False) == 1.0 + assert fieldset.U.eval(time=ds.time[0].values, z=[1.0], y=[30.0], x=[30.0], applyConversion=False) == 1.0 + assert fieldset.V.eval(time=ds.time[0].values, z=[1.0], y=[30.0], x=[30.0], applyConversion=False) == 1.0 + assert fieldset.W.eval(time=ds.time[0].values, z=[1.0], y=[30.0], x=[30.0], applyConversion=False) == 0.0 + assert fieldset.p.eval(time=ds.time[0].values, z=[1.0], y=[30.0], x=[30.0], applyConversion=False) == 1.0 def test_fesom2_square_delaunay_antimeridian_eval(): @@ -132,10 +132,19 @@ def test_fesom2_square_delaunay_antimeridian_eval(): Since the underlying data is constant, we can check that the values are as expected. """ ds = datasets_unstructured["fesom2_square_delaunay_antimeridian"] - P = Field(name="p", data=ds.p, grid=UxGrid(ds.uxgrid, z=ds.coords["nz"]), interp_method=UXPiecewiseLinearNode) + P = Field( + name="p", + data=ds.p, + grid=UxGrid(ds.uxgrid, z=ds.coords["nz"], mesh="spherical"), + interp_method=UXPiecewiseLinearNode, + ) fieldset = FieldSet([P]) - assert np.isclose(fieldset.p.eval(time=ds.time[0].values, z=1.0, y=30.0, x=-170.0, applyConversion=False), 1.0) - assert np.isclose(fieldset.p.eval(time=ds.time[0].values, z=1.0, y=30.0, x=-180.0, applyConversion=False), 1.0) - assert np.isclose(fieldset.p.eval(time=ds.time[0].values, z=1.0, y=30.0, x=180.0, applyConversion=False), 1.0) - assert np.isclose(fieldset.p.eval(time=ds.time[0].values, z=1.0, y=30.0, x=170.0, applyConversion=False), 1.0) + assert np.isclose( + fieldset.p.eval(time=ds.time[0].values, z=[1.0], y=[30.0], x=[-170.0], applyConversion=False), 1.0 + ) + assert np.isclose( + fieldset.p.eval(time=ds.time[0].values, z=[1.0], y=[30.0], x=[-180.0], applyConversion=False), 1.0 + ) + assert np.isclose(fieldset.p.eval(time=ds.time[0].values, z=[1.0], y=[30.0], x=[180.0], applyConversion=False), 1.0) + assert np.isclose(fieldset.p.eval(time=ds.time[0].values, z=[1.0], y=[30.0], x=[170.0], applyConversion=False), 1.0)