Exclude faces that touch NaN nodes from the SpatialHash table

src/parcels/_core/spatialhash.py

@@ -261,7 +261,10 @@ def _total_hash_entries(self, bitwidth):
         nx = xqhigh.astype(np.int64) - xqlow + 1
         ny = yqhigh.astype(np.int64) - yqlow + 1
         nz = zqhigh.astype(np.int64) - zqlow + 1
-        return int((nx * ny * nz).sum())
+        # NaN values are not allowed in the SpatialHash table, so faces with a NaN
+        # bounding box do not contribute to the entry count
+        valid_face = _generate_valid_mask(self._xlow, self._xhigh, self._ylow, self._yhigh, self._zlow, self._zhigh)
+        return int(np.where(valid_face, nx * ny * nz, 0).sum())
 
     def _initialize_hash_table(self):
         """Create a mapping that relates unstructured grid faces to hash indices by determining
@@ -302,7 +305,11 @@ def _initialize_hash_table(self):
         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(
+
+        # prevent NaN values from entering the SpatialHash table by setting their
+        # num_hash_per_face equal to 0
+        valid_face = _generate_valid_mask(self._xlow, self._xhigh, self._ylow, self._yhigh, self._zlow, self._zhigh)
+        num_hash_per_face = np.where(valid_face.ravel(), nx * ny * nz, 0).astype(
             np.int32, copy=False
         )  # Since nx, ny, nz are in the 10-bit range, their product fits in int32
         # Sums over faces can exceed int32, so accumulate in int64
@@ -756,3 +763,30 @@ def _encode_morton3d(x, y, z, xmin, xmax, ymin, ymax, zmin, zmax, bitwidth=1023)
 
     # Since our compact type fits in 30 bits, uint32 is enough.
     return code.astype(np.uint32)
+
+
+def _generate_valid_mask(xlow, xhigh, ylow, yhigh, zlow, zhigh):
+    """
+    Flag faces whose bounding box is fully defined, i.e. none of their 6 bounds
+    is NaN (a NaN indicates a corner node with a missing/masked coordinate).
+
+    Parameters
+    ----------
+    xlow, xhigh : array_like
+        Per-face bounding box in x.
+    ylow, yhigh : array_like
+        Per-face bounding box in y.
+    zlow, zhigh : array_like
+        Per-face bounding box in z.
+
+    Returns
+    -------
+    valid_face : ndarray of bool
+        Same shape as the inputs; True where the face's bounding box is finite,
+        False where it contains a NaN.
+    """
+    invalid_face = (
+        np.isnan(xlow) | np.isnan(xhigh) | np.isnan(ylow) | np.isnan(yhigh) | np.isnan(zlow) | np.isnan(zhigh)
+    )
+
+    return ~invalid_face

src/parcels/_reprs.py

@@ -287,6 +287,7 @@ def spatialhash_describe(spatialhash: SpatialHash) -> str:
     counts = hash_table["counts"]
 
     n_faces = int(np.size(spatialhash._xlow))
+    n_valid_faces = int(np.unique(hash_table["faces"]).size)
     n_entries = int(hash_table["faces"].size)
     n_occupied_cells = int(hash_table["keys"].size)
     n_total_cells = (spatialhash._bitwidth + 1) ** 3
@@ -295,10 +296,11 @@ def spatialhash_describe(spatialhash: SpatialHash) -> str:
         "Grid type": type(grid).__name__,
         "Mesh": grid._mesh,
         "Total mesh faces": f"{n_faces:,d}",
+        "Valid (non-NaN) mesh faces": f"{n_valid_faces:,d}",
         "Bitwidth (current / max)": f"{spatialhash._bitwidth} / 1023  (higher = finer resolution hash grid)",
         "Total hash cells": f"{n_total_cells:,d}",
         "Occupied hash cells": f"{n_occupied_cells:,d}, {n_occupied_cells / n_total_cells * 100:.4f}%",
-        "Total (hash cell --> gird face) entries": f"{n_entries:,d}",
+        "Total (hash cell --> grid face) entries": f"{n_entries:,d}",
         "Entries per occupied hash cell (avg)": f"{n_entries / n_occupied_cells:.2f}" if n_occupied_cells else "-",
         "Entries per face (avg)": f"{n_entries / n_faces:.2f}" if n_faces else "-",
         "Faces per occupied hash cell (min / mean / max)": (

tests/test_spatialhash.py

@@ -33,10 +33,11 @@ def test_spatialhash_describe():
 Grid type                                       : XGrid
 Mesh                                            : FlatMesh()
 Total mesh faces                                : 1,711
+Valid (non-NaN) mesh faces                      : 1,711
 Bitwidth (current / max)                        : 1023 / 1023  (higher = finer resolution hash grid)
 Total hash cells                                : 1,073,741,824
 Occupied hash cells                             : 796,054, 0.0741%
-Total (hash cell --> gird face) entries         : 1,080,194
+Total (hash cell --> grid face) entries         : 1,080,194
 Entries per occupied hash cell (avg)            : 1.36
 Entries per face (avg)                          : 631.32
 Faces per occupied hash cell (min / mean / max) : 1 / 1.36 / 4
@@ -119,3 +120,59 @@ def test_mixed_positions():
     assert i[0] == 14  # Actual value for 2d_left_rotated center
     assert j[1] == -3
     assert i[1] == -3
+
+
+def test_nan_node_invalidates_touching_faces():
+    """Any mesh face that touches a NaN node should not be added to the HashTable."""
+    ds = datasets["2d_left_rotated"]
+    grid = FieldSet.from_sgrid_conventions(ds, mesh="flat").data_g.grid
+    clat, clon, jj, ii = _cell_centers(grid)
+
+    # `grid._ds` shares its lon/lat arrays with the module-level `datasets` fixture
+    # (from_sgrid_conventions does not copy them), so deep-copy before mutating,
+    # otherwise the NaN injected below leaks into every other test in the session.
+    grid._ds = grid._ds.copy(deep=True)
+
+    # Set one interior node to NaN, and calculate the indexes of faces that touch it.
+    nj, ni = 10, 10
+    touching = [
+        (nj - 1, ni - 1),
+        (nj - 1, ni),
+        (nj, ni - 1),
+        (nj, ni),
+    ]
+    grid._ds["lon"].values[nj, ni] = np.nan
+    grid._ds["lat"].values[nj, ni] = np.nan
+    spatialhash = grid.get_spatial_hash(reconstruct=True)
+
+    # From the indexes of the faces that touch the NaN node, calculate their
+    # face_ids.
+    n_faces_x = clon.shape[1]
+    invalid_ids = set()
+    for j, i in touching:
+        invalid_ids.add(j * n_faces_x + i)
+
+    # Get a set of all of the face_ids that are in the table, and assert that the
+    # ones touching the NaN node are not among them.
+    faces_in_table = set(np.unique(spatialhash._hash_table["faces"]).tolist())
+    assert invalid_ids.isdisjoint(faces_in_table)
+
+    # The total number of mesh faces should be greater than the number in the table,
+    # since the NaN faces are filtered from the table.
+    n_total_faces = jj.size
+    assert n_total_faces > len(faces_in_table)
+
+    # Queries landing on those 4 faces should return GridSearchErrors (-3).
+    touching_lat = np.array([clat[j, i] for j, i in touching])
+    touching_lon = np.array([clon[j, i] for j, i in touching])
+    j_touch, i_touch, _ = spatialhash.query(touching_lat, touching_lon)
+    assert np.all(j_touch == -3)
+    assert np.all(i_touch == -3)
+
+    # All mesh cells not contacting the NaN node should resolve queries.
+    mask = np.ones(clat.shape, dtype=bool)
+    for j, i in touching:
+        mask[j, i] = False
+    j_rest, i_rest, _ = spatialhash.query(clat[mask], clon[mask])
+    assert np.array_equal(j_rest, jj[mask])
+    assert np.array_equal(i_rest, ii[mask])