Skip to content

Commit 35dd7d9

Browse files
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>
1 parent 34af903 commit 35dd7d9

1 file changed

Lines changed: 18 additions & 6 deletions

File tree

src/parcels/_core/spatialhash.py

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -251,7 +251,7 @@ def _initialize_hash_table(self):
251251
# box. Entries are generated in face-major order: face_ids maps each entry to its
252252
# face, and intra enumerates the cells of that face's box (0..num_hash_per_face-1).
253253
nface = np.size(self._xlow)
254-
face_ids = np.repeat(np.arange(nface, dtype=np.int32), num_hash_per_face)
254+
face_ids = np.repeat(np.arange(nface, dtype=np.uint32), num_hash_per_face)
255255
face_starts = np.concatenate(([0], np.cumsum(num_hash_per_face, dtype=np.int64)))[:-1]
256256
intra = np.arange(total_hash_entries, dtype=idx_dtype) - np.repeat(
257257
face_starts.astype(idx_dtype, copy=False), num_hash_per_face
@@ -273,11 +273,23 @@ def _initialize_hash_table(self):
273273

274274
# Vectorized morton encode for all entries at once, already in face-major order
275275
morton_codes = _encode_quantized_morton3d(xq, yq, zq)
276-
277-
# Sort face indices by morton code
278-
order = np.argsort(morton_codes)
279-
morton_codes_sorted = morton_codes[order]
280-
face_sorted = face_ids[order]
276+
del intra, rem, xi, yi, zi, ny_nz, nz_rep, xq, yq, zq
277+
278+
# Sort entries by morton code. Each (code, face) pair is fused into one uint64
279+
# with the code in the high 32 bits and the face id in the low 32 bits: unsigned
280+
# comparison then orders by code, with ties broken by ascending face id. Sorting
281+
# the fused array in place avoids the argsort permutation array and the gather
282+
# copies it would imply. Pairs are unique, so the ordering is deterministic.
283+
packed = morton_codes.astype(np.uint64)
284+
del morton_codes
285+
packed <<= np.uint64(32)
286+
np.bitwise_or(packed, face_ids, out=packed)
287+
del face_ids
288+
packed.sort()
289+
face_sorted = packed.astype(np.uint32) # truncating cast keeps the low 32 bits
290+
packed >>= np.uint64(32)
291+
morton_codes_sorted = packed.astype(np.uint32)
292+
del packed
281293
j_sorted, i_sorted = np.unravel_index(face_sorted, self._xlow.shape)
282294

283295
# Get a list of unique morton codes and their corresponding starts and counts (CSR format)

0 commit comments

Comments
 (0)