Spatialhash regional spherical improvements (#2744)

* Tighten bounding cube for regional spherical domains

When the mesh type is spherical, we now use the actual x,y,z bounding
box rather than the unit cube. For regional runs that leverage the
spherical barycentric coordinates and built in unit conversion in the
velocity field interpolators, this reduces the hit count per hash cell
and can improve the search time in the spatialhash.query()

Note that quantize_coordinates now clips in float space before casting to uint32.
This matters now that spherical queries can fall outside the hash-grid bounds
(previously impossible with the unit cube). A negative normalized value cast to
uint32 wraps to a huge number, which would have mapped below-range queries to the
top bin (with morton code 2^32) erroneously.

For spherical grids, points outside the regional box now quantize to edge bins
rather than interior bins. In this case, they either fail the exact Morton-key
match or are rejected by the point-in-cell check; either case results in the
same outcome as before (GRID_SEARCH_ERROR)

* Add structured grid spherical spatialhash query tests

Tests for both points inside the domain and explicitly outside the domain
to confirm correct behavior for out of bounds PIC checks

* Update documentation on the unstructured grid search to note change to spherical bounding box

* Remove dead valid-face branch and identity scatter from hash table build

num_hash_per_face is never zero (quantization is monotone), so the valid
mask was always all-True and the scatter an identity permutation.
Removing them drops three entry-length temporaries (~35% peak build
memory). Accumulate entry counts in int64 to avoid overflow.
Hash table output verified byte-identical.

Co-authored-by: Claude <noreply@anthropic.com>

* Sort packed (morton code, face id) uint64 pairs in place in hash table build

Fusing each pair as (code << 32) | face_id and sorting in place replaces
argsort, its int64 permutation array, and two gather passes. Pairs are
unique so the ordering is deterministic (ties by ascending face id).
Also free decode temporaries before the sort. Peak build memory drops
17.2 GB -> 6.35 GB (with previous commit) and build time 40 s -> 32 s on
a 101M-entry spherical UxGrid table. CSR arrays verified byte-identical;
per-cell face sets verified identical.

Co-authored-by: Claude <noreply@anthropic.com>

* Replace np.unique function with a simpler set of calls that exploits sorting

The np.unique function makes no assumption about whether or not the input
is sorted. Since we've done the work to sort  the list of morton codes and
face id pairings, computing the start, count, and unique keys of the hash table
can be done with just three operations:

1. Find the indices where the morton codes in the sorted list change value -
   prepending this list of indices with [0] gives the starting indices for a
   CSR format matrix

2. The keys for the hash table are just the unique morton code values. Once
   we know the `start` indices, we can quickly index for the unique morton codes

3. The count (the number of faces per morton code / hash table key) is just the
   np.diff of the `start` indices.

This results in a minor improvement in construction speed.

* Store the face id rather than i,j in the hash table to reduce memory footprint

Instead, we opt to unravel the face during the query. This cuts the
hash table size down by only storing a single uint64 array (with the start
and count) rather than two arrays. The unravel cost in the query is
small in comparison to particle in cell checks, which dominate the query
costs.

* Enforce a budget on the number of hash cells per face

Enforcing this upper bound on the number of hash cells per face helps considerably
in reducing the memory footprint of the hash table construction; this also significantly
speeds up the construction process. The max budget value was tuned to balance keeping the
number of PIC checks low while keeping the memory peak during construction small.

Claude code was used to design tests for this new feature

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Remove erroneous test that relied on bitwidth parameter for hash table construction

I've removed the bitwidth parameter for construction and opted for this to be calculated
internally for optimal memory footprint and query cost

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Erik van Sebille <e.vansebille@uu.nl>

Author: 4 people

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

@@ -11,6 +11,17 @@
 from parcels._core.warnings import FieldSetWarning
 from parcels._python import isinstance_noimport
 
+# Budget on the total number of (face, hash cell) pairs in the hash table:
+# max(_HASH_ENTRIES_PER_FACE * nfaces, _HASH_ENTRY_BUDGET_MIN).
+# When the hash grid cell size, set by the bitwidth, would result in too many
+# hash entries per face, the hash grid is coarsened by lowering the bitwidth
+# The target value for the number of hash entries per face is chosen to keep
+# the number of particle in cell checks as small as possible, while also
+# minimizing the hash table construction memory footprint.
+_HASH_ENTRIES_PER_FACE = 16
+_HASH_ENTRY_BUDGET_MIN = 2**22
+_HASH_MAX_BITWIDTH = 1023
+
 
 class SpatialHash:
     """Custom data structure that is used for performing grid searches using Spatial Hashing. This class constructs an overlying
@@ -31,7 +42,6 @@ class SpatialHash:
     def __init__(
         self,
         grid,
-        bitwidth=1023,
     ):
         if isinstance_noimport(grid, "XGrid"):
             self._point_in_cell = curvilinear_point_in_cell
@@ -41,21 +51,23 @@ def __init__(
             raise ValueError("Expected `grid` to be a parcels.XGrid or parcels.UxGrid")
 
         self._source_grid = grid
-        self._bitwidth = bitwidth  # Max integer to use per coordinate in quantization (10 bits = 0..1023)
+        self._bitwidth = _HASH_MAX_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
-                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 +161,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))
@@ -193,9 +206,60 @@ def __init__(
                 self._zlow = np.zeros_like(self._xlow)
                 self._zhigh = np.zeros_like(self._xlow)
 
+        # Cap the quantization resolution so the hash table stays within a fixed entry
+        # budget.
+        budget = max(_HASH_ENTRIES_PER_FACE * np.size(self._xlow), _HASH_ENTRY_BUDGET_MIN)
+        if self._total_hash_entries(self._bitwidth) > budget:
+            # Binary search for the largest bitwidth whose table fits the budget. The
+            # entry count is not perfectly monotone in bitwidth (cell-boundary flooring
+            # effects), so the result may sit marginally below the true maximum; any
+            # in-budget bitwidth is valid. At bitwidth 1 the count equals nfaces, which
+            # is always within budget, so the search cannot fail.
+            lo, hi = 1, self._bitwidth
+            while lo < hi:
+                mid = (lo + hi + 1) // 2
+                if self._total_hash_entries(mid) <= budget:
+                    lo = mid
+                else:
+                    hi = mid - 1
+            self._bitwidth = lo
+
         # Generate the mapping from the hash indices to unstructured grid elements
         self._hash_table = self._initialize_hash_table()
 
+    def _total_hash_entries(self, bitwidth):
+        """Total number of (face, hash cell) pairs the hash table would hold at a given
+        quantization resolution, i.e. the summed hash-cell count of all face bounding boxes.
+        """
+        xqlow, yqlow, zqlow = quantize_coordinates(
+            self._xlow,
+            self._ylow,
+            self._zlow,
+            self._xmin,
+            self._xmax,
+            self._ymin,
+            self._ymax,
+            self._zmin,
+            self._zmax,
+            bitwidth,
+        )
+        xqhigh, yqhigh, zqhigh = quantize_coordinates(
+            self._xhigh,
+            self._yhigh,
+            self._zhigh,
+            self._xmin,
+            self._xmax,
+            self._ymin,
+            self._ymax,
+            self._zmin,
+            self._zmax,
+            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())
+
     def _initialize_hash_table(self):
         """Create a mapping that relates unstructured grid faces to hash indices by determining
         which faces overlap with which hash cells
@@ -238,82 +302,77 @@ 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]
-
-        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]
-        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)
-        keys, starts, counts = np.unique(morton_codes_sorted, return_index=True, return_counts=True)
+        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
+        )
 
