Optimized connectivity: Simultaneous face_edges and edge_nodes - #1560
Optimized connectivity: Simultaneous face_edges and edge_nodes#1560cmdupuis3 wants to merge 30 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).
There was a problem hiding this comment.
Per @cmdupuis3 's latest comment:
This looks great, and it is more beneficial than just "hacking" I believe (please document it accordingly wherever possible in your code):
-
Using
xr.Dataset.compute()guarantees two major advantages:- Bundled instead of separate
.compute()calls on a couple DataArrays - uses same upstream dask graph and gets rid of duplicate overhead - Doesn't break optional Dask - If given chunked arrays, uses Dask; otherwise, skips compute under the xarray's hood and respects numpy arrays
- Bundled instead of separate
There was a problem hiding this comment.
Per @Sevans711's latest comment:
The intent with Grid's I/O design was already to have it support chunked features; however; full end-to-end "Grid supports chunkedness" is probably impossible - there are global-topology and spatial-index operations like what we see here in this PR that could probably never avoid materializing the whole array.
- Grid I/O is already made chunked - there might be things we can further fix though
- Some of the operations can be made chunk-compatible
I think a smart way forward would be to (1) keep (or make) the chunk-local geometric operations lazy, and (2) leave ( and explicitly document) the global-topology/KDTree/validation operations as materialized (whole array computations) rather than attempting to "fix" every .values or .compute() - this PR shows there is still room for optimization for even those though
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:
Benchmarks that have got worse:
|
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.
There was a problem hiding this comment.
Yeah, so this test was basically passing by coincidence, and whether it passes depends on which edge is at index 0. It also skips normalization. I only found this because prior to implementing sorting, this test was failing, which flagged it as something for Claude to look at, and it just happened to be wrong for another reason. With the new sorting algorithm, the original test would be passing though.
| ("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!
There was a problem hiding this comment.
The sorting was implemented to support compatibility with legacy behavior. If we decide to deprecate that behavior we could throw out all the sorting and revert this a little.
There was a problem hiding this comment.
Can you clarify a bit further? If you can provide at least one example of a legacy behavior this change was implemented to support, that should make it much easier for me to understand.
| # 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.
There was a problem hiding this comment.
Done, I think there could be some downstream effects in _ugrid though, but we should fix it there anyway.
|
|
||
| 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?
There was a problem hiding this comment.
Because xarray/dask would end up walking the entire structure twice if you compute them separately. Dask has the multi-arg dask.compute(a, b) which allows it to walk the structure once and iterate all the vars it needs simultaneously, but fsr xarray doesn't have native support for that. I opted to hack this structure in to backdoor in the dask behavior while staying dask-agnostic here.
There was a problem hiding this comment.
Noting that @erogluorhan clarified the reasoning above, see #1560 (comment)
Keeping my comment here as "unresolved" for now as a reminder to update comment within the code itself 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.)
There was a problem hiding this comment.
Removing .values keeps it lazy, and also helps with reducing the graph traversals on line 173
There was a problem hiding this comment.
I agree removing .values is a good idea, but this does more than that: it switches from using the numba _build_n_nodes_per_face to a numpy sum. How does this actually affect time and/or memory usage for large grids?
Also, I noticed a potential accuracy error here: should refer to axis by name, not number, because xarray arrays should support any axis order, and face_node_connectivity might be user-supplied, so it might be transposed compared to the order you expected here.
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:
|
|
pre-commit.ci autofix |
|
pre-commit.ci autofix |
I realize I never responded to this.
All that said, there seems like a big win on both performance and memory utilization in this PR. |
erogluorhan
left a comment
erogluorhan
left a comment
There was a problem hiding this comment.
@cmdupuis3 a quick observation:
- I believe this PR is not documenting this new way of combined constructions, so please mention it in both
face_edgeandedge_nodedocstrings for the user's reference.
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