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 26 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.
erogluorhan
changed the title
.values calls) in the code
erogluorhan
changed the title
.values calls) in the code.values calls)
|
pre-commit.ci autofix |
|
@erogluorhan Sam and I were wondering about your philosophy with respect to dask versus numpy routines. I was thinking that if we have a case where the numpy routine could be faster than a dask routine in some situations, we might want to keep both, or maybe you'd want to go all in on dask for simplicity. if isinstance(uxda.data, np.ndarray):
aggregated_var = _apply_node_to_face_aggregation_numpy(uxda, ...)
elif isinstance(uxda.data, da.Array):
# apply aggregation on dask array, TODO:
aggregated_var = _apply_node_to_face_aggregation_numpy(uxda, ...)
else:
raise ValueErrorThis is something in main's aggregate.py. I fleshed out the dask branch, but now we still have this logic, and I left the numpy routine in on this branch. |
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.
| # gather only the sampled faces (lazily for dask data): many pixels share a | ||
| # face, so deduplicate to materialize the minimal set rather than the whole | ||
| # field, then scatter back via the inverse map | ||
| unique_faces, inverse = np.unique(first_face[mask_has_face], return_inverse=True) |
There was a problem hiding this comment.
This was the one remaining case I had noted originally as "should have a regression test" which doesn't have such a test yet. Looking into it further, now I understand the result is going to be a numpy array anyway; this doesn't change inputs or outputs at all, it only changes internal computations a bit for a possible speedup.
I suspect this functionality is already covered by the test suite; a variety of tests call to_raster() already, and I confirmed (via adding a print statement) these lines of code are getting run >10 times during the test suite.
So, actually, no need to add any tests here, I think. Feel free to mark this as resolved once you took a quick look, if you agree it doesn't need tests.
|
|
||
| def test_node_to_face_dask_reproduces_numpy(gridpath): | ||
| # the numpy (eager) and dask (chunked) branches must agree | ||
| da = pytest.importorskip("dask") # dask-backed branch requires dask |
There was a problem hiding this comment.
Remove the importorskip calls everywhere; tests should crash if dask is not installed, not be skipped silently (well, mostly silently). Existing tests currently on main already use dask without importorskip, and the CI for the test suite installs dask, so it should be safe to assume dask is available.
There was a problem hiding this comment.
If we're assuming that dask is available for these tests though, why can't we assume that in the source code? My strategy lately has been to hedge with the assumption that dask will eventually be optional at the package level, but you're saying here I should be moving in the opposite direction?
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. |
4 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