Optimized connectivity: Simultaneous face_edges and edge_nodes - #1560
Optimized connectivity: Simultaneous face_edges and edge_nodes#1560cmdupuis3 wants to merge 20 commits into
Conversation
cmdupuis3
added
the
improvement
|
I think all the hard parts of the merge are done. At this point, the only failures seems to be sorting issues. |
cmdupuis3
added
scalability
@cmdupuis3 Thanks for your work, I will review it as soon as possible. And do you have any idea why the CIs are all failing? |
|
The CI is failing because the new algorithm returns the results in a different order. I was thinking we can add some sorting mechanism, at least for legacy behavior. Phillip was evidently aware of this issue too. It depends on if you need to support that... Personally I'd be okay with just changing it, but I think you would be a better judge of the situation. |
| pass | ||
|
|
||
| edge_nodes, face_edges = _build_edge_node_connectivity( | ||
| grid.face_node_connectivity.values, grid.n_nodes_per_face.values |
There was a problem hiding this comment.
Please think about if you can get rid of .values calls to be chunked-array compatible.
There was a problem hiding this comment.
On it, I think Philip was trying to scalarize here but it introduced this regression.
There was a problem hiding this comment.
So the best I could come up with is
face_nodes, n_nodes_per_face = dask.compute(
grid.face_node_connectivity.data, grid.n_nodes_per_face.data
)
edge_nodes, face_edges = _build_edge_node_connectivity(
face_nodes, n_nodes_per_face, grid.n_node
)...but I'm unclear on if we're committing to dask or not. Other parts of the repo imply that dask is still sort of an optional load, and doing this would pretty much require a dask import for most grids, unless we have some kind of numpy/dask conditional.
There was a problem hiding this comment.
Alright, I hacked it a bit to get this:
computed = xr.Dataset(
{
"face_nodes": grid.face_node_connectivity.variable,
"n_nodes_per_face": grid.n_nodes_per_face.variable,
}
).compute()
edge_nodes, face_edges = _build_edge_node_connectivity(
computed.face_nodes.data, computed.n_nodes_per_face.data, grid.n_node
)There was a problem hiding this comment.
Question for @erogluorhan, does uxarray have any plans to support chunked grids? I was under the impression that data variables can be chunked, but that it is always safe to assume the grid itself can be fully loaded into memory. Is that incorrect?
That said, using more dask-compatible syntax like this doesn't seem like it would have a downside (no need to block merging on this question).
The optimized edge builder deduped half edges with a numba hash map, which
numbered edges in first-encounter order. Edges had previously been numbered
lexicographically by their (min_node, max_node) pair, as a side effect of the
np.unique(..., axis=0) the hash map replaced.
Global edge index is a public identity: it indexes edge_lon/edge_lat, edge
centered data variables, and edge_node_distances, so renumbering silently
re-pairs user data with different physical edges. It also broke the five
TestQuadHexagon connectivity tests, which assert on edge_node, face_edge,
node_edge, edge_face and face_face -- all the same renumbering cascading
through the derived connectivities.
Sort as the dedup mechanism instead of hashing. Node indices are dense
integers in [0, n_node), so a counting sort buckets the half edges by their
first node without any comparisons, and sorting each bucket by its second node
leaves the duplicates adjacent -- the dedup then falls out of the same walk.
Buckets hold one entry per edge incident to a node, so on a real mesh they are
tiny (node degree, typically under ten) and an insertion sort finishes them.
A bucket above MAX_INSERTION_SORT_SIZE is heap sorted so that a degenerate
mesh cannot degrade the build quadratically; np.argsort is deliberately not
used there, as numba's implementation degrades badly on structured input.
Half edges are identified throughout by their flat face_node_connectivity
index, which is also the face_edge_connectivity slot they are written back to,
so the sort needs a single permutation array and no mapping back.
This is faster and leaner than the hash map it replaces. On a synthetic one
million face quad mesh, measured by peak RSS rather than tracemalloc, which
does not observe numba's typed dict allocations:
dict build 397.3 ms 239.5 MB
bucket sort 85.1 ms 91.6 MB
The five legacy tests now pass unchanged. Adds order invariant coverage for
the canonical ordering and the face_edge positional contract, plus high degree
nodes either side of the insertion sort threshold.
Also casts n_nodes_per_face back to INT_DTYPE, so that the builder is not
compiled a second time for int64, and restores a blank line dropped between
two top level functions.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
ASV BenchmarkingBenchmark Comparison ResultsBenchmarks that have improved:
Benchmarks that have stayed the same:
|
Sevans711
left a comment
Sevans711
left a comment
There was a problem hiding this comment.
These changes look clean overall! It is great to see the possibility for things that previously took hundreds of seconds to maybe be reduced to a few seconds or less (quoting those numbers from #1195).
I left some inline comments and also have a few bigger questions below:
- (1) Seeing some benchmarks improvements here is encouraging. Though, in order to get better confirmation these changes really solve the linked issues, would you be able to run it on at least one large grid on an HPC system and share the results? Both linked issues referred to large grids.
(You don't necessarily need to reproduce the same exact results or use the same exact grids from the original issues, just looking for at least one large grid showing what the timing looks like on main right now and what it looks like after these changes.)
- (2) Are the increases to memory usage in large grids acceptable here? E.g., if
edge_node_connectivityis requested but user has no plans to utilizeface_edge_connectivity,this PR leads to increased memory cost.
On grid=ux.tutorial.open_grid("outCSne30-vortex"), I checked grid._ds.nbytes after a few commands. Upon loading, it is 259236 bytes. On main after grid.edge_node_connectivity it becomes 432036 bytes, and after grid.face_edge_connectivity it becomes 604836 bytes. On this PR, just doing grid=ux.tutorial.open_grid("outCSne30-vortex"); grid.edge_node_connectivity leads to 648036 bytes. (The extra ~40000 bytes are from n_nodes_per_face, though that is basically just a rounding error when compared to the ~200000 extra bytes from storing face_edge_connectivity.) Of course, the numbers here are very small; this is a tiny example. But, they demonstrate that the concern about memory costs may be somewhat plausible, at least. (Unless everybody who asks for edge_node_connectivity also wants face_edge_connectivity?)
| centroid_x = grid.node_x.values[edge_nodes].mean(axis=1) | ||
| centroid_y = grid.node_y.values[edge_nodes].mean(axis=1) | ||
| centroid_z = grid.node_z.values[edge_nodes].mean(axis=1) | ||
| centroid_x, centroid_y, centroid_z = _normalize_xyz(centroid_x, centroid_y, centroid_z) |
There was a problem hiding this comment.
This seems like a substantive change to the test. Can you clarify why this change was needed? Reply here is fine, not trying to say it is wrong, just not understanding yet why it was changed here.
| ("ugrid", "quad-hexagon", "grid.nc"), | ||
| ("ugrid", "geoflow-small", "grid.nc")]) | ||
| def test_connectivity_edge_node_canonical_order(gridpath, grid_parts): | ||
| """Test that constructed edges are numbered in lexicographic node order.""" |
There was a problem hiding this comment.
I didn't realize lexicographic order was supposed to be guaranteed for constructed edges. I see now that the corresponding docstring for _build_edge_node_connectivity has been updated accordingly. Is that change being introduced intentionally by this PR?
If yes, replying with yes here is sufficient and the change looks good to me. In the future if you can mention changes like this too somewhere in the PR overview or as comments in thread, that would have helped with reducing time it takes to review!
| # Check edge coordinates already exist, if they do this might cause issues | ||
|
|
||
| if "n_edge" in grid.sizes: | ||
| # TODO: raise a warning or exception? |
There was a problem hiding this comment.
This TODO item should be completed before merging, to avoid introducing a new way for the code to fail silently. I would prefer if an exception is raised, not just a warning. Or (more work but maybe cleaner long-term?) figure out what could actually cause this, and debug it?
Minor style suggestion: change to "n_edge" in grid.dims instead, to make it clearer that the actual sizes of the dimensions are not relevant here.
|
|
||
| edge_node_attrs = ugrid.EDGE_NODE_CONNECTIVITY_ATTRS | ||
| edge_node_attrs["inverse_indices"] = inverse_indices | ||
| # HACK: this is lieu of an xarray equivalent to `da.compute(a, b)` |
There was a problem hiding this comment.
Clarifying: what is the purpose of this hack - why not just use something like:
computed_face_nodes = grid.face_node_connectivity.variable.compute()
computed_n_nodes_per_face = grid.n_nodes_per_face.variable.compute()I suppose there could theoretically be a speedup from doing a single compute() call instead of two compute calls?
Or, maybe it is some other reason. Could you add a comment in the code to clarify?
| n_nodes_per_face = _build_n_nodes_per_face( | ||
| grid.face_node_connectivity.values, grid.n_face, grid.n_max_face_nodes | ||
| n_nodes_per_face = ( | ||
| (grid.face_node_connectivity != INT_FILL_VALUE).sum(axis=1).astype(INT_DTYPE) |
There was a problem hiding this comment.
Briefly calling attention to this change, it looks like a huge speed up in the ASV benchmarking suite, which is awesome! It wasn't mentioned in the PR overview / comment thread so I was surprised to see it.
EDIT: actually, many of the n_nodes_per_face tests in the benchmarking suite seem unaffected by this change or even slightly slower because of it. Is this expected behavior?
Confirming: does it also speed things up for large grids on HPC clusters? (Testing with one large grid would be sufficient.)
Yeah, so I was wondering about that too, which is why I proposed having a dispatcher for combined connectivity routines #1559. However, in practice, I think the most likely situation with the grids we support is that both face/edge and edge/node connectivity are the missing connectivities, which is why they were interested in doing this. There could definitely be a discussion on the design, but I think that's out of scope for this PR. |
Ah, yes that makes sense, and seems like a good solution. I also agree a full design like that should probably be scoped to a different PR. Before this PR merges I would still be interested to hear from other uxarray developers/users to confirm whether the increased memory cost is an acceptable tradeoff for the possible huge speedup in the meantime. E.g., @erogluorhan, @rajeeja, @rljacob, if you get a chance to look into this, what are your thoughts? |
|
The speedups are very real. These are from the new connectivity benchmarks in PR #1633, so this is serial performance. I can add chunked performance to the benchmarks as well. I didn't notice any major memory degradation in glancing at top, if anything it seemed better with this branch. Maybe I can add peakmem benchmarks as well.
|
|
@erogluorhan I'm working on some peakmem benchmarks, here's the results for this PR. ASV main versus merge-OFE branch--- face_face:
edge_node:
|
5 participants
Would close #1138, #1196
Related to #1180
Supercedes #1195
Overview
This set of changes optimizes face_edge, edge_node, and face_face connectivity. face_edge and edge_node connectivity are combined into one routine, while face_face is optimized stand-alone.
PR Checklist
General
Testing
Documentation
_) and have been added todocs/internal_api/index.rst