From ba2b63dff192ba9fe7ee8f533276c3d93b201edb Mon Sep 17 00:00:00 2001 From: oshokothari07-ai Date: Fri, 31 Jul 2026 02:00:55 +0530 Subject: [PATCH] Fix interp performance regression from #9881 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 #10683 --- doc/whats-new.rst | 12 +++++++ xarray/core/missing.py | 20 ++++++++++- xarray/tests/test_missing.py | 66 +++++++++++++++++++++++++++++++++++- 3 files changed, 96 insertions(+), 2 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 63177c82cca..ab5d4320287 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -35,6 +35,18 @@ Deprecations ~~~~~~~~~~~~ +Performance +~~~~~~~~~~~ + +- Fix performance regression in :py:meth:`DataArray.interp` and + :py:meth:`Dataset.interp` introduced in :pull:`9881`. When the target coordinate + varied along a dimension of the interpolated object, ``np.vectorize`` looped in + Python over every remaining non-core dimension as well, so interpolating an + array with dimensions ``(t, r, z)`` along ``z`` at targets varying with ``t`` + made ``t * r`` calls into the interpolator instead of ``t`` (:issue:`10683`). + By `Osho Kothari `_. + + Bug Fixes ~~~~~~~~~ diff --git a/xarray/core/missing.py b/xarray/core/missing.py index 7f22d2df45c..97f97ec20fe 100644 --- a/xarray/core/missing.py +++ b/xarray/core/missing.py @@ -720,11 +720,28 @@ def interpolate_variable( # so it is in exclude_dims. vectorize_dims = (result_dims - all_in_core_dims) & set(var.dims) + # When vectorizing, apply_ufunc uses np.vectorize, which loops in Python over + # *every* non-core dimension, not just the ones that actually need vectorizing. + # In the example above, `q` is a core dimension and `lat`, `lon` are vectorized, + # but any remaining dimension of `var` is also peeled off one element at a time + # purely because it is not a core dimension. For + # `da[t, r, z].interp(z=target[t])` that means t*r calls instead of t (GH10683). + # Declaring those dimensions core hands them to the interpolator in bulk; + # _interpnd already supports leading "constant" dimensions. + # + # Only do this for in-memory arrays: making a dimension a core dimension forces + # dask to rechunk along it, which can blow up memory for chunked inputs. + bulk_dims: tuple[Hashable, ...] = () + if vectorize_dims and not is_chunked_array(var._data): + bulk_dims = tuple( + d for d in var.dims if d not in all_in_core_dims and d not in vectorize_dims + ) + # remove any output broadcast dimensions from the list of core dimensions output_core_dims = tuple(d for d in result_dims if d not in vectorize_dims) input_core_dims = ( # all coordinates on the input that we interpolate along - [tuple(indexes_coords)] + [bulk_dims + tuple(indexes_coords)] # the input coordinates are always 1D at the moment, so we just need to list out their names + [tuple(_.dims) for _ in in_coords] # The last set of inputs are the coordinates we are interpolating to. @@ -734,6 +751,7 @@ def interpolate_variable( ] ) output_sizes = {k: result_sizes[k] for k in output_core_dims} + output_core_dims = bulk_dims + output_core_dims # scipy.interpolate.interp1d always forces to float. dtype = float if not issubclass(var.dtype.type, np.inexact) else var.dtype diff --git a/xarray/tests/test_missing.py b/xarray/tests/test_missing.py index 9c2d6b47727..821b0be0889 100644 --- a/xarray/tests/test_missing.py +++ b/xarray/tests/test_missing.py @@ -9,7 +9,7 @@ import pytest import xarray as xr -from xarray.core import indexing +from xarray.core import indexing, missing from xarray.core.missing import ( NumpyInterpolator, ScipyInterpolator, @@ -799,3 +799,67 @@ def wrapper(self, indexer): ds["sigma_a"].interp(w=15000.5) actual_indexer = mock_func.mock_calls[0].args[1]._key assert actual_indexer == (slice(None), slice(None), slice(18404, 18408)) + + +def _interp_vectorize_example(): + # da[t, r, z] interpolated along `z` at targets that vary with `t`. + # `t` is a genuine vectorize dimension; `r` is a free dimension. + n_t, n_r, n_z = 4, 3, 5 + rng = np.random.default_rng(seed=0) + da = xr.DataArray( + rng.random((n_t, n_r, n_z)), + dims=["t", "r", "z"], + coords={ + "t": np.arange(n_t), + "r": np.arange(n_r), + "z": np.linspace(0.0, 1.0, n_z), + }, + ) + target = xr.DataArray(np.linspace(0.1, 0.9, n_t), dims=["t"], coords={"t": da["t"]}) + return da, target + + +@requires_scipy +def test_interp_vectorize_does_not_loop_over_free_dims(): + # regression test for GH10683 + # `da[t, r, z].interp(z=target[t])` must vectorize over `t` only. `r` is neither + # interpolated along nor varying with the target, so it should be handed to the + # interpolator in bulk rather than peeled off one element at a time by + # np.vectorize, which made this O(t * r) Python calls instead of O(t). + from scipy.interpolate import interp1d + + da, target = _interp_vectorize_example() + n_t = da.sizes["t"] + + with mock.patch.object( + missing, "_interpnd", side_effect=missing._interpnd, autospec=True + ) as mock_interpnd: + actual = da.interp(z=target) + + assert mock_interpnd.call_count == n_t + + # independent reference: interpolate each `t` separately with scipy + expected = np.stack( + [ + interp1d(da["z"].values, da.values[i], axis=-1)(target.values[i]) + for i in range(n_t) + ] + ) + assert actual.dims == ("t", "r") + np.testing.assert_allclose(actual.values, expected) + + +@requires_scipy +@requires_dask +def test_interp_vectorize_preserves_chunks_along_free_dims(): + # companion to GH10683: the bulk-dimension optimization must not apply to chunked + # arrays. Declaring `r` a core dimension would force dask to rechunk along it, + # so a preserved chunk structure is what shows the optimization was skipped. + da, target = _interp_vectorize_example() + + expected = da.interp(z=target) + actual = da.chunk({"r": 1}).interp(z=target) + + assert actual.chunks is not None + assert actual.chunks[actual.dims.index("r")] == (1, 1, 1) + assert_allclose(actual.compute(), expected)