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

Author: fluidnumericsJoe

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,7 +51,7 @@ 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)
@@ -196,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
@@ -355,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
 

tests/test_spatialhash.py

@@ -1,9 +1,18 @@
 import numpy as np
 
 from parcels._core.fieldset import FieldSet
+from parcels._core.spatialhash import _HASH_ENTRIES_PER_FACE, _HASH_ENTRY_BUDGET_MIN, SpatialHash
 from parcels._datasets.structured.generic import datasets
 
 
+def _cell_centers(grid):
+    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")
+    return clat, clon, jj, ii
+
+
 def test_spatialhash_init():
     ds = datasets["2d_left_rotated"]
     grid = FieldSet.from_sgrid_conventions(ds, mesh="flat").data_g.grid
@@ -40,10 +49,7 @@ def test_spherical_regional_bounds():
     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")
+    clat, clon, jj, ii = _cell_centers(grid)
     j, i, _ = spatialhash.query(clat.ravel(), clon.ravel())
     assert np.array_equal(j, jj.ravel())
     assert np.array_equal(i, ii.ravel())
@@ -54,6 +60,41 @@ def test_spherical_regional_bounds():
     assert np.all(i == -3)
 
 
+def test_query_honors_reduced_bitwidth():
+    """Queries must be encoded at the same quantization resolution the table was
+    built with; a mismatch makes every Morton key miss and all queries fail.
+    """
+    ds = datasets["2d_left_rotated"]
+    grid = FieldSet.from_sgrid_conventions(ds, mesh="spherical").data_g.grid
+    spatialhash = SpatialHash(grid, bitwidth=64)
+
+    clat, clon, jj, ii = _cell_centers(grid)
+    j, i, _ = spatialhash.query(clat.ravel(), clon.ravel())
+    assert np.array_equal(j, jj.ravel())
+    assert np.array_equal(i, ii.ravel())
+
+
+def test_hash_entry_budget():
+    """When the requested bitwidth would blow the hash-entry budget (e.g. tilted
+    regional spherical meshes, where face bounding boxes overlap 3-D blocks of
+    hash cells), the resolution is reduced to fit; queries still resolve exactly.
+    """
+    ds = datasets["2d_left_rotated"]
+    grid = FieldSet.from_sgrid_conventions(ds, mesh="spherical").data_g.grid
+    spatialhash = grid.get_spatial_hash()
+
+    budget = max(_HASH_ENTRIES_PER_FACE * np.size(spatialhash._xlow), _HASH_ENTRY_BUDGET_MIN)
+    assert spatialhash._total_hash_entries(1023) > budget  # this grid requires the cap
+    assert spatialhash._bitwidth < 1023
+    assert spatialhash._total_hash_entries(spatialhash._bitwidth) <= budget
+    assert spatialhash._hash_table["faces"].size <= budget
+
+    clat, clon, jj, ii = _cell_centers(grid)
+    j, i, _ = spatialhash.query(clat.ravel(), clon.ravel())
+    assert np.array_equal(j, jj.ravel())
+    assert np.array_equal(i, ii.ravel())
+
+
 def test_mixed_positions():
     ds = datasets["2d_left_rotated"]
     grid = FieldSet.from_sgrid_conventions(ds, mesh="flat").data_g.grid