[headtr1ck]

Fixes the performance regression from #9881 when the new coordinates vary along a dimension of the interpolated object, e.g. da[t, r, z].interp(z=new_z[t]).

• apply_ufunc accepts dimension names for vectorize, e.g. vectorize=["t"]: it only loops over these broadcast dimensions, in the given order, and passes all others to func in bulk (which must then broadcast its inputs following numpy rules).
• interp uses it to only loop over the dimensions its new coordinates vary along, i.e. t instead of t and r like np.vectorize. Works for numpy and dask alike, without changing chunks or the task graph.
• _localize and the datetime conversion run once per call before apply_ufunc instead of once per loop iteration.
• Fix interp with empty new coordinates for method="linear" / "nearest".
• New asv benchmarks InterpolationVectorized / InterpolationVectorizedDask.
• Docs: apply_ufunc docstring examples from basic to advanced. The apply_ufunc_vectorize_1d example notebook is removed in favor of the xarray-tutorial interpolation example. The new vectorize option is documented there and on the automatic vectorization page in apply_ufunc: only loop over some dimensions with vectorize=[dims] xarray-contrib/xarray-tutorial#379.

Looping over explicit dimensions could also enable multidimensional weights in polyfit (#9846): loop over the dimensions the weights vary along and solve all other series in one lstsq call.

This is an alternative to #11496, which declares the free dimensions as core dimensions (numpy only).

benchmark	before #9881	main	#11496	this PR
numpy, linear	1.7 ms	1173 ms	25 ms	9.0 ms
numpy, cubic	110 ms	605 ms	17 ms	18.1 ms
dask, linear	11.8 ms	1271 ms	1265 ms	30.7 ms
dask, cubic	102 ms	789 ms	807 ms	52.5 ms

The remaining gap for linear compared to before #9881 comes from the Python loop over t: each iteration builds its own scipy interpolator. Before #9881, the shared dimensions were added as extra interpolation dimensions, which needs merged dask chunks, only works for n-D capable methods with numeric coordinates, and leaks NaNs between slices.

• Closes interp() performance regression in 2025.1.0+ #10683, closes Fix interp performance regression from #9881 #11496
• Tests added
• User visible changes (including notable bug fixes) are documented in whats-new.rst

AI 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.
  • Tools: Claude Code with Opus 5.5

🤖 Generated with Claude Code

[headtr1ck]

It's somehow much uglier than just vectorize=True but it seems to work. Any ideas how to make this better understandable?

[dcherian]

I can take a look later. My conclusion a year ago was that there are no real nice solutions without reinventing parts of apply_ufunc. Perhaps there's a vectorization helper for apply-ufunc that we can use here?

But this is why i accepted the slowdown for #9881; indeed it was much worse than expected though.

[dcherian]

: each iteration builds its own scipy interpolator.

we can totally construct an interpolator outside the loop no?

[headtr1ck]

: each iteration builds its own scipy interpolator.

we can totally construct an interpolator outside the loop no?

I let Claude benchmark several variations of that and neither showed a clear performance improvement. There are always tradeoffs for small or large arrays while just adding more code boilerplate.

[headtr1ck]

I can take a look later. My conclusion a year ago was that there are no real nice solutions without reinventing parts of apply_ufunc. Perhaps there's a vectorization helper for apply-ufunc that we can use here?

But this is why i accepted the slowdown for #9881; indeed it was much worse than expected though.

Thanks!

I couln't find any such helper, but thinking about it we could totally turn the new interp code into such a helper. Potentially even add it as a new vectorized="shared" option.

[headtr1ck]

@dcherian I turned the loop into a helper and added it to apply_ufunc as vectorize="shared": it only loops over the broadcast dimensions along which more than one input varies and passes all other broadcast dimensions to func in bulk. interp now just uses apply_ufunc(..., vectorize="shared"), so _interpnd is almost back to its version on main. Performance is the same or slightly better than before, see the updated description.

[This is Claude Code on behalf of Michael Niklas]

[github-actions]

🚀 Some benchmarks got better. To infinity and beyond!

Change	Before [fd976b7]	After [2f92388]	Ratio	Benchmark (Parameter)
🚀	517±9ms	19.4±0.4ms	0.04	interp.InterpolationVectorized.time_interpolation_vectorized('cubic')
🚀	1.18±0s	7.62±0.3ms	0.01	interp.InterpolationVectorized.time_interpolation_vectorized('linear')
🚀	567±2ms	63.4±0.4ms	0.11	interp.InterpolationVectorizedDask.time_interpolation_vectorized('cubic')
🚀	1.44±0.02s	89.4±2ms	0.06	interp.InterpolationVectorizedDask.time_interpolation_vectorized('linear')

🚀 better · 🐌 worse · 💥 now failing · 🩹 fixed · 🤷 not comparable

Comparing 2f923884 (this PR merged into the base branch) with fd976b73 (base branch).
Only changes by more than a factor of 1.5 are shown, see the workflow run for all results.

[headtr1ck]

@dcherian do you like the new vectorize option? In xarray I didn't find any other usecase for it (most applications are differing only by none or a single dim and there np.vectorize works).
But I guess some very advanced users might be able to utilize it.

[headtr1ck]

@dcherian sorry for the repeated ping. Did you have time to think about this? I think it's a clean interface.

[dcherian]

Sorry for the slow reply, ironically the more interesting PRs tend to get procrastinated upon.

New apply_ufunc(..., vectorize="shared"): only loops over the broadcast dimensions along which more than one input varies, all others are passed to func in bulk (which must then broadcast its inputs following numpy rules).

interp uses it, so it only loops over t instead of t and r like np.vectorize.

Can we instead have vectorize: list[Dimension] instead? So the caller can fully override depending on what works for the function they have provided. I think that's far simpler to reason about, and in tune with the other kwargs.

I am unsure of set vs list, allowing control over loop order seems useful. but maybe it can be order of appearance in args.

cc @shoyer

[headtr1ck]

Sorry for the slow reply, ironically the more interesting PRs tend to get procrastinated upon.

No worries

Can we instead have vectorize: list[Dimension] instead? So the caller can fully override depending on what works for the function they have provided. I think that's far simpler to reason about, and in tune with the other kwargs.

I am unsure of set vs list, allowing control over loop order seems useful. but maybe it can be order of appearance in args.

cc @shoyer

Sounds also like a useful feature. Why not both, explicit dimensions and "shared"? I'm sure both options share some similarities and don't need separate code paths.

[dcherian]

Why not both, explicit dimensions and "shared"?

I worry that "shared" is hard to reason about. So I'd like to start with the most explicit option.

[headtr1ck]

Ok done. Seems sensible to me.

[dcherian]

i thought the point here was to do this for dask on a chunk-by-chunk basis?

[headtr1ck]

Per chunk, _localize never reduced I/O: the core dimension is a single chunk, so each block already contains all of x and localizing only shrank the in-memory array passed to scipy. For linear/nearest that doesn't save time either, since scipy only does a binary search into x (assume_sorted=True), while _localize itself is O(n) (x.to_index() and get_indexer) and ran for every block.

With 1 GB of data (64 × 2M points, 16 chunks) and chunked new coordinates, this PR is as fast or faster than main:

new coordinates	method	main	this PR
narrow (500 of 2M)	linear	228 ms	161 ms
narrow (500 of 2M)	nearest	231 ms	241 ms
wide (all of x)	linear	280 ms	199 ms
wide (all of x)	nearest	325 ms	245 ms

If the new coordinates are in memory, localizing once before apply_ufunc is better, because the lazy isel also drops the chunks of var that aren't needed.

Or was there another reason to localize per chunk that I'm missing?

[dcherian]

yes i believe its about controlling the memory used by scipy. Can you check memray

[headtr1ck]

I checked with memray, localizing per chunk doesn't reduce memory either. ScipyInterpolator passes copy=False and scipy's reshape of the data was a view in all layouts I tried, including the interpolated dimension in the middle ((y, x, w)), so scipy never copies the block. For linear/nearest, scipy only allocates arrays of size (new points × other dims), independent of the length of x. The peak is dask concatenating blocks, same on both branches:

data (dask, chunked new coords)	method	main	this PR
(y, x) 64 × 2M, 16 blocks, sync scheduler	linear	282 MB	273 MB
	nearest	282 MB	298 MB
(y, x, w) 16 × 200k × 10, 4 blocks	linear	212 MB	212 MB
	nearest	209 MB	210 MB

(narrow new coordinates covering 500 points of x; covering all of x gives the same numbers)

Do you have a link to an older issue where the memory problem came up? Then I can test that case specifically.

[headtr1ck]

Found a case where it matters: integer data. scipy casts y to float64 (interp1d.__init__, and RegularGridInterpolator._check_values for n-D), so without localizing it copies the whole block instead of the small slice that is needed. Float data isn't affected, which is why my earlier test showed no difference.

memray, synchronous scheduler, chunked new coordinates covering a narrow range, peak in MiB (one block is 61 MiB for 1-D and 31 MiB for n-D):

data	method	#4155	main	PR before dda83a0	PR after dda83a0
1-D float (y, x) 64 × 2M	linear / nearest	157 / 157	189 / 190	174 / 189	189 / 189
1-D int	linear / nearest	157 / 157	189 / 189	235 / 250	189 / 189
n-D pointwise float (t, y, x) 16 × 2000²	linear / nearest	36 / 36	82 / 82	82 / 82	83 / 82
n-D pointwise int	linear / nearest	36 / 36	82 / 82	113 / 113	82 / 82

The #4155 column (where per-chunk localizing was added, run with 2020 dependencies) is lower because until #9881 _localize computed chunked new coordinates eagerly (new_x.values) and sliced the dask array before interpolating. Since #9881 skips chunked new coordinates there, per-chunk localizing is the only localization left for them.

So I restored per-chunk localizing in _interpnd when the new coordinates are chunked, and kept localizing once before apply_ufunc when they are in memory (dda83a0). With that, this PR matches main in all cases (last column). I also added a peakmem asv benchmark for this case, which goes from ~593 MiB to ~700 MiB for int data without the fix.

[dcherian]

Ok that's not very important then. Perhaps some upstream work has improved the situation

[dcherian]

Didn't review _vectorize_dims too closely but in general this seems ok. left some comments to be addressed.

Lets wait a couple of days to see if Stephan has any input. If not, lets merge.