Devalue uxarray to make it less expensive - #1588
Conversation
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.
| # apply aggregation on dask array, TODO: | ||
| aggregated_var = _apply_node_to_face_aggregation_numpy( | ||
| # apply aggregation lazily on a dask array | ||
| return _apply_node_to_face_aggregation_dask( |
There was a problem hiding this comment.
Style/documentation: it may improve readability to make _apply_node_to_face_aggregation_dask and _apply_node_to_face_aggregation_numpy have the same return type (e.g., both return xarray.DataArray).
If there's a compelling reason to not do that, please leave a small comment explaining why, and add a note in the docstring of each method saying the type of the returned object
There was a problem hiding this comment.
Yeah that's a good idea, I'll move the ux.DataArray type inside
| uxda, aggregation_func, aggregation_func_kwargs | ||
| ): | ||
| """Applies a Node to Face Topological aggregation on a Dask array, lazily""" | ||
| import xarray as xr |
| aggregation_var = _apply_node_to_edge_aggregation_numpy( | ||
| uxda, NUMPY_AGGREGATIONS[aggregation], aggregation_func_kwargs | ||
| # apply aggregation lazily on a dask array | ||
| return _apply_node_to_edge_aggregation_dask( |
There was a problem hiding this comment.
Same style/documentation comment as above, but for node to edge case; make numpy and dask have same return type and/or clarify in function docstrings
| dimensions are processed blockwise. Each edge reduces over its two endpoint | ||
| nodes. | ||
| """ | ||
| import xarray as xr |
There was a problem hiding this comment.
(same as before, move xarray import to top of file, see #1575)
| ) | ||
|
|
||
| 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)
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