Skip to content

Commit badcf7d

Browse files
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.
1 parent 35dd7d9 commit badcf7d

1 file changed

Lines changed: 13 additions & 3 deletions

File tree

src/parcels/_core/spatialhash.py

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -285,15 +285,25 @@ def _initialize_hash_table(self):
285285
packed <<= np.uint64(32)
286286
np.bitwise_or(packed, face_ids, out=packed)
287287
del face_ids
288+
# Perform a single sort on the packed (morton_code | face_id ) list
288289
packed.sort()
289-
face_sorted = packed.astype(np.uint32) # truncating cast keeps the low 32 bits
290+
# Trunctating back to a uint32 keeps the lower 32 bits (the face_id's)
291+
face_sorted = packed.astype(np.uint32)
292+
# Purge the face ids from the packed list to retain only the morton codes
290293
packed >>= np.uint64(32)
294+
# Cast the morton codes back to uint32
291295
morton_codes_sorted = packed.astype(np.uint32)
292296
del packed
293297
j_sorted, i_sorted = np.unravel_index(face_sorted, self._xlow.shape)
294298

295-
# Get a list of unique morton codes and their corresponding starts and counts (CSR format)
296-
keys, starts, counts = np.unique(morton_codes_sorted, return_index=True, return_counts=True)
299+
# Get a list of unique morton codes and their corresponding starts and counts (CSR format).
300+
# The codes are already sorted at this point, first by morton code, then by face_id
301+
# Starting indices of the matrix rows are located by finding indices where the morton codes differ
302+
starts = np.concatenate(([0], np.flatnonzero(morton_codes_sorted[1:] != morton_codes_sorted[:-1]) + 1))
303+
# The unique keys for the hash table are the unique morton codes
304+
keys = morton_codes_sorted[starts]
305+
# The number of faces per hash keys (morton codes) is easily calculated as the difference betwee the start values
306+
counts = np.diff(np.concatenate((starts, [morton_codes_sorted.size])))
297307

298308
hash_table = {
299309
"keys": keys,

0 commit comments

Comments
 (0)