Spatialhash regional spherical improvements#2744
Conversation
When the mesh type is spherical, we now use the actual x,y,z bounding box rather than the unit cube. For regional runs that leverage the spherical barycentric coordinates and built in unit conversion in the velocity field interpolators, this reduces the hit count per hash cell and can improve the search time in the spatialhash.query() Note that quantize_coordinates now clips in float space before casting to uint32. This matters now that spherical queries can fall outside the hash-grid bounds (previously impossible with the unit cube). A negative normalized value cast to uint32 wraps to a huge number, which would have mapped below-range queries to the top bin (with morton code 2^32) erroneously. For spherical grids, points outside the regional box now quantize to edge bins rather than interior bins. In this case, they either fail the exact Morton-key match or are rejected by the point-in-cell check; either case results in the same outcome as before (GRID_SEARCH_ERROR)
Tests for both points inside the domain and explicitly outside the domain to confirm correct behavior for out of bounds PIC checks
…o spherical bounding box
|
I still need to test this out on the example you shared @erikvansebille, but wanted to get this in and make sure that existing tests all pass |
|
looks like there's a few tests to clean up here. |
|
Thanks for this PR, @fluidnumericsJoe! Seems like there are some segmentation errors in the tests. Do you want to try fix them; or can I help? |
I'll get to it today |
|
So, I've narrowed down what's going on here. To recap what we primarily changed here - we've changed the spatialhash initialization to use a more narrow bounding box, based on the actual x,y,z extents for spherical mesh types. effectively, for regional domains this makes the hash grid size smaller. This happens because the number of hash grid cells is fixed, but we're now using a smaller domain extent for the hash grid. The upshot of this change is that we have fewer faces in the "parent" mesh that we need to check, when a particle is found in a given hash cell. This reduces the number of calls to the particle in-cell checks. The issue here is in the construction of the hash table that maps a hash cell index (the morton code) to a list of faces. During construction we have to use the bounding box of a face in the parent grid to find the range of hash cells that overlap. This information is used to pre-allocate a lookup table that maps from parent grid face to hash grid cells; the hash table is contructed as the inverse of that table. If the hash cells are significantly smaller than the parent model face sizes, we end up with a fairly large memory footprint (high number of hash cells per face). I'm looking at tidying up that initializer here to reduce the memory footprint.. |
This figure shows the antimeridian spanning grid that the CI is failing on with OOM errors. The grid locations are plotted in cartesian (x,y,z) space after mapping the latitude, longitude through a spherical coordinate transformation. A box shows a region that we'll zoom in on in the next figure
This zoomed in region shows the hash grid (light gray lines) and the parent grid edges (blue lines). The bouding cube for the cell shaded light blue has 55x59x46 hash cells. Summing the number of hash cells per face across all faces (to get the size of the work arrays for constructing the hash table), each work array is roughly 193 million entries (too big).
In addition to tidying up in the initializer, I'll probably look at putting in a lever to manipulate the hash cell size - with the morton cells, this is controlled by the bit width . A higher bitwidth gives more hash cells per direction. At the moment, we default to 1024. Another thing we could possibly do is be a bit smarter about tossing out hash cells that are clearly outside of elements (bounding box overlaps are kind of wasteful). In both plane views in the zoomed in plot, there's a large number of hash cells that are completely outside the element. |
num_hash_per_face is never zero (quantization is monotone), so the valid mask was always all-True and the scatter an identity permutation. Removing them drops three entry-length temporaries (~35% peak build memory). Accumulate entry counts in int64 to avoid overflow. Hash table output verified byte-identical. Co-authored-by: Claude <noreply@anthropic.com>
…e 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>
|
I pushed up a couple of commits created with claude that focus on memory footprint reduction in the The first commit removes a conditional path that is never traversed; that branch would only be executed if a face is found that doesn't overlap a single hash cell. Since the hash grid bounds are the bounding box of the whole grid, it's safe to assume each face will overlap at least 1 hash cell. Relying on this assumption allows us to trim back on a number of intermediate arrays. The second commit clears previously used intermediate arrays from memory after all the morton codes are calculated. Then since an entry in the hash table lookup consists of the morton code with a corresponding face id and each is a Using the
Even with these improvements, I still think we might need to provide controls to coarsen the hash grid. |
|
Thanks for this further digging, @fluidnumericsJoe! I must say I don't entirely understand what's going on, but taking a step back I think this is the problem:
|
It was working. Reducing the extent of the hash grid ballooned the number of hash cells per face. This happens because reducing the hash grid extent while keeping the number of hash cells the same causes the hash cells to become smaller. With smaller hash cells, more of them are withing each parent mesh face. The number of hash cells per face determines the size of work arrays during hash table construction. For some grids, this led to oom errors during construction. Now the game is to reduce the memory footprint of the hash table construction.
Ultimately, once the hash table is constructed, since we use a compressed sparse row matrix format, hash entries that have no face overlaps don't get stored in the lookup table.
Yes. This optimizes query, but balloons the memory footprint for the hash table creation
Antimeridian is not the issue. It's a ratio of hash grid cell size to parent mesh face/cell size.
This is not an issue..for polar caps we have to use spherical mesh types. Flat mesh types would end up with colinear cells . The changes made in this pr will improve construction time of the hash grid for these cases and help reduce the number of particle in cell checks. |
…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.
…com:Parcels-Code/Parcels into spatialhash-regional-spherical-improvements
…footprint Instead, we opt to unravel the face during the query. This cuts the hash table size down by only storing a single uint64 array (with the start and count) rather than two arrays. The unravel cost in the query is small in comparison to particle in cell checks, which dominate the query costs.
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
for more information, see https://pre-commit.ci
|
I put some effort into having the hash grid cell spacing change according to a budget number of hash cells per parent model face. Rather than passing the choice to users to set the Here's the full progression on the fesom2_square_delaunay_uniform_z_coordinate spherical benchmark (721 faces), each commit measured in its own process:
That last change, with the adjustable hash cell size, makes a huge difference in peak memory consumption and runtime of the hash table construction. In terms of the number of particle in cell checks, below are some stats on the number of faces per hash table key; this is directly related to the number of particle in cell checks that are run for finding particles.
|
…e construction I've removed the bitwidth parameter for construction and opted for this to be calculated internally for optimal memory footprint and query cost
…com:Parcels-Code/Parcels into spatialhash-regional-spherical-improvements
for more information, see https://pre-commit.ci
|
@erikvansebille - I'd be interested to know how this branch works on the example you shared. I won't be able to give that a go until monday, but if you have the time to try it out before me, I'd be interested to know what your experience is like |
2 participants
Description
This PR changes the spatialhash initialization for
sphericalmesh types. Previously the bounding box was set to the unit cube for spherical mesh types. However, spherical mesh types might be desirable for regional simulations (e.g. for improved accuracy in barycentric coordinates, dealing with regions close to the poles, or for leveraging built in unit conversions in Parcels velocity field interpolators).The main problem with the unit cube for regional configurations is that the hash grid resolution is often far too coarse for a regional domain. This results in a high number of mesh cells being aligned with a single morton key. Calling
spatialhash.query()then results in an abnormally high number of particle in cell checks.This is remedied here by choosing the extents of the hash grid as the bounding cube (x,y,z) of the regional domain.
Checklist
mainfor normal development,v3-supportfor v3 support)AI Disclosure