Merge pull request #2249 from OceanParcels/bugfix/2239-oob-spatialhash

Bugfix/2239 oob spatialhash

Author: fluidnumerics-joe

parcels/_core/index_search.py

@@ -224,7 +224,7 @@ def uxgrid_point_in_cell(grid, y: np.ndarray, x: np.ndarray, yi: np.ndarray, xi:
             ),
             axis=-1,
         )
-        points = np.stack((x, y))
+        points = np.stack((x, y), axis=-1)
 
     M = len(points)
 
@@ -252,9 +252,6 @@ def _triangle_area(A, B, C):
         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):

parcels/_core/spatialhash.py

@@ -41,6 +41,7 @@ def __init__(
         self._bitwidth = bitwidth  # Max integer to use per coordinate in quantization (10 bits = 0..1023)
 
         if isinstance_noimport(grid, "XGrid"):
+            self._coord_dim = 2  # Number of computational coordinates is 2 (bilinear interpolation)
             if self._source_grid._mesh == "spherical":
                 # Boundaries of the hash grid are the unit cube
                 self._xmin = -1.0
@@ -126,6 +127,7 @@ def __init__(
                 self._zhigh = np.zeros_like(self._xlow)
 
         elif isinstance_noimport(grid, "UxGrid"):
+            self._coord_dim = grid.uxgrid.n_max_face_nodes  # Number of barycentric coordinates
             if self._source_grid._mesh == "spherical":
                 # Boundaries of the hash grid are the unit cube
                 self._xmin = -1.0
@@ -344,7 +346,11 @@ def query(self, y, x):
         pos = np.searchsorted(keys, query_codes)  # pos is shape (num_queries,)
 
         # Valid hits: inside range with finite query coordinates and query codes give exact morton code match.
-        valid = (pos < len(keys)) & np.isfinite(x) & np.isfinite(y) & (query_codes == keys[pos])
+        valid = (pos < len(keys)) & np.isfinite(x) & np.isfinite(y)
+        # Clip pos to valid range to avoid out-of-bounds indexing
+        pos = np.clip(pos, 0, len(keys) - 1)
+        # Further filter out false positives from searchsorted by checking for exact code match
+        valid[valid] &= query_codes[valid] == keys[pos[valid]]
 
         # Pre-allocate i and j indices of the best match for each query
         # Default values to -1 (no match case)
@@ -357,7 +363,7 @@ def query(self, y, x):
             return (
                 j_best.reshape(query_codes.shape),
                 i_best.reshape(query_codes.shape),
-                np.full((num_queries, 2), -1.0, dtype=np.float32),
+                np.full((num_queries, self._coord_dim), -1.0, dtype=np.float32),
             )
 
         # Now, for each query, we need to gather the candidate (j,i) indices from the hash table