Replace curvilinear grid search with spatial hashing

[fluidnumerics-joe]

The initial search for particles in a curvilinear grid starts in the middle of the model domain and searches outward, based on the computed barycentric coordinates, for the element that a particle resides in. We currently believe that a more sophisticated approach, like spatial hashing could yield faster initialization of the particles.

To resolve this issue, put together a benchmark to estimate the time taken for initial particle search on a curvilinear grid. Implement spatial hashing and compare runtimes for a range of problem sizes (grid size, particle count) with the original search methodology. If spatial hashing proves to be more performant, move this implementation into v4.

[erikvansebille]

We just had a deeper discussion about the spatial hashing on curvilinear grids; and then specifically the so-called ORCA-grid used in NEMO (and thus in the MOi dataset).

The issue is that, because of the grid refinement over Canada and Siberia, some hash-cells contain thousands of numerical cells. When creating the hash-map, append operations drastically slow down the hashtable generation but preallocation is not possible because it would require ~2TB arrays.

@fluidnumerics-joe discussed a few alternatives

1. Use a landmask (most NEMO grid files have an flask field) and simply don't hash the land points. This would avoid the large number of grid cells in a few hash cells, because they are (by definition) on land, so particles will not get there. The hash table would also be much smaller, because it would contain less grid cells. And this mimics the unstructured grids, which also don't contain 'land cells'. The disadvantage is that users would need to provide a landmask
2. Use a different, more exotic hash function (for example a Morton or Z-order curve) to better balance the number of grid cells per hash-cell
3. Use a KD-tree instead of a hash-map

We decided to first try option 1, and see how much performance improvement that gives us. If it works, we would need to add some plumbing so that users can easily provide landmasks etc

[fluidnumerics-joe]

WRT point 1. When we don't use a land mask and attempt to pre-allocate and append, the maximum adjacency (the number of parent model grid cells found within a hash cell) is 7450. Output from the spatial hash initialization on the MOi benchmark is shown below:

SpatialHash: Using spherical mesh with 13213618 faces and 74618461 hash cells.
SpatialHash: Maximum adjacency count is 7450.
Traceback (most recent call last):
  File "/home/joe/Projects/Geomar-Utrecht/parcels-benchmarks/MOi-Curvilinear/benchmark_index_search.py", line 130, in <module>
    main()
  File "/home/joe/Projects/Geomar-Utrecht/parcels-benchmarks/MOi-Curvilinear/benchmark_index_search.py", line 126, in main
    run_benchmark(args.memory)
  File "/home/joe/Projects/Geomar-Utrecht/parcels-benchmarks/MOi-Curvilinear/benchmark_index_search.py", line 92, in run_benchmark
    _ = fieldset.U.grid.get_spatial_hash(dh_factor=dh_factor, reconstruct=True)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/joe/miniconda3/envs/parcels/lib/python3.12/site-packages/parcels/xgrid.py", line 373, in get_spatial_hash
    self._spatialhash = SpatialHash(self, dh_factor, reconstruct)
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/joe/miniconda3/envs/parcels/lib/python3.12/site-packages/parcels/spatialhash.py", line 117, in __init__
    self._hash_table = self._initialize_hash_table()
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/joe/miniconda3/envs/parcels/lib/python3.12/site-packages/parcels/spatialhash.py", line 196, in _initialize_hash_table
    face_ids = np.empty((nhashcells, max_adjacency), dtype=np.int32)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
numpy._core._exceptions._ArrayMemoryError: Unable to allocate 2.02 TiB for an array with shape (74618461, 7450) and data type int32

[erikvansebille]

And here is the fmask field from the NEMO-curvilinear data

[fluidnumerics-joe]

@erikvansebille - the MOi dataset I have does not appear to have this field anywhere.. Would you happen to know where to look ?

[erikvansebille]

Yep, I uploaded the 1/12degree version to https://surfdrive.surf.nl/files/index.php/s/A12xp90sZc4V4Ht. Does that work?

[fluidnumerics-joe]

Yep, I uploaded the 1/12degree version to https://surfdrive.surf.nl/files/index.php/s/A12xp90sZc4V4Ht. Does that work?

Indeed this works. I'm putting in the patch now to see if this fixes the issue with spatial hash initialization

[fluidnumerics-joe]

While this definitely solves the pre-allocation issue, the initialization time is still a bit long. I've been playing with a "toy" implementation that uses Morton code of xyz position to buiild the hash table, and that takes just about 1s for the MOi grid.

I've also got some comparisons of the lookup and search/query functionality with a slow naive approach that I know works, and this looks promising.

(parcels) [joe@claymation parcels-benchmarks]$ python moi_morton_code.py 
lon.shape: (3059, 4322), lat.shape: (3059, 4322)
ny, nx = 3058, 4321
Spatial hash construction time: 1.0796 seconds
Number of unique Morton codes: 3773768 / 13213618  (28.559687437611714 %)  
The Morton code for (0.5, 0.5000000000000001, 0.7071067811865475) is: 1015004895
The nearest grid point to (0.5, 0.5000000000000001, 0.7071067811865475) is at index (np.int64(2183), np.int64(3973)) with coordinates (0.5002017021179199, 0.4998987913131714, 0.7070354223251343) and Morton code 1015004895
Distance to nearest grid point: 0.00023668339562623087
The nearest grid point by Morton code to (0.5, 0.5000000000000001, 0.7071067811865475) is at index (np.int64(2183), np.int64(3973)) with coordinates (0.5002017021179199, 0.4998987913131714, 0.7070354223251343) and Morton code 1015004895
Distance to nearest grid point by Morton code: 0.00023668339562623087

I'm looking in to see how incorporating the land mask will change initialization and search time. This approach may be a useful upgrade to the spatial hashing..

[fluidnumerics-joe]

Here's a rough visualization of the morton codes calculated from x,y,z positions on the unit sphere, using the MOi grid lat/lon as input. In essence, the morton code quantizes the x,z,y positions to an integer range of 0-1023 and then computes an integer that interleaves the bits of the quantized x,y,z positions.. From this, we can quickly get a table that relates the morton code of an x,y,z value to a short list of (j,i) indices to check, effectivley building a hash table much quicker.

This would in effect have us swap out our hash function to use a morton code calculation. I'm interested to see how this performs on a large list of particles.. more soon..

[erikvansebille]

Very exciting to see, @fluidnumerics-joe! Is this already at a stage where you can open a PR< so that I can also explore the code and performance (asynchronously) for myself? Keen to see how this Morton code indeed works!

[fluidnumerics-joe]

I will get a branch open this morning and share.