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 16 commits into
.values calls)#1588Conversation
cmdupuis3
added
scalability
ASV BenchmarkingBenchmark Comparison ResultsBenchmarks that have improved:
Benchmarks that have stayed the same:
Benchmarks that have got worse:
|
|
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.
| ) | ||
|
|
||
| if self.values.size == self.uxgrid.n_face: | ||
| if self.size == self.uxgrid.n_face: |
There was a problem hiding this comment.
These changes look good here; no further revision needed in this area, just flagging #1616 as related. Fixing 1616 will lead to these lines being replaced by something like if self._face_centered() instead.
| "i,...i", self.uxgrid.face_areas.values, self.values | ||
| ) | ||
| else: | ||
| # dask-backed data: xr.dot keeps the reduction lazy |
There was a problem hiding this comment.
Does integration need a regression test which is sensitive to dask versus numpy? Perhaps, assert that the result is a dask array, and that the values equal the numpy version, similar to the tests already added by this PR into test_topological_agg.
|
|
||
| assert numpy_result.dims == dask_result.dims | ||
| assert np.allclose(numpy_result.values, dask_result.values, equal_nan=True) | ||
|
|
There was a problem hiding this comment.
This test currently passes on main, so it does not serve as a regression test yet. One possible fix is to add: import dask; assert isinstance(dask_result.data, dask.array.Array). Then, it fails on main but passes on this branch.
This applies to the node_to_edge case below as well.
There was a problem hiding this comment.
The idea is that if someone is expecting bitwise reproducibility and switching to dask flips bits, these tests will catch it.
There was a problem hiding this comment.
That makes sense. Clarifying, I think the bitwise reproducibility part of these tests are important to include and your implementation of them looks good!
Separately, I think there should also be tests to ensure the result is actually a dask array. But, those don't need to be separate test functions, for example it would be sufficient to add them via something like the assert statement above.
There was a problem hiding this comment.
Separately, I think there should also be tests to ensure the result is actually a dask array. But, those don't need to be separate test functions, for example it would be sufficient to add them via something like the assert statement above.
This makes sense! However, either they need to be in separate tests or the naming of the test case here should slightly be adjusted (currently the name suggests whether input output values and dims are all close)
erogluorhan
changed the title
.values calls) in the code
erogluorhan
changed the title
.values calls) in the code.values calls)
|
pre-commit.ci autofix |
for more information, see https://pre-commit.ci
3 participants
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