Avoid converting chunked data to Numpy arrays (i.e. .values calls) - #1588
Avoid converting chunked data to Numpy arrays (i.e. .values calls)#1588cmdupuis3 wants to merge 31 commits into
.values calls)#1588Conversation
cmdupuis3
added
scalability
ASV BenchmarkingBenchmark Comparison ResultsBenchmarks that have improved:
Benchmarks that have stayed the same:
|
|
pre-commit.ci autofix |
for more information, see https://pre-commit.ci
|
There's some spurious asv results from machine variability, but the benchmarks consistenly show peak-mem reductions and some speedups for cross-sections. |
|
pre-commit.ci autofix |
Sevans711
left a comment
Sevans711
left a comment
There was a problem hiding this comment.
Hi @cmdupuis3, thank you for proposing these changes! Overall these look like good clean changes which should help improve the scalability of uxarray.
I have some notes/suggestions/requested changes. Primarily, I noticed that there are missing regression tests. This may be especially important for places with significant changes (more significant than just replacing obj.values with obj.data):
UxDataArray.integrate()UxDataArrayCrossSectionAccessor.__call__()uxarray.plot.matplotlib._nearest_neighbor_resample()RemapAccessor.apply_weights()
It may be nice to add a regression test for UxDataset.to_xarray() too, even though the changes here are minor, since that feels like a very core part of the functionality.
The tests could be similar to what you added already in test_topological_agg.py (plus my additional request on that file): ensure that numpy and dask inputs ultimately give the same values (plus assert that the dask inputs lead to dask outputs).
Other parts of the code changes here might benefit from similar tests, but I don't know if that should be necessary. For example, the changes in _geos just replace obj.values with obj.data; does a regression test need to be added for that or no? Curious to hear from @erogluorhan and/or @rajeeja on this question in particular.
Misc. note: I'm not sure if I fully understand the changes to uxarray/cross_sections files. I need to make sure to take a closer look at those during a subsequent review.
erogluorhan
changed the title
.values calls) in the code
erogluorhan
changed the title
.values calls) in the code.values calls)
If I am understanding this correctly, question is about the use of routines, not data structures. If so, using standard NumPy functions directly on chunked Dask-backed Xarray objects should be completely safe with chunks. Let me know if this helps. |
A stash pop was committed in 9cfd53e with its markers intact, making the module a syntax error. Both sides were additive, so keep both. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Sevans711
left a comment
Sevans711
left a comment
There was a problem hiding this comment.
This looks almost ready to merge, I just have a few minor comments remaining!
Requires response:
- I noticed the
mpas_ocean.Integrate.time_integrate('480km')benchmark has worsened slightly. Is this expected behavior? Do you think this is real, a rounding error, or something else? - I would suggest to remove all of the
pytest.importorskipcalls; I left an inline comment with more details.
The rest of this comment does not require response. It includes the notes I wrote down, for future reference.
Other Benchmarking:
- The cross sections benchmarks seem to have improved, which is great!
- I'm ignoring peakmem "improvements" (as per #1605)
- Even though other benchmarks didn't show notable improvements, I think that is fine, because they didn't get worse, either. Also, this PR isn't fully closing the original issue, anyway, it is just addressing part of it.
Testing:
I checked all the cases where I previously requested tests for correctness (numpy & dask give same result) and regression tests (check dask inputs cause dask outputs); here is what I found:
- (Looks done to me! Confirmed tests pass here, fail on main.)
UxDataArray.integrate() - (Looks done to me! Confirmed tests pass here, fail on main. Also, the tests do a good job covering a variety of possible failure cases, by including multiple dimensions and some nan / some non-nan values.)
UxDataArrayCrossSectionAccessor.__call__() - (Not done, but now I'm less sure if it's necessary; left an inline comment to discuss.)
uxarray.plot.matplotlib._nearest_neighbor_resample() - (Looks done to me! Confirmed tests pass here, fail on main. Also, the tests do a good job covering nontrivial cases, like handling multiple dimensions.)
RemapAccessor.apply_weights() - (Not done, but optional; no need to block merging for this.)
UxDataset.to_xarray()
Previously I noted I need to look closer at cross sections. After doing so and considering the new cross sections tests, I am feeling convinced that the cross sections changes should be working as intended. Noting that the extra dims included in those tests might also help with #1461.
It isn't real, you can rerun the benchmarks and it'll probably disappear. There's always a chance with the smaller benchmarks that they'll trigger due to machine variability. 1.1x is pretty suspect. |
|
While reviewing this I noticed the numpy aggregation paths drop coords and attrs while your new dask paths preserve them via |
|
pre-commit.ci autofix |
for more information, see https://pre-commit.ci
Sevans711
left a comment
Sevans711
left a comment
There was a problem hiding this comment.
Looks to me like this is ready to merge into main!
Any guesses as to why the mpas_ocean.CrossSections.time_const_lat benchmarks show improvements now, but didn't show improvements in prior iterations of this PR? (This question is non-blocking, since it's an improvement! Just curious if you have insights about this.)
|
@Sevans711 They've shown improvements for a while now, that's not recent. There are some changes to how cross-sections are caluculated on chunked grids that served as the |
Ah, woops! I confused myself by looking at my prior comment mentioning the benchmarks. Thank you for clarifying! |
|
I have copilot access on github and will ask it to do a review. |
There was a problem hiding this comment.
Pull request overview
Improves lazy execution by replacing eager .values conversions with Dask/xarray-compatible operations.
Changes:
- Adds lazy remapping, aggregation, integration, and cross-section paths.
- Preserves chunked backing arrays during grid and format conversions.
- Adds Dask parity tests for key operations.
Reviewed changes
Copilot reviewed 17 out of 17 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
uxarray/remap/structured.py |
Preserves lazy arrays during reshaping. |
uxarray/remap/apply_weights.py |
Applies sparse weights blockwise. |
uxarray/remap/accessor.py |
Documents lazy weight application. |
uxarray/plot/matplotlib.py |
Loads only sampled faces. |
uxarray/io/_scrip.py |
Preserves lazy SCRIP coordinates. |
uxarray/io/_geos.py |
Preserves GEOS backing arrays. |
uxarray/grid/slice.py |
Adds lazy connectivity remapping. |
uxarray/grid/grid.py |
Avoids unnecessary grid materialization. |
uxarray/cross_sections/sample.py |
Removes eager Numba gathering. |
uxarray/cross_sections/dataarray_accessor.py |
Implements lazy cross-section gathering. |
uxarray/core/dataset.py |
Preserves backing arrays in duals. |
uxarray/core/dataarray.py |
Adds lazy integration and transformations. |
uxarray/core/aggregation.py |
Implements Dask topological aggregation. |
test/test_cross_sections.py |
Tests lazy cross-sections. |
test/precomputed_weights_test.py |
Tests lazy sparse remapping. |
test/grid/integrate/test_basic.py |
Tests lazy integration. |
test/core/test_topological_agg.py |
Tests Dask aggregations. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| def _build_remap(orig_indices, size): | ||
| remap = np.full(size, INT_FILL_VALUE, dtype=INT_DTYPE) | ||
| remap[orig_indices] = np.arange(len(orig_indices), dtype=INT_DTYPE) | ||
| return remap |
| # Apply remapping (vectorized; stays lazy when the connectivity is dask) | ||
| ds[conn_name] = xr.apply_ufunc( | ||
| _remap_indices, | ||
| ds[conn_name], | ||
| kwargs={"remap": remap}, | ||
| dask="parallelized", | ||
| output_dtypes=[INT_DTYPE], |
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Partly addresses #1583
Overview
This PR is to resolve suboptimal usage of
.valuesthroughout the repo, primarily by deferring to lazy xarray and dask operations. See the table of usage sites in Issue #1583.This specific PR is for the first two tables of usage sites; dead code and bugfixes aren't really included here. Aside from simple cases, some cases were solved by branching on whether the data type was already chunked, so scalability for some parts will depend on if you're using numpy or xarray/dask arrays at those points.
PR Checklist
General
Testing
Documentation