Fast and exact multidimensional cubic/quintic interp

asv_bench/benchmarks/interp.py

@@ -1,3 +1,5 @@
+from typing import Any
+
 import numpy as np
 import pandas as pd
 
@@ -18,6 +20,51 @@
 new_y_long = np.linspace(0.1, 0.9, 500)
 
 
+class InterpolationSplineND:
+    # spline interpolation over multiple dimensions (GH11342)
+    def setup(self, *args: Any, **kwargs: Any) -> None:
+        nt, ny, nx = 50, 40, 50
+        self.da = xr.DataArray(
+            randn((nt, ny, nx)),
+            dims=("t", "y", "x"),
+            coords={
+                "t": np.arange(nt),
+                "y": np.linspace(0, 1, ny),
+                "x": np.linspace(0, 1, nx),
+            },
+        )
+        self.destinations = {
+            "point": {"t": 25.5, "y": 0.45, "x": 0.12},
+            "grid": {
+                "t": np.linspace(5, 45, 20),
+                "y": np.linspace(0.1, 0.9, 20),
+                "x": np.linspace(0.1, 0.9, 20),
+            },
+            "pointwise": {
+                dim: xr.DataArray(np.linspace(lo, hi, 100), dims="p")
+                for dim, lo, hi in [("t", 5, 45), ("y", 0.1, 0.9), ("x", 0.9, 0.1)]
+            },
+        }
+
+    @parameterized(
+        ["method", "destination"],
+        (["cubic", "quintic"], ["point", "grid", "pointwise"]),
+    )
+    def time_interpolation_spline_nd(self, method: str, destination: str) -> None:
+        self.da.interp(self.destinations[destination], method=method).compute()
+
+    @parameterized(["method"], (["cubic", "quintic"]))
+    def peakmem_interpolation_spline_nd(self, method: str) -> None:
+        self.da.interp(self.destinations["point"], method=method).compute()
+
+
+class InterpolationSplineNDDask(InterpolationSplineND):
+    def setup(self, *args: Any, **kwargs: Any) -> None:
+        requires_dask()
+        super().setup(**kwargs)
+        self.da = self.da.chunk({"t": 10})
+
+
 class Interpolation:
     def setup(self, *args, **kwargs):
         self.ds = xr.Dataset(

doc/whats-new.rst

@@ -189,6 +189,14 @@ Documentation
 
 Performance
 ~~~~~~~~~~~
+- :py:meth:`Dataset.interp` and :py:meth:`DataArray.interp` with ``method="cubic"``
+  or ``"quintic"`` over multiple dimensions are much faster and need much less
+  memory, e.g. 15x faster and 25x less memory for a single point in a
+  ``(288, 71, 73)`` array. The results are now exact instead of only accurate to
+  about ``1e-6``: the spline coefficients are solved for one dimension after
+  the other instead of with an iterative solver for all dimensions at once
+  (:issue:`11342`, :pull:`11708`).
+  By `Michael Niklas <https://github.com/headtr1ck>`_.
 - Speed up the ``repr`` of objects with many coordinates by about 2.5x, which
   got about twice as slow when the coordinates were ordered by dimension in
   v2025.10.0 (:pull:`10778`), and their HTML repr by about 15% (:pull:`11691`): the

xarray/core/missing.py

@@ -7,7 +7,7 @@
 from collections.abc import Callable, Generator, Hashable, Sequence
 from functools import partial
 from numbers import Number
-from typing import TYPE_CHECKING, Any, TypeVar, get_args, override
+from typing import TYPE_CHECKING, Any, Literal, TypeVar, get_args, override
 
 import numpy as np
 import pandas as pd
@@ -34,6 +34,8 @@
 from xarray.namedarray.pycompat import is_chunked_array
 
 if TYPE_CHECKING:
+    from numpy.typing import ArrayLike
+
     from xarray.core.dataarray import DataArray
     from xarray.core.dataset import Dataset
 
@@ -544,7 +546,55 @@ def _get_interpolator(
     return interp_class, kwargs
 
 
-def _get_interpolator_nd(method, **kwargs):
+_SPLINE_DEGREES = {"cubic": 3, "quintic": 5}
+
+
+def _interpn_spline(
+    points: Sequence[ArrayLike],
+    values: ArrayLike,
+    xi: ArrayLike,
+    *,
+    method: Literal["cubic", "quintic"],
+    bounds_error: bool = False,
+    fill_value: float | None = np.nan,
+) -> np.ndarray:
+    """Tensor-product spline interpolation on a grid, like scipy.interpolate.interpn.
+
+    For spline methods, scipy solves for the coefficients along all dimensions at
+    once with an iterative sparse solver, which is slow, needs a lot of memory and
+    is only accurate to ~1e-6 (GH11342). The coefficients of a tensor-product
+    spline can be solved for one dimension after the other instead, which is exact.
+    """
+    from scipy.interpolate import NdBSpline, make_interp_spline
+
+    k = _SPLINE_DEGREES[method]
+    grid = [np.asarray(x, dtype=np.float64) for x in points]
+    xi = np.asarray(xi, dtype=np.float64)
+
+    knots = []
+    coeffs = np.asarray(values)
+    for axis, x in enumerate(grid):
+        spline = make_interp_spline(x, coeffs, k=k, axis=axis)
+        knots.append(spline.t)
+        coeffs = np.moveaxis(spline.c, 0, axis)
+
+    nans = np.isnan(xi).any(axis=-1)
+    out_of_bounds = np.zeros(xi.shape[:-1], dtype=bool)
+    for i, x in enumerate(grid):
+        out_of_bounds |= (xi[..., i] < x[0]) | (xi[..., i] > x[-1])
+    if bounds_error and out_of_bounds.any():
+        raise ValueError("One of the requested xi is out of bounds")
+
+    result = NdBSpline(tuple(knots), coeffs, k)(xi)
+    if fill_value is not None:
+        result[out_of_bounds] = fill_value
+    result[nans] = np.nan
+    return result
+
+
+def _get_interpolator_nd(
+    method: InterpOptions, **kwargs: Any
+) -> tuple[InterpCallable, dict[str, Any]]:
     """helper function to select the appropriate interpolator class
 
     returns interpolator class and keyword arguments for the class
@@ -553,7 +603,11 @@ def _get_interpolator_nd(method, **kwargs):
     if method in valid_methods:
         kwargs.update(method=method)
         kwargs.setdefault("bounds_error", False)
-        interp_class = _import_interpolant("interpn", method)
+        if method in _SPLINE_DEGREES:
+            _import_interpolant("NdBSpline", method)
+            interp_class = _interpn_spline
+        else:
+            interp_class = _import_interpolant("interpn", method)
     else:
         raise ValueError(
             f"{method} is not a valid interpolator for interpolating "

xarray/tests/test_interp.py

@@ -2,11 +2,12 @@
 
 import contextlib
 from itertools import combinations, permutations, product
-from typing import cast, get_args
+from typing import Any, cast, get_args
 
 import numpy as np
 import pandas as pd
 import pytest
+from numpy.typing import ArrayLike
 
 import xarray as xr
 from xarray.coding.cftimeindex import _parse_array_of_cftime_strings
@@ -284,6 +285,36 @@ def func(obj, dim, new_x, method):
     assert_allclose(actual, expected.transpose("z", "w", "y", transpose_coords=True))
 
 
+def _interpn_reference(
+    points: tuple[ArrayLike, ...],
+    values: ArrayLike,
+    xi: ArrayLike,
+    method: InterpnOptions,
+) -> np.ndarray:
+    """scipy.interpolate.interpn, but exact for spline methods.
+
+    For spline methods, interpn solves for the coefficients iteratively with a loose
+    default tolerance, which is only accurate to ~1e-5 for quintic. A direct solver
+    (spsolve) is exact but takes minutes on these grids, so tighten the tolerance
+    of the iterative solver instead.
+    """
+    from scipy.sparse.linalg import gcrotmk
+
+    # scipy-stubs types solver as returning an array and solver_args as a tuple,
+    # but scipy calls iterative solvers as ``x, info = solver(A, b, **solver_args)``
+    solver_kwargs: dict[str, Any] = {}
+    if method in ("cubic", "quintic"):
+        solver_kwargs["solver"] = gcrotmk
+        solver_kwargs["solver_args"] = {"rtol": 1e-12, "atol": 1e-14}
+    return scipy.interpolate.RegularGridInterpolator(
+        [np.asarray(x, dtype=np.float64) for x in points],
+        np.asarray(values, dtype=np.float64),
+        method=method,
+        bounds_error=False,
+        **solver_kwargs,
+    )(np.asarray(xi, dtype=np.float64))
+
+
 @requires_scipy
 @pytest.mark.parametrize("method", get_args(InterpnOptions))
 @pytest.mark.parametrize(
@@ -305,12 +336,8 @@ def test_interpolate_nd(case: int, method: InterpnOptions, nd_interp_coords) ->
     # coordinate is unchanged
     # TODO: test this?
     actual = da.interp(x=xdestnp, y=ydestnp, z=zdestnp, method=method)
-    expected_data = scipy.interpolate.interpn(
-        points=(da.x, da.y, da.z),
-        values=da.load().data,
-        xi=grid_grid_points,
-        method=method,
-        bounds_error=False,
+    expected_data = _interpn_reference(
+        (da.x, da.y, da.z), da.load().data, grid_grid_points, method
     ).reshape((len(xdestnp), len(ydestnp), len(zdestnp)))
     expected = xr.DataArray(
         expected_data,
@@ -330,12 +357,8 @@ def test_interpolate_nd(case: int, method: InterpnOptions, nd_interp_coords) ->
     zdest = nd_interp_coords["zdest"]
     grid_oned_points = nd_interp_coords["grid_oned_points"]
     actual = da.interp(x=xdest, y=ydest, z=zdest, method=method)
-    expected_data_1d: np.ndarray = scipy.interpolate.interpn(
-        points=(da.x, da.y, da.z),
-        values=da.data,
-        xi=grid_oned_points,
-        method=method,
-        bounds_error=False,
+    expected_data_1d: np.ndarray = _interpn_reference(
+        (da.x, da.y, da.z), da.data, grid_oned_points, method
     ).reshape([len(xdest), len(zdest)])
     expected = xr.DataArray(
         expected_data_1d,
@@ -485,6 +508,60 @@ def test_interpolate_nd_scalar(method: InterpOptions, case: int) -> None:
     assert_allclose(actual, expected)
 
 
+@requires_scipy
+@pytest.mark.parametrize("method", ["cubic", "quintic"])
+@pytest.mark.parametrize("fill_value", [np.nan, None, -5.0])
+def test_interpolate_nd_spline_exact(
+    method: InterpnOptions, fill_value: float | None, use_dask: bool
+) -> None:
+    # GH11342: scipy solves for the spline coefficients iteratively, which is slow
+    # and inexact, so compare against scipy with a direct solver instead
+    from scipy.sparse.linalg import spsolve
+
+    rng = np.random.default_rng(0)
+    x = np.sort(rng.uniform(0, 10, 12))
+    y = np.linspace(-1, 1, 9)
+    z = np.arange(8.0)
+    da = xr.DataArray(
+        rng.random((12, 9, 8, 3)),
+        dims=("x", "y", "z", "w"),
+        coords={"x": x, "y": y, "z": z},
+    )
+    if use_dask:
+        da = da.chunk(x=4, w=1)
+    xdest = xr.DataArray(rng.uniform(-1, 11, 20), dims="p")
+    ydest = xr.DataArray(rng.uniform(-1.2, 1.2, 20), dims="p")
+    zdest = xr.DataArray(rng.uniform(0, 7, 20), dims="p")
+    xdest[0] = np.nan
+
+    actual = da.interp(
+        x=xdest, y=ydest, z=zdest, method=method, kwargs={"fill_value": fill_value}
+    )
+    expected = scipy.interpolate.RegularGridInterpolator(
+        (x, y, z),
+        da.values,
+        method=method,
+        bounds_error=False,
+        fill_value=fill_value,
+        solver=spsolve,
+    )(np.stack([xdest, ydest, zdest], axis=-1))
+    np.testing.assert_allclose(actual.transpose("p", "w").values, expected, rtol=1e-9)
+
+    # orthogonal destinations
+    actual = da.interp(x=[2.5, 7.1], y=0.3, z=[1.2, 5.5, 6.9], method=method)
+    expected = scipy.interpolate.RegularGridInterpolator(
+        (x, y, z), da.values, method=method, bounds_error=False, solver=spsolve
+    )(np.stack(np.meshgrid([2.5, 7.1], 0.3, [1.2, 5.5, 6.9], indexing="ij"), axis=-1))
+    np.testing.assert_allclose(
+        actual.transpose("x", "z", "w").values, expected.squeeze(1), rtol=1e-9
+    )
+
+    with pytest.raises(ValueError, match="out of bounds"):
+        da.interp(
+            x=11, y=0, z=0, method=method, kwargs={"bounds_error": True}
+        ).compute()
+
+
 @requires_scipy
 def test_nans(use_dask: bool) -> None:
     da = xr.DataArray([0, 1, np.nan, 2], dims="x", coords={"x": range(4)})