diff --git a/asv_bench/benchmarks/interp.py b/asv_bench/benchmarks/interp.py index ca1d0a2dd89..f728365c2e5 100644 --- a/asv_bench/benchmarks/interp.py +++ b/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( diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 21bbb428600..eb400460c47 100644 --- a/doc/whats-new.rst +++ b/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 `_. - 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 diff --git a/xarray/core/missing.py b/xarray/core/missing.py index ebf93bb0e28..f3e1d09a713 100644 --- a/xarray/core/missing.py +++ b/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 " diff --git a/xarray/tests/test_interp.py b/xarray/tests/test_interp.py index f779c6fe396..5de676e15fa 100644 --- a/xarray/tests/test_interp.py +++ b/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)})