diff --git a/src/parcels/_core/spatialhash.py b/src/parcels/_core/spatialhash.py index b3e6aa677..15e4e47d0 100644 --- a/src/parcels/_core/spatialhash.py +++ b/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 diff --git a/src/parcels/_reprs.py b/src/parcels/_reprs.py index 03323b2a7..9377cbc41 100644 --- a/src/parcels/_reprs.py +++ b/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)": ( diff --git a/tests/test_spatialhash.py b/tests/test_spatialhash.py index 78ae1311b..2f6222f97 100644 --- a/tests/test_spatialhash.py +++ b/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])