+        # 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
+        # Perform a single sort on the packed (morton_code | face_id ) list
+        packed.sort()
+        # Trunctating back to a uint32 keeps the lower 32 bits (the face_id's)
+        face_sorted = packed.astype(np.uint32)
+        # Purge the face ids from the packed list to retain only the morton codes
+        packed >>= np.uint64(32)
+        # Cast the morton codes back to uint32
+        morton_codes_sorted = packed.astype(np.uint32)
+        del packed
+
+        # Get a list of unique morton codes and their corresponding starts and counts (CSR format).
+        # The codes are already sorted at this point, first by morton code, then by face_id
+        # Starting indices of the matrix rows are located by finding indices where the morton codes differ
+        starts = np.concatenate(([0], np.flatnonzero(morton_codes_sorted[1:] != morton_codes_sorted[:-1]) + 1))
+        # The unique keys for the hash table are the unique morton codes
+        keys = morton_codes_sorted[starts]
+        # The number of faces per hash keys (morton codes) is easily calculated as the difference betwee the start values
+        counts = np.diff(np.concatenate((starts, [morton_codes_sorted.size])))
+
+        # The flat face id is stored (4 bytes per entry); query() unravels the gathered
+        # candidates to (j, i) on demand, instead of holding two precomputed int64
+        # index arrays (16 bytes per entry) for the lifetime of the grid.
         hash_table = {
             "keys": keys,
             "starts": starts,
             "counts": counts,
-            "i": i_sorted,
-            "j": j_sorted,
+            "faces": face_sorted,
         }
         return hash_table
 
@@ -341,8 +400,7 @@ def query(self, y, x):
         keys = self._hash_table["keys"]
         starts = self._hash_table["starts"]
         counts = self._hash_table["counts"]
-        i = self._hash_table["i"]
-        j = self._hash_table["j"]
+        faces = self._hash_table["faces"]
 
         y = np.asarray(y)
         x = np.asarray(x)
@@ -358,7 +416,16 @@ def query(self, y, x):
             qz = np.zeros_like(qx)
 
         query_codes = _encode_morton3d(
-            qx, qy, qz, self._xmin, self._xmax, self._ymin, self._ymax, self._zmin, self._zmax
+            qx,
+            qy,
+            qz,
+            self._xmin,
+            self._xmax,
+            self._ymin,
+            self._ymax,
+            self._zmin,
+            self._zmax,
+            bitwidth=self._bitwidth,
         ).ravel()
         num_queries = query_codes.size
 
@@ -418,9 +485,10 @@ def query(self, y, x):
         # use to quickly gather the (i,j) pairs for each query
         source_idx = starts[hash_positions].astype(np.int32) + intra
 
-        # Gather all candidate (j,i) pairs in one shot
-        j_all = j[source_idx]
-        i_all = i[source_idx]
+        # Gather all candidate face ids in one shot and unravel them to (j, i) pairs;
+        # only the gathered candidates are unraveled, not the whole table
+        face_all = faces[source_idx]
+        j_all, i_all = np.unravel_index(face_all, self._xlow.shape)
 
         # Now we need to construct arrays that repeats the y and x coordinates for each candidate
         # to enable vectorized point-in-cell checks
@@ -588,10 +656,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