Fix interp performance regression from #9881

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 <https://github.com/oshokothari07-ai>`_.
+
+
 Bug Fixes
 ~~~~~~~~~
 

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

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)