Feature/morton hashing

[fluidnumerics-joe]

• [ X ] Chose the correct base branch (main for v3 changes, v4-dev for v4 changes)
• [ X ] Fixes Spatial hashing and curvilinear search performance improvements #2144

Overview

This PR brings in a fairly substantial change to the spatial hashing. In experimenting with the spatial hashing in the MOi benchmark, we found that the spatial hashing spent a large amount of time in initialization. Further, the search times were not providing an improvement over the naive implementation for searching in curvilinear grids. This was found to be related to a few key issues

1. The chosen hash function that relates hash keys to a short list of faces/cells in the curvilinear grid produced long searches in areas where the curvilinear grid approached the poles (e.g. over canada and siberia in the tripolar grid)
2. The list of lists structure for the hash table required a large number of append operations (within nested for loops) during hash table creation. Attempts to pre-allocate and use a CSR format for the hash table are hampered by point (1); a large number of overlaps of curvilinear cells with the hash grid resulted in large allocations that are prohibitive
3. Querying was a challenge to vectorize with the list of lists hash table and computation of barycentric coordinates for particle in cell checking was found to be expensive

This PR addresses these issues by changing the hashing strategy. The "hash function" is a function that takes in a position on the grid and returns an integer value that has some correlation with spatial locality. In other words the hash function of two nearby points in space, should return integer values that are close to each other in value.

Morton encoding for hash table construction

Here, we opt to use Morton encoding of three floating point values to compute the hash of a position. In spherical meshes, the latitude and longitude points are converted to x,y,z locations on the unit sphere. On "flat" meshes, z is fixed to zero and x and y are taken as is, with no conversions. The hash table is constructed by relating the (j,i) indices of the curvilinear grid to the morton code of the centroid of each curvilinear cell. The hash table is quick to calculate and is simply stored as a compressed sparse row (CSR) matrix and sorted in order of increasing morton number. This last step (sorting) adds some cost to the initialization but saves us considerably in the querying stage since np.searchsorted can be used (instead of np.search).

A comparison of the construction time for the MOi benchmark is shown below for v4-dev and the new proposed morton encoding spatial hash function.

Morton encoding for particle position queries

With the hash table as a CSR matrix, it's a bit easier to vectorize queries over a large number of particles. Essentially, the morton code of all particle positions are calculated and searched (using np.searchsorted) against the hash table. This returns, for each particle, a short list of (j,i) indices to check for each particle. We then perform distance minimization between each particle and it list of curvilinear cells where the distance is measured from the curvilinear cell centroid. This avoids the need for barycentric coordinate calculation. At the end, the method returns, as before, a list of (j,i) indices, one for each particle position, where the particle is found.

Below is a figure comparing the original v4-dev naive search, with the previous spatial hashing search, and the new morton spatial hashing for the MOi benchmark across a range of particle counts. Note that that new morton hashing search times include the spatial hash reconstruction, which is minimal (~2s on system76 Lemur with NVMe drives)

The above figure shows improved scaling in the morton hash based search functionality. However, the 'v4 search indices' and 'v4 hash table' runtimes were generated on @erikvansebille 's work-station and the 'v4 morton' runtimes were on mine, making direct comparison challenging.

Below shows the spatial hashing search (what's currently on v4-dev) and the propose "v4 morton " implementation (x-axis is log scale).

[erikvansebille]

This is fantastic, @fluidnumerics-joe! The performance improvements are amazing!! For 1M particles, I now get a 2 seconds time, instead of the 600 seconds with the spatial hash map 🤩.

What a great leap towards a fast v4!

Is this implementation also useful for unstructured grids?

[fluidnumerics-joe]

Is this implementation also useful for unstructured grids?

I want to follow up on this after this pr. I suspect yes. This SpatialHash class just requires face centroids which are easy to get from uxarray.

The query method as it is written right now works with j,i indices. This would need to be adapted to work with single dimension face indices. It shouldn't be too hard to do

[erikvansebille]

I want to follow up on this after this pr. I suspect yes.

Yes, fully agree that this should be a separate PR