Spatialhash regional spherical improvements

docs/development/unstructured_grid_search.md

@@ -42,12 +42,12 @@ The algorithm has two phases: **initialisation** (once, when the grid is first u
 
 The hash grid is three-dimensional regardless of the source grid:
 
-| Source grid | Mesh type | Hash-grid space                       |
-| ----------- | --------- | ------------------------------------- |
-| `XGrid`     | spherical | Cartesian unit cube (lon/lat → x,y,z) |
-| `XGrid`     | flat      | 2-D lon/lat bounding box (z set to 0) |
-| `UxGrid`    | spherical | Cartesian unit cube                   |
-| `UxGrid`    | flat      | 2-D lon/lat bounding box              |
+| Source grid | Mesh type | Hash-grid space                          |
+| ----------- | --------- | ---------------------------------------- |
+| `XGrid`     | spherical | Cartesian bounding box (lon/lat → x,y,z) |
+| `XGrid`     | flat      | 2-D lon/lat bounding box (z set to 0)    |
+| `UxGrid`    | spherical | Cartesian bounding box (lon/lat → x,y,z) |
+| `UxGrid`    | flat      | 2-D lon/lat bounding box                 |
 
 For spherical grids, node coordinates are converted from degrees to radians and then to Cartesian (x, y, z) on the unit sphere using `parcels._core.index_search._latlon_rad_to_xyz`:
 
@@ -57,7 +57,7 @@ y = sin(lon) * cos(lat)
 z = sin(lat)
 ```
 
-The hash grid then spans the unit cube `[-1, 1]³`. Working in Cartesian space avoids the longitude wrap-around discontinuity that would otherwise cause the bounding boxes of cells crossing the antimeridian to erroneously span the entire domain.
+The hash grid then spans the axis-aligned Cartesian bounding box of the transformed grid nodes. For a global grid this is (nearly) the unit cube `[-1, 1]³`; for a regional grid it is much tighter, so the full quantization resolution (1024 bins per axis) is spent on the region actually covered by the grid rather than the whole sphere. Working in Cartesian space avoids the longitude wrap-around discontinuity that would otherwise cause the bounding boxes of cells crossing the antimeridian to erroneously span the entire domain.
 
 For flat meshes the hash grid simply spans `[lon_min, lon_max] × [lat_min, lat_max]`, with z fixed at 0.
 

src/parcels/_core/spatialhash.py

@@ -46,16 +46,18 @@ def __init__(
         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
-                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)
+                # Boundaries of the hash grid are the Cartesian bounding box of the
+                # transformed grid, so that regional domains retain full quantization
+                # resolution instead of spreading it over the whole unit cube
+                self._xmin = x.min()
+                self._xmax = x.max()
+                self._ymin = y.min()
+                self._ymax = y.max()
+                self._zmin = z.min()
+                self._zmax = z.max()
                 _xbound = np.stack(
                     (
                         x[:-1, :-1],
@@ -149,19 +151,20 @@ def __init__(
         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
-                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))
+                # Boundaries of the hash grid are the Cartesian bounding box of the
+                # transformed grid, so that regional domains retain full quantization
+                # resolution instead of spreading it over the whole unit cube
+                self._xmin = _xbound.min()
+                self._xmax = _xbound.max()
+                self._ymin = _ybound.min()
+                self._ymax = _ybound.max()
+                self._zmin = _zbound.min()
+                self._zmax = _zbound.max()
 
                 # Compute bounding box of each face
                 self._xlow = np.atleast_2d(np.min(_xbound, axis=-1))
@@ -238,71 +241,55 @@ def _initialize_hash_table(self):
         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
+        # Sums over faces can exceed int32, so accumulate in int64
+        total_hash_entries = int(num_hash_per_face.sum(dtype=np.int64))
+        # Entry indices fit in int32 for all but extreme cases; fall back to int64 when needed
+        idx_dtype = np.int64 if total_hash_entries > np.iinfo(np.int32).max else np.int32
+
+        # Every face overlaps at least one hash cell (nx, ny, nz >= 1 since quantization
+        # is monotone), and contributes one hash entry per cell of its quantized bounding
+        # box. Entries are generated in face-major order: face_ids maps each entry to its
+        # face, and intra enumerates the cells of that face's box (0..num_hash_per_face-1).
         nface = np.size(self._xlow)
-        face_ids = np.repeat(np.arange(nface, dtype=np.int32), num_hash_per_face)
-        offsets = np.concatenate(([0], np.cumsum(num_hash_per_face))).astype(np.int32)[:-1]
+        face_ids = np.repeat(np.arange(nface, dtype=np.uint32), num_hash_per_face)
+        face_starts = np.concatenate(([0], np.cumsum(num_hash_per_face, dtype=np.int64)))[:-1]
+        intra = np.arange(total_hash_entries, dtype=idx_dtype) - np.repeat(
+            face_starts.astype(idx_dtype, copy=False), num_hash_per_face
+        )
 
-        valid = num_hash_per_face != 0
-        if not np.any(valid):
-            # nothing to do
-            pass
-        else:
-            # Grab only valid faces to avoid empty arrays
-            nx_v = np.asarray(nx[valid], dtype=np.int32)
-            ny_v = np.asarray(ny[valid], dtype=np.int32)
-            nz_v = np.asarray(nz[valid], dtype=np.int32)
-            xlow_v = np.asarray(xqlow[valid], dtype=np.int32)
-            ylow_v = np.asarray(yqlow[valid], dtype=np.int32)
-            zlow_v = np.asarray(zqlow[valid], dtype=np.int32)
-            starts_v = np.asarray(offsets[valid], dtype=np.int32)
-
-            # Count of elements per valid face (should match num_hash_per_face[valid])
-            counts = (nx_v * ny_v * nz_v).astype(np.int32)
-            total = int(counts.sum())
-
-            # Map each global element to its face and output position
-            start_for_elem = np.repeat(starts_v, counts)  # shape (total,)
-
-            # Intra-face linear index for each element (0..counts_i-1)
-            # Offsets per face within the concatenation of valid faces:
-            face_starts_local = np.cumsum(np.r_[0, counts[:-1]])
-            intra = np.arange(total, dtype=np.int32) - np.repeat(face_starts_local, counts)
-
-            # Derive (zi, yi, xi) from intra using per-face sizes
-            ny_nz = np.repeat(ny_v * nz_v, counts)
-            nz_rep = np.repeat(nz_v, counts)
-
-            xi = intra // ny_nz
-            rem = intra % ny_nz
-            yi = rem // nz_rep
-            zi = rem % nz_rep
-
-            # Add per-face lows
-            x0 = np.repeat(xlow_v, counts)
-            y0 = np.repeat(ylow_v, counts)
-            z0 = np.repeat(zlow_v, counts)
-
-            xq = x0 + xi
-            yq = y0 + yi
-            zq = z0 + zi
-
-            # Vectorized morton encode for all elements at once
-            codes_all = _encode_quantized_morton3d(xq, yq, zq)
-
-            # Scatter into the preallocated output using computed absolute indices
-            out_idx = start_for_elem + intra
-            morton_codes[out_idx] = codes_all
-
-        # Sort face indices by morton code
-        order = np.argsort(morton_codes)
-        morton_codes_sorted = morton_codes[order]
-        face_sorted = face_ids[order]
+        # Derive (xi, yi, zi) cell offsets within each face's box from intra,
+        # then shift by the per-face low corner to get quantized cell coordinates
+        ny_nz = np.repeat(ny * nz, num_hash_per_face)
+        nz_rep = np.repeat(nz, num_hash_per_face)
+
+        xi = intra // ny_nz
+        rem = intra % ny_nz
+        yi = rem // nz_rep
+        zi = rem % nz_rep
+
+        xq = np.repeat(xqlow, num_hash_per_face) + xi
+        yq = np.repeat(yqlow, num_hash_per_face) + yi
+        zq = np.repeat(zqlow, num_hash_per_face) + zi
+
+        # Vectorized morton encode for all entries at once, already in face-major order
+        morton_codes = _encode_quantized_morton3d(xq, yq, zq)
+        del intra, rem, xi, yi, zi, ny_nz, nz_rep, xq, yq, zq
+
+        # Sort entries by morton code. Each (code, face) pair is fused into one uint64
+        # with the code in the high 32 bits and the face id in the low 32 bits: unsigned
+        # comparison then orders by code, with ties broken by ascending face id. Sorting
+        # the fused array in place avoids the argsort permutation array and the gather
+        # copies it would imply. Pairs are unique, so the ordering is deterministic.
+        packed = morton_codes.astype(np.uint64)
+        del morton_codes
+        packed <<= np.uint64(32)
+        np.bitwise_or(packed, face_ids, out=packed)
+        del face_ids
+        packed.sort()
+        face_sorted = packed.astype(np.uint32)  # truncating cast keeps the low 32 bits
+        packed >>= np.uint64(32)
+        morton_codes_sorted = packed.astype(np.uint32)
+        del packed
         j_sorted, i_sorted = np.unravel_index(face_sorted, self._xlow.shape)
 
         # Get a list of unique morton codes and their corresponding starts and counts (CSR format)
@@ -588,10 +575,15 @@ def quantize_coordinates(x, y, z, xmin, xmax, ymin, ymax, zmin, zmax, bitwidth=1
         zn = np.where(dz != 0, (z - zmin) / dz, 0.0)
 
     # --- 2) Quantize to (0..bitwidth). ---
-    # Multiply by bitwidth, round down, and clip to be safe against slight overshoot.
-    xq = np.clip((xn * bitwidth).astype(np.uint32), 0, bitwidth)
-    yq = np.clip((yn * bitwidth).astype(np.uint32), 0, bitwidth)
-    zq = np.clip((zn * bitwidth).astype(np.uint32), 0, bitwidth)
+    # Multiply by bitwidth, round down, and clip to be safe against overshoot.
+    # Clip in float space before casting: out-of-range queries (e.g., points outside
+    # a regional domain) would otherwise wrap around when a negative float is cast to uint32.
+    # NaN queries produce arbitrary codes here; they are discarded downstream by the
+    # finite-coordinate mask in SpatialHash.query.
+    with np.errstate(invalid="ignore"):
+        xq = np.clip(xn * bitwidth, 0, bitwidth).astype(np.uint32)
+        yq = np.clip(yn * bitwidth, 0, bitwidth).astype(np.uint32)
+        zq = np.clip(zn * bitwidth, 0, bitwidth).astype(np.uint32)
 
     return xq, yq, zq
 

tests/test_spatialhash.py

@@ -20,6 +20,40 @@ def test_invalid_positions():
     assert np.all(i == -3)
 
 
+def test_spherical_regional_bounds():
+    """Hash-grid bounds for spherical meshes are the Cartesian bounding box of the
+    (regional) grid, not the whole unit cube, so quantization resolution is not
+    wasted on parts of the sphere the grid does not cover.
+    """
+    ds = datasets["2d_left_rotated"]
+    grid = FieldSet.from_sgrid_conventions(ds, mesh="spherical").data_g.grid
+    spatialhash = grid.get_spatial_hash()
+
+    extents = np.array(
+        [
+            spatialhash._xmax - spatialhash._xmin,
+            spatialhash._ymax - spatialhash._ymin,
+            spatialhash._zmax - spatialhash._zmin,
+        ]
+    )
+    assert np.all(extents > 0.0)
+    assert np.all(extents < 2.0)  # strictly tighter than the unit cube
+
+    # Queries at cell centers must still resolve to the correct cell
+    lon, lat = grid.lon, grid.lat
+    clon = 0.25 * (lon[:-1, :-1] + lon[:-1, 1:] + lon[1:, :-1] + lon[1:, 1:])
+    clat = 0.25 * (lat[:-1, :-1] + lat[:-1, 1:] + lat[1:, :-1] + lat[1:, 1:])
+    jj, ii = np.meshgrid(np.arange(clat.shape[0]), np.arange(clat.shape[1]), indexing="ij")
+    j, i, _ = spatialhash.query(clat.ravel(), clon.ravel())
+    assert np.array_equal(j, jj.ravel())
+    assert np.array_equal(i, ii.ravel())
+
+    # Points far outside the regional domain must not match any cell
+    j, i, _ = spatialhash.query([-60.0, 80.0], [120.0, -150.0])
+    assert np.all(j == -3)
+    assert np.all(i == -3)
+
+
 def test_mixed_positions():
     ds = datasets["2d_left_rotated"]
     grid = FieldSet.from_sgrid_conventions(ds, mesh="flat").data_g.grid