Fix interp performance regression from #9881 - #11496
Draft
oshokothari07-ai wants to merge 1 commit into
Draft
Conversation
When the target coordinate varies along a dimension of the interpolated object, interpolate_variable passes vectorize=True to apply_ufunc, which uses np.vectorize. np.vectorize loops in Python over every non-core dimension, not only the dimensions that actually needed vectorizing, so da[t, r, z].interp(z=target[t]) made t*r calls into the interpolator instead of t. Declare the remaining dimensions as core dimensions so they are handed to the interpolator in bulk; _interpnd already treats leading dimensions as constant. This is restricted to in-memory arrays because making a dimension a core dimension forces dask to rechunk along it. For the example in the issue this reduces _interpnd calls from 7930 to 122 (a factor of n_r = 65). Fixes pydata#10683
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Closes #10683.
This PR addresses the performance regression in
DataArray.interpandDataset.interpintroduced by #9881 when the target coordinate varies along a dimension of the source array.The regression occurs because
interpolate_variableenablesvectorize=Truewhen callingapply_ufunc. Withnp.vectorize, every non-core dimension is iterated in Python, including dimensions that do not actually require vectorization. For cases such asda[t, r, z].interp(z=target[t]), this causes the interpolator to be called once for every(t, r)pair instead of once pert.The fix promotes the remaining non-vectorized dimensions to core dimensions so they are passed to
_interpndin bulk._interpndalready treats these leading dimensions as constant, so this avoids the unnecessary Python-level looping while preserving interpolation behavior. This optimization is applied only to in-memory arrays because making these dimensions core for chunked arrays would force dask to rechunk along them.On the example from the issue, this reduces
_interpndcalls from 7930 to 122 and reduces runtime from roughly 2.2 s to around 110 ms in my local testing.Testing
xarray/tests/test_interp.pyandxarray/tests/test_missing.py: 283 passed, 74 skipped, 1 xfailed.ruff checkandruff formatpass.Limitations
This is a targeted workaround rather than the more general redesign of the
apply_ufuncinterpolation path discussed in the issue. It recovers most of the regression but does not restore the performance seen in 2024.11.0.I was not able to run the full test suite on
mainlocally because the repository currently requires Python 3.11+, while my local environment is Python 3.10.11. The implementation was validated against byte-identical code for the affected functions, but CI will provide the final verification across supported Python versions, platforms, the full dask test matrix, and minimum-version configurations.I also did not run
mypylocally. The only typed addition isbulk_dims: tuple[Hashable, ...], and CI will validate type checking against both the current and minimum supported environments.Checklist
doc/whats-new.rstapi.rstupdate not requiredAI Disclosure
This PR contains AI-generated content.
I have tested any AI-generated content in my PR.
I take responsibility for any AI-generated content in my PR.
I used Claude as an engineering assistant during the investigation and implementation. I reviewed the changes, validated the implementation and tests, and take responsibility for the final code and this pull request.