From db74ab3f7d35018cbe7a761f1ab6bc9f7bd3155a Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Tue, 6 Oct 2026 21:57:29 +0200 Subject: [PATCH 01/18] Fix interp performance regression for vectorized interpolation np.vectorize looped over all non-core dimensions, although the new coordinates only vary along some of them (GH10683). _interpnd now loops over only those dimensions itself and passes all others to the interpolator in bulk, for numpy and dask arrays alike. The localization and conversion of the coordinates now happens once per call instead of once per loop iteration. Also fix interp with empty new coordinates for the linear and nearest methods. Co-authored-by: Claude Co-Authored-By: Claude Opus 5.5 --- asv_bench/benchmarks/interp.py | 27 +++++++++ doc/whats-new.rst | 11 ++++ xarray/core/missing.py | 100 ++++++++++++++++++++++++++------- xarray/tests/test_interp.py | 28 +++++++++ 4 files changed, 147 insertions(+), 19 deletions(-) diff --git a/asv_bench/benchmarks/interp.py b/asv_bench/benchmarks/interp.py index ca1d0a2dd89..d995d565f79 100644 --- a/asv_bench/benchmarks/interp.py +++ b/asv_bench/benchmarks/interp.py @@ -63,3 +63,30 @@ def setup(self, *args, **kwargs): requires_dask() super().setup(**kwargs) self.ds = self.ds.chunk({"t": 50}) + + +class InterpolationVectorized: + # the destination coordinates vary along a dimension of the source (GH10683) + def setup(self, *args, **kwargs): + nt, nr, nz = 100, 50, 50 + self.da = xr.DataArray( + randn((nt, nr, nz)), + dims=("t", "r", "z"), + coords={ + "t": np.arange(nt), + "r": np.arange(nr), + "z": np.linspace(0, 1, nz), + }, + ) + self.new_z = xr.DataArray(np.linspace(0.1, 0.9, nt), dims="t") + + @parameterized(["method"], (["linear", "cubic"])) + def time_interpolation_vectorized(self, method): + self.da.interp(z=self.new_z, method=method).compute() + + +class InterpolationVectorizedDask(InterpolationVectorized): + def setup(self, *args, **kwargs): + requires_dask() + super().setup(**kwargs) + self.da = self.da.chunk({"t": 25, "r": 25}) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 25649d3e427..3afe6ee4ce5 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -40,6 +40,10 @@ Deprecations Bug Fixes ~~~~~~~~~ +- Fix :py:meth:`Dataset.interp` and :py:meth:`DataArray.interp` with + ``method="linear"`` or ``"nearest"`` raising a ``ValueError`` for empty new + coordinates (:pull:`PRNUM`). + By `Michael Niklas `_. - Fix ``InvalidIndexError`` in :py:meth:`DataArrayGroupBy.median` and other groupby reductions using flox's ``method="blockwise"`` on dask arrays, when members of a group are spread over multiple chunks and are not contiguous. @@ -100,6 +104,13 @@ Performance ones instead of sorting them. This avoids copying the data before interpolating (:issue:`9758`, :pull:`11658`). By `Bhaskar Gurram `_. +- Fix a performance regression in :py:meth:`Dataset.interp` and + :py:meth:`DataArray.interp` when the new coordinates vary along a dimension + of the interpolated object, e.g. ``da.interp(z=new_z)`` where ``new_z`` has a + dimension ``t`` of ``da``. The interpolation now only loops over the + dimensions the new coordinates depend on, for numpy and dask arrays alike + (:issue:`10683`, :pull:`PRNUM`). + By `Michael Niklas `_. Internal Changes diff --git a/xarray/core/missing.py b/xarray/core/missing.py index 02e9e4edd6a..249fb0bfd3f 100644 --- a/xarray/core/missing.py +++ b/xarray/core/missing.py @@ -583,7 +583,7 @@ def _localize[T](obj: T, indexes_coords: SourceDest) -> tuple[T, SourceDest]: """ indexes = {} for dim, [x, new_x] in indexes_coords.items(): - if is_chunked_array(new_x._data): + if is_chunked_array(new_x._data) or new_x.size == 0: continue new_x_loaded = new_x.data minval = np.nanmin(new_x_loaded) @@ -760,7 +760,6 @@ def interpolate_variable( # TODO: deprecate and have the user rechunk themselves dask_gufunc_kwargs=dict(output_sizes=output_sizes, allow_rechunk=True), output_dtypes=[dtype], - vectorize=bool(vectorize_dims), keep_attrs=True, ) return result @@ -795,26 +794,30 @@ def _interpnd( Core nD array interpolation routine. The first half arrays in `coords` are original coordinates, the other half are destination coordinates. + + `data` has the loop dimensions first, followed by the dimensions we + interpolate along. The destination coordinates may vary along some of the + loop dimensions (following numpy broadcasting rules), e.g. for + ``da[t, r, z].interp(z=target[t])``. We loop over only those dimensions, + all other loop dimensions are passed to the interpolator in bulk. """ n_x = len(coords) // 2 - ndim = data.ndim - nconst = ndim - n_x + nloop = data.ndim - n_x + # names that do not clash with the dimension names of the destination coordinates + loop_dims = tuple(f"__interpnd_loop_dim_{i}__" for i in range(nloop)) + core_dims = tuple(f"__interpnd_core_dim_{i}__" for i in range(n_x)) # Convert everything to Variables, since that makes applying # `_localize` and `_floatize_x` much easier + var = Variable(loop_dims + core_dims, data, fastpath=True) x = [ - Variable([f"dim_{nconst + dim}"], _x, fastpath=True) - for dim, _x in enumerate(coords[:n_x]) + Variable([dim], _x, fastpath=True) + for dim, _x in zip(core_dims, coords[:n_x], strict=True) + ] + new_x = [ + Variable(loop_dims[nloop - (_x.ndim - len(dims)) :] + dims, _x, fastpath=True) + for dims, _x in zip(result_coord_core_dims, coords[n_x:], strict=True) ] - new_x = list( - broadcast_variables( - *( - Variable(dims, _x, fastpath=True) - for dims, _x in zip(result_coord_core_dims, coords[n_x:], strict=True) - ) - ) - ) - var = Variable([f"dim_{dim}" for dim in range(ndim)], data, fastpath=True) if interp_kwargs.get("method") in ["linear", "nearest"]: indexes_coords = { @@ -827,18 +830,77 @@ def _interpnd( for _ in zip(*(indexes_coords[d] for d in indexes_coords), strict=True) ) - x_list, new_x_list = _floatize_x(x, new_x) + x, new_x = _floatize_x(x, new_x) + + def drop_loop_dims(v: Variable) -> Variable: + # the remaining loop dimensions all have size 1 on the destination coordinates + return v.squeeze([dim for dim in loop_dims if dim in v.dims]) + + # the loop dimensions along which the destination coordinates vary + vectorize_dims = [ + dim + for dim in loop_dims + if any(_new_x.sizes.get(dim, 1) != 1 for _new_x in new_x) + ] + if not vectorize_dims: + return _interpnd_core( + var, + x, + [drop_loop_dims(_new_x) for _new_x in new_x], + interp_func, + interp_kwargs, + ) + + # Loop over the vectorized dimensions only. This is much faster than + # np.vectorize, which would loop over all loop dimensions (GH10683). + result_sizes: dict[Hashable, int] = {} + for _new_x in new_x: + for dim, size in _new_x.sizes.items(): + if dim not in loop_dims: + result_sizes.setdefault(dim, size) + # scipy.interpolate.interp1d always forces to float. + dtype = float if not issubclass(data.dtype.type, np.inexact) else data.dtype + result = np.empty(var.shape[:nloop] + tuple(result_sizes.values()), dtype=dtype) + + vectorize_axes = [loop_dims.index(dim) for dim in vectorize_dims] + key: list[int | slice] = [slice(None)] * nloop + for index in np.ndindex(*(var.sizes[dim] for dim in vectorize_dims)): + indexers = dict(zip(vectorize_dims, index, strict=True)) + for axis, i in zip(vectorize_axes, index, strict=True): + key[axis] = i + result[tuple(key)] = _interpnd_core( + var.isel(indexers), + x, + [ + drop_loop_dims(_new_x.isel(indexers, missing_dims="ignore")) + for _new_x in new_x + ], + interp_func, + interp_kwargs, + ) + return result + + +def _interpnd_core( + var: Variable, + x: list[Variable], + new_x: list[Variable], + interp_func: Interpolator | InterpCallable, + interp_kwargs, +) -> np.ndarray: + """Interpolate the last ``len(x)`` dimensions of `var` at the points `new_x`.""" + new_x = list(broadcast_variables(*new_x)) if len(x) == 1: # TODO: narrow interp_func to interpolator here - return _interp1d(var, x_list, new_x_list, interp_func, interp_kwargs) # type: ignore[arg-type] + return _interp1d(var, x, new_x, interp_func, interp_kwargs) # type: ignore[arg-type] # move the interpolation axes to the start position data = transpose(var._data, range(-len(x), var.ndim - len(x))) # stack new_x to 1 vector, with reshape - xi = stack([ravel(x1.data) for x1 in new_x_list], axis=-1) - rslt: np.ndarray = interp_func(x_list, data, xi, **interp_kwargs) # type: ignore[assignment] + xi = stack([ravel(x1.data) for x1 in new_x], axis=-1) + rslt: np.ndarray = interp_func(x, data, xi, **interp_kwargs) # type: ignore[assignment] # move back the interpolation axes to the last position rslt = transpose(rslt, range(-rslt.ndim + 1, 1)) return reshape(rslt, rslt.shape[:-1] + new_x[0].shape) diff --git a/xarray/tests/test_interp.py b/xarray/tests/test_interp.py index f779c6fe396..e4e2a02deb7 100644 --- a/xarray/tests/test_interp.py +++ b/xarray/tests/test_interp.py @@ -1300,3 +1300,31 @@ def test_dataset_interp_datetime_dask() -> None: computed = result.compute() expected_time = np.datetime64("2024-01-01") + np.timedelta64(3, "D") np.testing.assert_equal(computed["time"].values[0, 0], expected_time) + + +@requires_scipy +@pytest.mark.parametrize("method", ["linear", "nearest", "cubic"]) +def test_interp_empty_new_coords(method) -> None: + da = xr.DataArray(np.arange(5.0), dims="x", coords={"x": np.arange(5)}) + actual = da.interp(x=np.array([], dtype=float), method=method) + expected = xr.DataArray( + np.array([], dtype=float), dims="x", coords={"x": np.array([], dtype=float)} + ) + assert_identical(actual, expected) + + +@requires_scipy +@pytest.mark.parametrize("chunk", [False, pytest.param(True, marks=requires_dask)]) +def test_interp_vectorized_empty_dim(chunk: bool) -> None: + # the new coordinates vary along the empty dimension "t" + da = xr.DataArray( + np.zeros((0, 3, 5)), + dims=("t", "r", "z"), + coords={"z": np.linspace(0, 1, 5)}, + ) + if chunk: + da = da.chunk() + new_z = xr.DataArray(np.array([], dtype=float), dims="t") + actual = da.interp(z=new_z) + expected = xr.DataArray(np.zeros((0, 3)), dims=("t", "r"), coords={"z": new_z}) + assert_identical(actual.compute(), expected) From ef9d63de19bfb45a990e1c434d8284c897038fc9 Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Tue, 6 Oct 2026 21:57:48 +0200 Subject: [PATCH 02/18] Add PR number to whats-new Co-authored-by: Claude Co-Authored-By: Claude Opus 5.5 --- doc/whats-new.rst | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 3afe6ee4ce5..fc9a4d6c1c0 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -42,7 +42,7 @@ Bug Fixes ~~~~~~~~~ - Fix :py:meth:`Dataset.interp` and :py:meth:`DataArray.interp` with ``method="linear"`` or ``"nearest"`` raising a ``ValueError`` for empty new - coordinates (:pull:`PRNUM`). + coordinates (:pull:`11681`). By `Michael Niklas `_. - Fix ``InvalidIndexError`` in :py:meth:`DataArrayGroupBy.median` and other groupby reductions using flox's ``method="blockwise"`` on dask arrays, when @@ -109,7 +109,7 @@ Performance of the interpolated object, e.g. ``da.interp(z=new_z)`` where ``new_z`` has a dimension ``t`` of ``da``. The interpolation now only loops over the dimensions the new coordinates depend on, for numpy and dask arrays alike - (:issue:`10683`, :pull:`PRNUM`). + (:issue:`10683`, :pull:`11681`). By `Michael Niklas `_. From 43c38ff7cd5d00b3c87413d35f49a52a4906ca4f Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Tue, 6 Oct 2026 21:58:13 +0200 Subject: [PATCH 03/18] Use the use_dask fixture in test_interp_vectorized_empty_dim Co-authored-by: Claude Co-Authored-By: Claude Opus 5.5 --- xarray/tests/test_interp.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/xarray/tests/test_interp.py b/xarray/tests/test_interp.py index e4e2a02deb7..2bb8f4e97ee 100644 --- a/xarray/tests/test_interp.py +++ b/xarray/tests/test_interp.py @@ -1314,15 +1314,14 @@ def test_interp_empty_new_coords(method) -> None: @requires_scipy -@pytest.mark.parametrize("chunk", [False, pytest.param(True, marks=requires_dask)]) -def test_interp_vectorized_empty_dim(chunk: bool) -> None: +def test_interp_vectorized_empty_dim(use_dask: bool) -> None: # the new coordinates vary along the empty dimension "t" da = xr.DataArray( np.zeros((0, 3, 5)), dims=("t", "r", "z"), coords={"z": np.linspace(0, 1, 5)}, ) - if chunk: + if use_dask: da = da.chunk() new_z = xr.DataArray(np.array([], dtype=float), dims="t") actual = da.interp(z=new_z) From c22092466353bef8c0d6378e9a51b5114468dd53 Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Tue, 6 Oct 2026 22:15:02 +0200 Subject: [PATCH 04/18] Fix interp with scalar new coordinates on numpy<2 Co-authored-by: Claude Co-Authored-By: Claude Opus 5.5 --- xarray/core/missing.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/xarray/core/missing.py b/xarray/core/missing.py index 249fb0bfd3f..29b71001f90 100644 --- a/xarray/core/missing.py +++ b/xarray/core/missing.py @@ -4,7 +4,7 @@ import itertools import warnings from collections import ChainMap -from collections.abc import Callable, Generator, Hashable, Sequence +from collections.abc import Callable, Generator, Hashable, Mapping, Sequence from functools import partial from numbers import Number from typing import TYPE_CHECKING, Any, TypeVar, get_args, override @@ -832,9 +832,12 @@ def _interpnd( x, new_x = _floatize_x(x, new_x) - def drop_loop_dims(v: Variable) -> Variable: - # the remaining loop dimensions all have size 1 on the destination coordinates - return v.squeeze([dim for dim in loop_dims if dim in v.dims]) + def select_loop_index(v: Variable, indexers: Mapping[Any, int]) -> Variable: + # Select `indexers` and drop all loop dimensions of a destination coordinate. + # The non-vectorized loop dimensions all have size 1. + key = {dim: indexers.get(dim, 0) for dim in loop_dims if dim in v.dims} + # no isel without key: v might wrap a numpy scalar (numpy<2) + return v.isel(key) if key else v # the loop dimensions along which the destination coordinates vary vectorize_dims = [ @@ -846,7 +849,7 @@ def drop_loop_dims(v: Variable) -> Variable: return _interpnd_core( var, x, - [drop_loop_dims(_new_x) for _new_x in new_x], + [select_loop_index(_new_x, {}) for _new_x in new_x], interp_func, interp_kwargs, ) @@ -871,10 +874,7 @@ def drop_loop_dims(v: Variable) -> Variable: result[tuple(key)] = _interpnd_core( var.isel(indexers), x, - [ - drop_loop_dims(_new_x.isel(indexers, missing_dims="ignore")) - for _new_x in new_x - ], + [select_loop_index(_new_x, indexers) for _new_x in new_x], interp_func, interp_kwargs, ) From c79b40f6694016512e778e1065950f482229c8e8 Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Wed, 7 Oct 2026 20:40:20 +0200 Subject: [PATCH 05/18] Add vectorize="shared" to apply_ufunc and use it in interp vectorize="shared" only loops over the broadcast dimensions along which more than one input varies and passes all others to func in bulk. interp now uses it instead of its own loop, and localizes and converts datetime coordinates once per call before apply_ufunc. Co-authored-by: Claude Co-Authored-By: Claude Opus 5.5 --- doc/whats-new.rst | 6 ++ xarray/computation/apply_ufunc.py | 156 +++++++++++++++++++++++++++++- xarray/core/missing.py | 124 +++++++----------------- xarray/tests/test_computation.py | 89 +++++++++++++++++ 4 files changed, 280 insertions(+), 95 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index ab28090b8da..cfb8d78646a 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -13,6 +13,12 @@ v2026.09.1 (unreleased) New Features ~~~~~~~~~~~~ +- :py:func:`apply_ufunc` supports ``vectorize="shared"``, which only loops over + the broadcast dimensions along which more than one input varies and passes + all other broadcast dimensions to ``func`` in bulk. This is much faster than + ``vectorize=True`` if ``func`` can broadcast its inputs, for numpy and dask + arrays alike (:pull:`11681`). + By `Michael Niklas `_. Breaking Changes diff --git a/xarray/computation/apply_ufunc.py b/xarray/computation/apply_ufunc.py index d0b2e1d7b79..0d5fbee88cd 100644 --- a/xarray/computation/apply_ufunc.py +++ b/xarray/computation/apply_ufunc.py @@ -704,6 +704,121 @@ def _vectorize(func, signature, output_dtypes, exclude_dims): return func +def _vectorize_shared( + func: Callable, + signature: _UFuncSignature, + output_dtypes: Sequence | None, + output_core_sizes: Mapping[Hashable, int], +) -> Callable: + """Vectorize ``func`` over the loop dimensions shared by several inputs. + + Like :py:class:`numpy.vectorize` with a signature, but ``func`` is only + looped over the loop (non-core) dimensions along which more than one input + varies, i.e. has a size other than 1. All other loop dimensions are passed + to ``func`` in bulk, so ``func`` must accept inputs with leading loop + dimensions that broadcast against each other following numpy's rules, and + return outputs with the broadcast loop dimensions followed by the output + core dimensions. + + ``output_core_sizes`` maps output core dimensions to their sizes; it is only + needed if the loop is empty, so that ``func`` is never called. + """ + input_core_ndims = [len(dims) for dims in signature.input_core_dims] + num_outputs = signature.num_outputs + + @functools.wraps(func) + def wrapper(*args): + loop_shapes = [ + np.shape(arg)[: np.ndim(arg) - ndim] + for arg, ndim in zip(args, input_core_ndims, strict=True) + ] + broadcast_shape = np.broadcast_shapes(*loop_shapes) + nloop = len(broadcast_shape) + # loop dimensions are aligned to the right, as in numpy broadcasting + offsets = [nloop - len(shape) for shape in loop_shapes] + loop_axes = [ + axis + for axis in range(nloop) + if sum( + axis >= offset and shape[axis - offset] != 1 + for shape, offset in zip(loop_shapes, offsets, strict=True) + ) + > 1 + ] + if not loop_axes: + return func(*args) + + outputs = None + output_key: list[int | slice] = [slice(None)] * nloop + for index in np.ndindex(*(broadcast_shape[axis] for axis in loop_axes)): + loop_index = dict(zip(loop_axes, index, strict=True)) + sliced_args = [] + for arg, shape, offset in zip(args, loop_shapes, offsets, strict=True): + # select the loop index, or the only element along a size-1 axis + key = tuple( + (loop_index[axis] if shape[axis - offset] != 1 else 0) + if axis in loop_index + else slice(None) + for axis in range(offset, nloop) + ) + sliced_args.append(arg[key] if key else arg) + + results = func(*sliced_args) + if num_outputs == 1: + results = (results,) + elif not isinstance(results, tuple) or len(results) != num_outputs: + raise ValueError( + f"applied function does not have the number of outputs " + f"specified in the ufunc signature. Expected {num_outputs} " + f"outputs, got {results!r}" + ) + + if outputs is None: + outputs = tuple( + np.empty_like( + result, + shape=broadcast_shape + + np.shape(result)[np.ndim(result) - len(core_dims) :], + dtype=None if output_dtypes is None else output_dtypes[i], + ) + for i, (result, core_dims) in enumerate( + zip(results, signature.output_core_dims, strict=True) + ) + ) + + for axis, i in loop_index.items(): + output_key[axis] = i + for output, result in zip(outputs, results, strict=True): + output[tuple(output_key)] = result + + if outputs is None: + # the loop is empty, so func was never called + missing_dims = { + dim + for dims in signature.output_core_dims + for dim in dims + if dim not in output_core_sizes + } + if output_dtypes is None or missing_dims: + raise ValueError( + "cannot vectorize over an empty loop dimension unless " + "``output_dtypes`` is set and the sizes of all output core " + "dimensions are known" + ) + outputs = tuple( + np.empty( + broadcast_shape + tuple(output_core_sizes[dim] for dim in dims), + dtype=dtype, + ) + for dims, dtype in zip( + signature.output_core_dims, output_dtypes, strict=True + ) + ) + return outputs[0] if num_outputs == 1 else outputs + + return wrapper + + def apply_variable_ufunc( func, *args, @@ -711,7 +826,7 @@ def apply_variable_ufunc( exclude_dims=frozenset(), dask="forbidden", output_dtypes=None, - vectorize=False, + vectorize: bool | Literal["shared"] = False, keep_attrs="override", dask_gufunc_kwargs=None, ) -> Variable | tuple[Variable, ...]: @@ -794,12 +909,20 @@ def apply_variable_ufunc( f"dimension '{key}' in 'output_core_dims' needs corresponding (dim, size) in 'output_sizes'" ) + if vectorize == "shared": + numpy_func = _vectorize_shared( + numpy_func, + signature, + output_dtypes, + output_core_sizes={**dim_sizes, **output_sizes}, + ) + def func(*arrays): res = chunkmanager.apply_gufunc( numpy_func, signature.to_gufunc_string(exclude_dims), *arrays, - vectorize=vectorize, + vectorize=vectorize is True, output_dtypes=output_dtypes, **dask_gufunc_kwargs, ) @@ -812,6 +935,14 @@ def func(*arrays): raise ValueError( f"unknown setting for chunked array handling in apply_ufunc: {dask}" ) + elif vectorize == "shared": + output_sizes = (dask_gufunc_kwargs or {}).get("output_sizes", {}) + func = _vectorize_shared( + func, + signature, + output_dtypes, + output_core_sizes={**dim_sizes, **output_sizes}, + ) elif vectorize: func = _vectorize( func, signature, output_dtypes=output_dtypes, exclude_dims=exclude_dims @@ -899,7 +1030,7 @@ def apply_ufunc( input_core_dims: Sequence[Sequence] | None = None, output_core_dims: Sequence[Sequence] | None = ((),), exclude_dims: AbstractSet = frozenset(), - vectorize: bool = False, + vectorize: bool | Literal["shared"] = False, join: JoinOptions = "exact", dataset_join: str = "exact", dataset_fill_value: object = _NO_FILL_VALUE, @@ -959,11 +1090,23 @@ def apply_ufunc( will be dropped. Each excluded dimension must also appear in ``input_core_dims`` for at least one argument. Only dimensions listed here are allowed to change size between input and output objects. - vectorize : bool, optional + vectorize : bool or "shared", optional If True, then assume ``func`` only takes arrays defined over core dimensions as input and vectorize it automatically with :py:func:`numpy.vectorize`. This option exists for convenience, but is almost always slower than supplying a pre-vectorized function. + + If "shared", only loop over the broadcast dimensions along which more + than one input varies, and pass all other broadcast dimensions to + ``func`` in bulk. ``func`` then has to accept inputs with leading + broadcast dimensions that broadcast against each other following + numpy's rules, e.g. an input of shape ``(t, r, z)`` together with an + input of shape ``(1, z)``, and return outputs with the broadcast + dimensions followed by the output core dimensions. This is much faster + than ``vectorize=True`` if some broadcast dimensions only appear on a + single input. If the loop is empty, ``output_dtypes`` and the sizes of + all output core dimensions (from the inputs or ``output_sizes`` in + ``dask_gufunc_kwargs``) are required. join : {"outer", "inner", "left", "right", "exact"}, default: "exact" Method for joining the indexes of the passed objects along each dimension, and the variables of Dataset objects with mismatched @@ -1175,6 +1318,11 @@ def apply_ufunc( signature = _UFuncSignature(input_core_dims, output_core_dims) + if vectorize not in (True, False, "shared"): + raise ValueError(f"vectorize must be a bool or 'shared', got {vectorize!r}") + if vectorize != "shared": + vectorize = bool(vectorize) + if exclude_dims: if not isinstance(exclude_dims, set): raise TypeError( diff --git a/xarray/core/missing.py b/xarray/core/missing.py index 29b71001f90..52ef4cf1be2 100644 --- a/xarray/core/missing.py +++ b/xarray/core/missing.py @@ -4,7 +4,7 @@ import itertools import warnings from collections import ChainMap -from collections.abc import Callable, Generator, Hashable, Mapping, Sequence +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 @@ -700,7 +700,13 @@ def interpolate_variable( else: func, kwargs = _get_interpolator_nd(method, **kwargs) - in_coords, result_coords = zip(*(v for v in indexes_coords.values()), strict=True) + if kwargs.get("method") in ["linear", "nearest"]: + # simple speed up for the local interpolation + var, indexes_coords = _localize(var, dict(indexes_coords)) + + in_coords, result_coords = _floatize_x( + *(list(_) for _ in zip(*indexes_coords.values(), strict=True)) + ) # input coordinates along which we are interpolation are core dimensions # the corresponding output coordinates may or may not have the same name, @@ -760,6 +766,9 @@ def interpolate_variable( # TODO: deprecate and have the user rechunk themselves dask_gufunc_kwargs=dict(output_sizes=output_sizes, allow_rechunk=True), output_dtypes=[dtype], + # only loop over the dimensions along which the destination coordinates vary, + # np.vectorize would loop over all of them (GH10683) + vectorize="shared", keep_attrs=True, ) return result @@ -796,100 +805,33 @@ def _interpnd( the other half are destination coordinates. `data` has the loop dimensions first, followed by the dimensions we - interpolate along. The destination coordinates may vary along some of the - loop dimensions (following numpy broadcasting rules), e.g. for - ``da[t, r, z].interp(z=target[t])``. We loop over only those dimensions, - all other loop dimensions are passed to the interpolator in bulk. + interpolate along. ``apply_ufunc(..., vectorize="shared")`` only loops over + the loop dimensions along which the destination coordinates vary, e.g. ``t`` + for ``da[t, r, z].interp(z=target[t])``, so the destination coordinates have + size 1 along all remaining loop dimensions. """ n_x = len(coords) // 2 - nloop = data.ndim - n_x - # names that do not clash with the dimension names of the destination coordinates - loop_dims = tuple(f"__interpnd_loop_dim_{i}__" for i in range(nloop)) - core_dims = tuple(f"__interpnd_core_dim_{i}__" for i in range(n_x)) - - # Convert everything to Variables, since that makes applying - # `_localize` and `_floatize_x` much easier - var = Variable(loop_dims + core_dims, data, fastpath=True) - x = [ - Variable([dim], _x, fastpath=True) - for dim, _x in zip(core_dims, coords[:n_x], strict=True) - ] - new_x = [ - Variable(loop_dims[nloop - (_x.ndim - len(dims)) :] + dims, _x, fastpath=True) - for dims, _x in zip(result_coord_core_dims, coords[n_x:], strict=True) - ] + ndim = data.ndim + nconst = ndim - n_x - if interp_kwargs.get("method") in ["linear", "nearest"]: - indexes_coords = { - _x.dims[0]: (_x, _new_x) for _x, _new_x in zip(x, new_x, strict=True) - } - # simple speed up for the local interpolation - var, indexes_coords = _localize(var, indexes_coords) - x, new_x = tuple( - list(_) - for _ in zip(*(indexes_coords[d] for d in indexes_coords), strict=True) - ) - - x, new_x = _floatize_x(x, new_x) - - def select_loop_index(v: Variable, indexers: Mapping[Any, int]) -> Variable: - # Select `indexers` and drop all loop dimensions of a destination coordinate. - # The non-vectorized loop dimensions all have size 1. - key = {dim: indexers.get(dim, 0) for dim in loop_dims if dim in v.dims} - # no isel without key: v might wrap a numpy scalar (numpy<2) - return v.isel(key) if key else v - - # the loop dimensions along which the destination coordinates vary - vectorize_dims = [ - dim - for dim in loop_dims - if any(_new_x.sizes.get(dim, 1) != 1 for _new_x in new_x) + # Convert everything to Variables, since that makes broadcasting the + # destination coordinates much easier + x = [ + Variable([f"dim_{nconst + dim}"], _x, fastpath=True) + for dim, _x in enumerate(coords[:n_x]) ] - if not vectorize_dims: - return _interpnd_core( - var, - x, - [select_loop_index(_new_x, {}) for _new_x in new_x], - interp_func, - interp_kwargs, - ) - - # Loop over the vectorized dimensions only. This is much faster than - # np.vectorize, which would loop over all loop dimensions (GH10683). - result_sizes: dict[Hashable, int] = {} - for _new_x in new_x: - for dim, size in _new_x.sizes.items(): - if dim not in loop_dims: - result_sizes.setdefault(dim, size) - # scipy.interpolate.interp1d always forces to float. - dtype = float if not issubclass(data.dtype.type, np.inexact) else data.dtype - result = np.empty(var.shape[:nloop] + tuple(result_sizes.values()), dtype=dtype) - - vectorize_axes = [loop_dims.index(dim) for dim in vectorize_dims] - key: list[int | slice] = [slice(None)] * nloop - for index in np.ndindex(*(var.sizes[dim] for dim in vectorize_dims)): - indexers = dict(zip(vectorize_dims, index, strict=True)) - for axis, i in zip(vectorize_axes, index, strict=True): - key[axis] = i - result[tuple(key)] = _interpnd_core( - var.isel(indexers), - x, - [select_loop_index(_new_x, indexers) for _new_x in new_x], - interp_func, - interp_kwargs, + new_x = list( + broadcast_variables( + *( + # drop the size-1 loop dimensions + Variable( + dims, reshape(_x, _x.shape[_x.ndim - len(dims) :]), fastpath=True + ) + for dims, _x in zip(result_coord_core_dims, coords[n_x:], strict=True) + ) ) - return result - - -def _interpnd_core( - var: Variable, - x: list[Variable], - new_x: list[Variable], - interp_func: Interpolator | InterpCallable, - interp_kwargs, -) -> np.ndarray: - """Interpolate the last ``len(x)`` dimensions of `var` at the points `new_x`.""" - new_x = list(broadcast_variables(*new_x)) + ) + var = Variable([f"dim_{dim}" for dim in range(ndim)], data, fastpath=True) if len(x) == 1: # TODO: narrow interp_func to interpolator here diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py index 0c5868229bd..49491546c2b 100644 --- a/xarray/tests/test_computation.py +++ b/xarray/tests/test_computation.py @@ -1645,6 +1645,95 @@ def test_vectorize_exclude_dims_dask() -> None: assert_identical(expected, actual) +def test_vectorize_shared(use_dask: bool) -> None: + rng = np.random.default_rng(0) + a = xr.DataArray(rng.random((4, 3, 5)), dims=("t", "r", "z")) + b = xr.DataArray(rng.random((4, 2, 5)), dims=("t", "s", "z")) + if use_dask: + a, b = a.chunk(t=2), b.chunk(t=2) + shapes = [] + + def dot(x: np.ndarray, y: np.ndarray) -> np.ndarray: + shapes.append((x.shape, y.shape)) + return (x * y).sum(axis=-1) + + actual = apply_ufunc( + dot, + a, + b, + input_core_dims=[["z"], ["z"]], + vectorize="shared", + dask="parallelized", + output_dtypes=[float], + ).compute() + expected = (a * b).sum("z").transpose("t", "r", "s").compute() + xr.testing.assert_allclose(actual, expected) + # only "t" is looped over, "r" and "s" are passed in bulk + assert len(shapes) == 4 + assert set(shapes) == {((3, 1, 5), (1, 2, 5))} + + +def test_vectorize_shared_without_loop() -> None: + a = xr.DataArray(np.arange(6.0).reshape(2, 3), dims=("x", "z")) + b = xr.DataArray(np.arange(3.0), dims="z") + shapes = [] + + def dot(x: np.ndarray, y: np.ndarray) -> np.ndarray: + shapes.append((x.shape, y.shape)) + return (x * y).sum(axis=-1) + + actual = apply_ufunc(dot, a, b, input_core_dims=[["z"], ["z"]], vectorize="shared") + assert_identical(actual, xr.DataArray([5.0, 14.0], dims="x")) + assert shapes == [((2, 3), (3,))] + + +def test_vectorize_shared_multiple_outputs() -> None: + a = xr.DataArray(np.arange(12.0).reshape(2, 2, 3), dims=("x", "y", "z")) + b = xr.DataArray(np.ones((2, 3)), dims=("x", "z")) + + def func(x: np.ndarray, y: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + return (x * y).sum(axis=-1), (x * y).max(axis=-1) + + actual_sum, actual_max = apply_ufunc( + func, + a, + b, + input_core_dims=[["z"], ["z"]], + output_core_dims=[[], []], + vectorize="shared", + ) + assert_identical(actual_sum, a.sum("z")) + assert_identical(actual_max, a.max("z")) + + +def test_vectorize_shared_empty_loop() -> None: + a = xr.DataArray(np.zeros((0, 3, 2)), dims=("x", "y", "z")) + b = xr.DataArray(np.zeros((0, 2)), dims=("x", "z")) + + def func(x: np.ndarray, y: np.ndarray) -> np.ndarray: + raise AssertionError("must not be called") + + actual = apply_ufunc( + func, + a, + b, + input_core_dims=[["z"], ["z"]], + vectorize="shared", + output_dtypes=[float], + ) + assert actual.sizes == {"x": 0, "y": 3} + assert actual.dtype == float + + with pytest.raises(ValueError, match=r"empty loop dimension"): + apply_ufunc(func, a, b, input_core_dims=[["z"], ["z"]], vectorize="shared") + + +def test_vectorize_invalid() -> None: + a = xr.DataArray([1.0], dims="x") + with pytest.raises(ValueError, match=r"vectorize must be a bool or 'shared'"): + apply_ufunc(np.sum, a, vectorize="all") # type: ignore[arg-type] + + def test_corr_only_dataarray() -> None: with pytest.raises(TypeError, match=r"Only xr.DataArray is supported"): xr.corr(xr.Dataset(), xr.Dataset()) # type: ignore[type-var] From b743bf12bc6ddd908330c2afa1ba9fade5dd9c52 Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Wed, 7 Oct 2026 20:45:54 +0200 Subject: [PATCH 06/18] Add apply_ufunc examples from basic to advanced Co-authored-by: Claude Co-Authored-By: Claude Opus 5.5 --- xarray/computation/apply_ufunc.py | 122 +++++++++++++++++++++++++----- 1 file changed, 101 insertions(+), 21 deletions(-) diff --git a/xarray/computation/apply_ufunc.py b/xarray/computation/apply_ufunc.py index 0d5fbee88cd..d9096d1b2c5 100644 --- a/xarray/computation/apply_ufunc.py +++ b/xarray/computation/apply_ufunc.py @@ -1193,7 +1193,9 @@ def apply_ufunc( Examples -------- - Calculate the vector magnitude of two arguments: + The examples below go from basic to advanced usage. + + **Element-wise functions.** Calculate the vector magnitude of two arguments: >>> def magnitude(a, b): ... func = lambda x, y: np.sqrt(x**2 + y**2) @@ -1223,17 +1225,19 @@ def apply_ufunc( Coordinates: * x (x) float64 24B 0.1 0.2 0.3 - Other examples of how you could use ``apply_ufunc`` to write functions to - (very nearly) replicate existing xarray functionality: - - Compute the mean (``.mean``) over one dimension: - - >>> def mean(obj, dim): - ... # note: apply always moves core dimensions to the end - ... return apply_ufunc( - ... np.mean, obj, input_core_dims=[[dim]], kwargs={"axis": -1} - ... ) - ... + **Core dimensions.** Functions that operate along a dimension, e.g. a + reduction, need ``input_core_dims``. Core dimensions are moved to the end + of the arrays passed to ``func``, so here we reduce over ``axis=-1``, like + ``.mean("y")``: + + >>> da = xr.DataArray( + ... np.arange(6.0).reshape(2, 3), coords={"x": [10, 20], "y": [0, 1, 2]} + ... ) + >>> xr.apply_ufunc(np.mean, da, input_core_dims=[["y"]], kwargs={"axis": -1}) + Size: 16B + array([1., 4.]) + Coordinates: + * x (x) int64 16B 10 20 Inner product over a specific dimension (like :py:func:`dot`): @@ -1242,35 +1246,85 @@ def apply_ufunc( ... return result[..., 0, 0] ... >>> def inner_product(a, b, dim): - ... return apply_ufunc(_inner, a, b, input_core_dims=[[dim], [dim]]) + ... return xr.apply_ufunc(_inner, a, b, input_core_dims=[[dim], [dim]]) ... + **Output core dimensions.** If ``func`` returns a dimension, list it in + ``output_core_dims``. If its size changes, it also has to be listed in + ``exclude_dims``, and its coordinates are dropped: + + >>> xr.apply_ufunc( + ... np.diff, + ... da, + ... input_core_dims=[["y"]], + ... output_core_dims=[["y"]], + ... exclude_dims={"y"}, + ... kwargs={"axis": -1}, + ... ) + Size: 32B + array([[1., 1.], + [1., 1.]]) + Coordinates: + * x (x) int64 16B 10 20 + Dimensions without coordinates: y + Stack objects along a new dimension (like :py:func:`concat`): >>> def stack(objects, dim, new_coord): ... # note: this version does not stack coordinates ... func = lambda *x: np.stack(x, axis=-1) - ... result = apply_ufunc( + ... result = xr.apply_ufunc( ... func, ... *objects, ... output_core_dims=[[dim]], ... join="outer", - ... dataset_fill_value=np.nan + ... dataset_fill_value=np.nan, ... ) ... result[dim] = new_coord ... return result ... - If your function is not vectorized but can be applied only to core - dimensions, you can use ``vectorize=True`` to turn into a vectorized - function. This wraps :py:func:`numpy.vectorize`, so the operation isn't - terribly fast. Here we'll use it to calculate the distance between - empirical samples from two probability distributions, using a scipy + **Multiple outputs.** A function returning a tuple needs one entry in + ``output_core_dims`` per output: + + >>> quotient, remainder = xr.apply_ufunc( + ... np.divmod, da, 4, output_core_dims=[[], []] + ... ) + >>> remainder + Size: 48B + array([[0., 1., 2.], + [3., 0., 1.]]) + Coordinates: + * x (x) int64 16B 10 20 + * y (y) int64 24B 0 1 2 + + **Dask arrays.** With ``dask="parallelized"``, ``func`` is applied to each + block of a chunked array. Core dimensions must not be split into several + chunks: + + >>> xr.apply_ufunc( + ... np.mean, + ... da.chunk(x=1), + ... input_core_dims=[["y"]], + ... kwargs={"axis": -1}, + ... dask="parallelized", + ... output_dtypes=[float], + ... ).compute() + Size: 16B + array([1., 4.]) + Coordinates: + * x (x) int64 16B 10 20 + + **Vectorizing functions.** If your function is not vectorized but can be + applied only to core dimensions, you can use ``vectorize=True`` to turn into + a vectorized function. This wraps :py:func:`numpy.vectorize`, so the + operation isn't terribly fast. Here we'll use it to calculate the distance + between empirical samples from two probability distributions, using a scipy function that needs to be applied to vectors: >>> import scipy.stats >>> def earth_mover_distance(first_samples, second_samples, dim="ensemble"): - ... return apply_ufunc( + ... return xr.apply_ufunc( ... scipy.stats.wasserstein_distance, ... first_samples, ... second_samples, @@ -1279,6 +1333,32 @@ def apply_ufunc( ... ) ... + **Vectorizing over shared dimensions only.** ``vectorize=True`` calls + ``func`` once for every combination of the broadcast dimensions. If + ``func`` can handle some of them itself, ``vectorize="shared"`` only loops + over the dimensions along which more than one input varies and passes all + others in bulk. Here, ``np.quantile`` computes a single quantile per call + for all ``x`` at once, but the quantile differs along ``time``: + + >>> data = xr.DataArray( + ... np.arange(24.0).reshape(2, 3, 4), dims=("time", "x", "sample") + ... ) + >>> q = xr.DataArray([0.25, 0.75], dims="time") + >>> def quantile(a, q): + ... # a has shape (x, sample), q has size 1 along x + ... return np.quantile(a, q.item(), axis=-1) + ... + >>> xr.apply_ufunc( + ... quantile, data, q, input_core_dims=[["sample"], []], vectorize="shared" + ... ) + Size: 48B + array([[ 0.75, 4.75, 8.75], + [14.25, 18.25, 22.25]]) + Dimensions without coordinates: time, x + + This calls ``quantile`` twice, once per ``time``, while ``vectorize=True`` + would call it six times, once per ``time`` and ``x``. + Most of NumPy's builtin functions already broadcast their inputs appropriately for use in ``apply_ufunc``. You may find helper functions such as :py:func:`numpy.broadcast_arrays` helpful in writing your function. ``apply_ufunc`` also From 28ee4a196dea761971bb1efb17f23a1c4c6474fb Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Wed, 7 Oct 2026 21:05:12 +0200 Subject: [PATCH 07/18] Add vectorize="shared" to the apply_ufunc example notebook Tag the cells that fail on purpose with raises-exception so that the notebook is executed in the documentation build again, fix outdated text and links, and print simple values in the numba example. Co-authored-by: Claude Co-Authored-By: Claude Opus 5.5 --- doc/conf.py | 1 - doc/examples/apply_ufunc_vectorize_1d.ipynb | 120 +++++++++++++++++--- doc/whats-new.rst | 5 + 3 files changed, 109 insertions(+), 17 deletions(-) diff --git a/doc/conf.py b/doc/conf.py index c7e06758962..33a62f0c90c 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -79,7 +79,6 @@ copybutton_prompt_is_regexp = True # myst-nb -nb_execution_excludepatterns = ["examples/apply_ufunc_vectorize_1d.ipynb"] # RTD does not publish the .err.log report, so print tracebacks in the build log nb_execution_show_tb = True diff --git a/doc/examples/apply_ufunc_vectorize_1d.ipynb b/doc/examples/apply_ufunc_vectorize_1d.ipynb index 6348f690d4e..ca4bf89a304 100644 --- a/doc/examples/apply_ufunc_vectorize_1d.ipynb +++ b/doc/examples/apply_ufunc_vectorize_1d.ipynb @@ -139,7 +139,10 @@ "ExecuteTime": { "end_time": "2020-01-15T14:45:59.768626Z", "start_time": "2020-01-15T14:45:59.543808Z" - } + }, + "tags": [ + "raises-exception" + ] }, "outputs": [], "source": [ @@ -187,7 +190,10 @@ "ExecuteTime": { "end_time": "2020-01-15T14:46:02.187012Z", "start_time": "2020-01-15T14:46:02.105563Z" - } + }, + "tags": [ + "raises-exception" + ] }, "outputs": [], "source": [ @@ -247,7 +253,10 @@ "ExecuteTime": { "end_time": "2020-01-15T14:46:05.031672Z", "start_time": "2020-01-15T14:46:04.947588Z" - } + }, + "tags": [ + "raises-exception" + ] }, "outputs": [], "source": [ @@ -351,7 +360,10 @@ "ExecuteTime": { "end_time": "2020-01-15T14:46:13.808646Z", "start_time": "2020-01-15T14:46:13.680098Z" - } + }, + "tags": [ + "raises-exception" + ] }, "outputs": [], "source": [ @@ -381,7 +393,7 @@ "source": [ "That's a hard-to-interpret error but our `print` call helpfully printed the shapes of the input data: \n", "\n", - " data: (10, 53, 25) | x: (25,) | xi: (100,)\n", + " data: (4, 3, 25) | x: (25,) | xi: (100,)\n", "\n", "We see that `air` has been passed as a 3D numpy array which is not what `np.interp` expects. Instead we want loop over all combinations of `lon` and `time`; and apply our function to each corresponding vector of data along `lat`.\n", "`apply_ufunc` makes this easy by specifying `vectorize=True`:\n", @@ -391,9 +403,8 @@ " dimensions as input and vectorize it automatically with\n", " :py:func:`numpy.vectorize`. This option exists for convenience, but is\n", " almost always slower than supplying a pre-vectorized function.\n", - " Using this option requires NumPy version 1.12 or newer.\n", " \n", - "Also see the documentation for `np.vectorize`: https://docs.scipy.org/doc/numpy/reference/generated/numpy.vectorize.html. Most importantly\n", + "Also see the documentation for `np.vectorize`: https://numpy.org/doc/stable/reference/generated/numpy.vectorize.html. Most importantly\n", "\n", " The vectorize function is provided primarily for convenience, not for performance. \n", " The implementation is essentially a for loop." @@ -406,7 +417,10 @@ "ExecuteTime": { "end_time": "2020-01-15T14:46:26.633233Z", "start_time": "2020-01-15T14:46:26.515209Z" - } + }, + "tags": [ + "raises-exception" + ] }, "outputs": [], "source": [ @@ -550,6 +564,82 @@ "\n" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Vectorizing over shared dimensions only with `vectorize=\"shared\"`" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "`vectorize=True` calls our function once for every combination of the non-core dimensions, here 12 times for 4 `time` and 3 `lon` values. That is slow for large arrays, but often not necessary: many functions can process a whole stack of vectors at once, as long as the other inputs are the same for all of them.\n", + "\n", + "Let's interpolate to different latitudes at every time step. The new latitudes now vary along `time`, but not along `lon`:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "newlat_t = xr.DataArray(\n", + " np.linspace(20, 70, 100) + np.arange(4)[:, np.newaxis],\n", + " dims=(\"time\", \"new_lat\"),\n", + " coords={\"time\": air.time},\n", + ")\n", + "expected_t = air.interp(lat=newlat_t).drop_vars(\"lat\")\n", + "newlat_t" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "With `vectorize=\"shared\"`, `apply_ufunc` only loops over the non-core dimensions along which more than one input varies, here `time`. All other non-core dimensions, here `lon`, are passed to our function in bulk. So our function now has to handle extra leading dimensions, following numpy's broadcasting rules. `np.interp` only supports 1D data, so we use scipy's `make_interp_spline` with `k=1` (linear) instead, which interpolates all vectors along the last axis at once:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from scipy.interpolate import make_interp_spline\n", + "\n", + "\n", + "def interp1d_batched(data, x, xi):\n", + " print(f\"data: {data.shape} | x: {x.shape} | xi: {xi.shape}\")\n", + " # xi does not vary along lon, so it has size 1 along that dimension\n", + " xi = xi.reshape(xi.shape[-1])\n", + " return make_interp_spline(x, data, k=1, axis=-1)(xi)\n", + "\n", + "\n", + "interped_t = xr.apply_ufunc(\n", + " interp1d_batched,\n", + " air,\n", + " air.lat,\n", + " newlat_t,\n", + " input_core_dims=[[\"lat\"], [\"lat\"], [\"new_lat\"]],\n", + " output_core_dims=[[\"new_lat\"]],\n", + " exclude_dims={\"lat\"},\n", + " vectorize=\"shared\", # only loop over dimensions shared by several inputs\n", + ")\n", + "xr.testing.assert_allclose(expected_t.transpose(*interped_t.dims), interped_t)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Our function is now called only 4 times, once per `time`, and receives all `lon` at once: `data` has shape `(lon, lat)`. `xi` keeps a size-1 axis for `lon`, because non-core dimensions follow numpy's broadcasting rules, so we drop it before interpolating. This also works with `dask=\"parallelized\"`, where the loop runs within each block.\n", + "\n", + "If none of the non-core dimensions are shared, e.g. for `newlat` from above, the function is called only once with all data." + ] + }, { "cell_type": "markdown", "metadata": {}, @@ -564,7 +654,7 @@ "source": [ "`np.vectorize` is a very convenient function but is unfortunately slow. It is only marginally faster than writing a for loop in Python and looping. A common way to get around this is to write a base interpolation function that can handle nD arrays in a compiled language like Fortran and then pass that to `apply_ufunc`.\n", "\n", - "Another option is to use the numba package which provides a very convenient `guvectorize` decorator: https://numba.pydata.org/numba-doc/latest/user/vectorize.html#the-guvectorize-decorator\n", + "Another option is to use the numba package which provides a very convenient `guvectorize` decorator: https://numba.readthedocs.io/en/stable/user/vectorize.html#the-guvectorize-decorator\n", "\n", "Any decorated function gets compiled and will loop over any non-core dimension in parallel when necessary. We need to specify some extra information:\n", "\n", @@ -573,7 +663,8 @@ " 2. `guvectorize` needs to know the dtypes of the input and output. This is specified in string form as the first argument. Each element of the tuple corresponds to each argument of the function. In this case, we specify `float64` for all inputs and outputs: `\"(float64[:], float64[:], float64[:], float64[:])\"` corresponding to `data, x, xi, out`\n", " \n", " 3. Now we need to tell numba the size of the dimensions the function takes as inputs and returns as output i.e. core dimensions. This is done in symbolic form i.e. `data` and `x` are vectors of the same length, say `n`; `xi` and the output `out` have a different length, say `m`. So the second argument is (again as a string)\n", - " `\"(n), (n), (m) -> (m).\"` corresponding again to `data, x, xi, out`\n" + " `\"(n), (n), (m) -> (m).\"` corresponding again to `data, x, xi, out`\n", + "" ] }, { @@ -592,11 +683,8 @@ "\n", "@guvectorize(\"(float64[:], float64[:], float64[:], float64[:])\", \"(n), (n), (m) -> (m)\")\n", "def interp1d_np_gufunc(data, x, xi, out):\n", - " # numba doesn't really like this.\n", - " # seem to support fstrings so do it the old way\n", - " print(\n", - " \"data: \" + str(data.shape) + \" | x:\" + str(x.shape) + \" | xi: \" + str(xi.shape)\n", - " )\n", + " # numba can only print simple values, so we print the lengths of the vectors\n", + " print(\"data:\", data.shape[0], \"| x:\", x.shape[0], \"| xi:\", xi.shape[0])\n", " out[:] = np.interp(xi, x, data)\n", " # gufuncs don't return data\n", " # instead you assign to a the last arg\n", @@ -607,7 +695,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "The warnings are about object-mode compilation relating to the `print` statement. This means we don't get much speed up: https://numba.pydata.org/numba-doc/latest/user/performance-tips.html#no-python-mode-vs-object-mode. We'll keep the `print` statement temporarily to make sure that `guvectorize` acts like we want it to." + "We'll keep the `print` statement temporarily to make sure that `guvectorize` acts like we want it to." ] }, { diff --git a/doc/whats-new.rst b/doc/whats-new.rst index cfb8d78646a..951a8c4c852 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -98,6 +98,11 @@ Bug Fixes Documentation ~~~~~~~~~~~~~ +- Add examples from basic to advanced usage to the :py:func:`apply_ufunc` + docstring, and a section on ``vectorize="shared"`` to the + :doc:`apply_ufunc example notebook `, + which is executed again in the documentation build (:pull:`11681`). + By `Michael Niklas `_. - Clarified the ``rename`` docstrings so they no longer describe the result as a "new" object, which could be read as implying it no longer shares memory with the original (:issue:`9432`, :pull:`11644`). From 13ff770fe6680d7e6811a77e588fa2f998678659 Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Thu, 8 Oct 2026 11:20:34 +0200 Subject: [PATCH 08/18] Use the numpy docstring style for _vectorize_shared Co-authored-by: Claude Co-Authored-By: Claude Opus 5.5 --- xarray/computation/apply_ufunc.py | 27 +++++++++++++++++++++------ 1 file changed, 21 insertions(+), 6 deletions(-) diff --git a/xarray/computation/apply_ufunc.py b/xarray/computation/apply_ufunc.py index d9096d1b2c5..695fe226aab 100644 --- a/xarray/computation/apply_ufunc.py +++ b/xarray/computation/apply_ufunc.py @@ -715,13 +715,28 @@ def _vectorize_shared( Like :py:class:`numpy.vectorize` with a signature, but ``func`` is only looped over the loop (non-core) dimensions along which more than one input varies, i.e. has a size other than 1. All other loop dimensions are passed - to ``func`` in bulk, so ``func`` must accept inputs with leading loop - dimensions that broadcast against each other following numpy's rules, and - return outputs with the broadcast loop dimensions followed by the output - core dimensions. + to ``func`` in bulk. - ``output_core_sizes`` maps output core dimensions to their sizes; it is only - needed if the loop is empty, so that ``func`` is never called. + Parameters + ---------- + func : callable + Function to vectorize. It must accept inputs with leading loop + dimensions that broadcast against each other following numpy's rules, + and return outputs with the broadcast loop dimensions followed by the + output core dimensions. + signature : _UFuncSignature + Core dimensions signature for the operation. + output_dtypes : sequence of dtype or None + Dtypes of the outputs. If None, they are taken from the results of + ``func``, so they are required if the loop is empty. + output_core_sizes : mapping of hashable to int + Sizes of the output core dimensions. Only needed if the loop is empty, + so that ``func`` is never called. + + Returns + ------- + callable + Vectorized version of ``func``, taking the same arguments. """ input_core_ndims = [len(dims) for dims in signature.input_core_dims] num_outputs = signature.num_outputs From 597bdd7375a15f4937f3928c5371bbe0befbf9a6 Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Thu, 8 Oct 2026 11:21:19 +0200 Subject: [PATCH 09/18] Document the numpy docstring style in CLAUDE.md Co-authored-by: Claude Co-Authored-By: Claude Opus 5.5 --- CLAUDE.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/CLAUDE.md b/CLAUDE.md index b4c0061bb43..dfe97559a98 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -32,6 +32,11 @@ uv run dmypy run # Type checking with mypy 2. Related third-party imports 3. Local application/library specific imports +### Docstrings + +- Use the numpy docstring style (`Parameters`, `Returns`, ... sections), also + for private functions + ## GitHub Interaction Guidelines - **NEVER impersonate the user on GitHub**, always sign off with something like From 3b4fa2810b2f993b112b9a4c27d47bf64b42e8c4 Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Thu, 8 Oct 2026 11:25:12 +0200 Subject: [PATCH 10/18] Type the function vectorized by _vectorize_shared Co-authored-by: Claude Co-Authored-By: Claude Opus 5.5 --- xarray/computation/apply_ufunc.py | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/xarray/computation/apply_ufunc.py b/xarray/computation/apply_ufunc.py index 695fe226aab..1842e0dc3be 100644 --- a/xarray/computation/apply_ufunc.py +++ b/xarray/computation/apply_ufunc.py @@ -20,7 +20,7 @@ from collections.abc import ( Set as AbstractSet, ) -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any, Literal, cast import numpy as np @@ -704,12 +704,12 @@ def _vectorize(func, signature, output_dtypes, exclude_dims): return func -def _vectorize_shared( - func: Callable, +def _vectorize_shared[T]( + func: Callable[..., T], signature: _UFuncSignature, output_dtypes: Sequence | None, output_core_sizes: Mapping[Hashable, int], -) -> Callable: +) -> Callable[..., T]: """Vectorize ``func`` over the loop dimensions shared by several inputs. Like :py:class:`numpy.vectorize` with a signature, but ``func`` is only @@ -742,7 +742,7 @@ def _vectorize_shared( num_outputs = signature.num_outputs @functools.wraps(func) - def wrapper(*args): + def wrapper(*args: Any) -> T: loop_shapes = [ np.shape(arg)[: np.ndim(arg) - ndim] for arg, ndim in zip(args, input_core_ndims, strict=True) @@ -778,14 +778,17 @@ def wrapper(*args): ) sliced_args.append(arg[key] if key else arg) - results = func(*sliced_args) + result = func(*sliced_args) + results: tuple[Any, ...] if num_outputs == 1: - results = (results,) - elif not isinstance(results, tuple) or len(results) != num_outputs: + results = (result,) + elif isinstance(result, tuple) and len(result) == num_outputs: + results = result + else: raise ValueError( f"applied function does not have the number of outputs " f"specified in the ufunc signature. Expected {num_outputs} " - f"outputs, got {results!r}" + f"outputs, got {result!r}" ) if outputs is None: @@ -829,7 +832,7 @@ def wrapper(*args): signature.output_core_dims, output_dtypes, strict=True ) ) - return outputs[0] if num_outputs == 1 else outputs + return cast(T, outputs[0] if num_outputs == 1 else outputs) return wrapper From da39dcad94c359136b77da0b603229e094c2877a Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Tue, 20 Oct 2026 22:48:58 +0200 Subject: [PATCH 11/18] Replace vectorize="shared" with explicit loop dimensions, drop example notebook apply_ufunc now takes the dimensions to loop over, e.g. vectorize=["time"], instead of the implicit "shared" mode. The apply_ufunc example notebook is removed in favor of the xarray-tutorial interpolation example. Co-Authored-By: Claude Opus 5.5 --- doc/examples/apply_ufunc_vectorize_1d.ipynb | 829 -------------------- doc/gallery.rst | 1 - doc/gallery.yml | 4 +- doc/whats-new.rst | 13 +- xarray/computation/apply_ufunc.py | 118 +-- xarray/core/missing.py | 10 +- xarray/tests/test_computation.py | 71 +- 7 files changed, 140 insertions(+), 906 deletions(-) delete mode 100644 doc/examples/apply_ufunc_vectorize_1d.ipynb diff --git a/doc/examples/apply_ufunc_vectorize_1d.ipynb b/doc/examples/apply_ufunc_vectorize_1d.ipynb deleted file mode 100644 index ca4bf89a304..00000000000 --- a/doc/examples/apply_ufunc_vectorize_1d.ipynb +++ /dev/null @@ -1,829 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": { - "deletable": true, - "editable": true, - "slideshow": { - "slide_type": "" - }, - "tags": [] - }, - "source": [ - "# Applying unvectorized functions with `apply_ufunc`" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This example will illustrate how to conveniently apply an unvectorized function `func` to xarray objects using `apply_ufunc`. `func` expects 1D numpy arrays and returns a 1D numpy array. Our goal is to conveniently apply this function along a dimension of xarray objects that may or may not wrap dask arrays with a signature.\n", - "\n", - "We will illustrate this using `np.interp`: \n", - "\n", - " Signature: np.interp(x, xp, fp, left=None, right=None, period=None)\n", - " Docstring:\n", - " One-dimensional linear interpolation.\n", - "\n", - " Returns the one-dimensional piecewise linear interpolant to a function\n", - " with given discrete data points (`xp`, `fp`), evaluated at `x`.\n", - "\n", - "and write an `xr_interp` function with signature\n", - "\n", - " xr_interp(xarray_object, dimension_name, new_coordinate_to_interpolate_to)\n", - "\n", - "## Load data\n", - "\n", - "First let's load an example dataset" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "ExecuteTime": { - "end_time": "2020-01-15T14:45:51.659160Z", - "start_time": "2020-01-15T14:45:50.528742Z" - } - }, - "outputs": [], - "source": [ - "import xarray as xr\n", - "import numpy as np\n", - "\n", - "xr.set_options(display_style=\"html\") # fancy HTML repr\n", - "\n", - "air = (\n", - " xr.tutorial.load_dataset(\"air_temperature\")\n", - " .air.sortby(\"lat\") # np.interp needs coordinate in ascending order\n", - " .isel(time=slice(4), lon=slice(3))\n", - ") # choose a small subset for convenience\n", - "air" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The function we will apply is `np.interp` which expects 1D numpy arrays. This functionality is already implemented in xarray so we use that capability to make sure we are not making mistakes." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "ExecuteTime": { - "end_time": "2020-01-15T14:45:55.431708Z", - "start_time": "2020-01-15T14:45:55.104701Z" - } - }, - "outputs": [], - "source": [ - "newlat = np.linspace(15, 75, 100)\n", - "air.interp(lat=newlat)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Let's define a function that works with one vector of data along `lat` at a time." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "ExecuteTime": { - "end_time": "2020-01-15T14:45:57.889496Z", - "start_time": "2020-01-15T14:45:57.792269Z" - } - }, - "outputs": [], - "source": [ - "def interp1d_np(data, x, xi):\n", - " return np.interp(xi, x, data)\n", - "\n", - "\n", - "interped = interp1d_np(air.isel(time=0, lon=0), air.lat, newlat)\n", - "expected = air.interp(lat=newlat)\n", - "\n", - "# no errors are raised if values are equal to within floating point precision\n", - "np.testing.assert_allclose(expected.isel(time=0, lon=0).values, interped)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## No errors are raised so our interpolation is working.\n", - "\n", - "This function consumes and returns numpy arrays, which means we need to do a lot of work to convert the result back to an xarray object with meaningful metadata. This is where `apply_ufunc` is very useful.\n", - "\n", - "### `apply_ufunc`\n", - "\n", - " Apply a vectorized function for unlabeled arrays on xarray objects.\n", - "\n", - " The function will be mapped over the data variable(s) of the input arguments using \n", - " xarray’s standard rules for labeled computation, including alignment, broadcasting, \n", - " looping over GroupBy/Dataset variables, and merging of coordinates.\n", - " \n", - "`apply_ufunc` has many capabilities but for simplicity this example will focus on the common task of vectorizing 1D functions over nD xarray objects. We will iteratively build up the right set of arguments to `apply_ufunc` and read through many error messages in doing so." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "ExecuteTime": { - "end_time": "2020-01-15T14:45:59.768626Z", - "start_time": "2020-01-15T14:45:59.543808Z" - }, - "tags": [ - "raises-exception" - ] - }, - "outputs": [], - "source": [ - "xr.apply_ufunc(\n", - " interp1d_np, # first the function\n", - " air.isel(time=0, lon=0), # now arguments in the order expected by 'interp1_np'\n", - " air.lat,\n", - " newlat,\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "`apply_ufunc` needs to know a lot of information about what our function does so that it can reconstruct the outputs. In this case, the size of dimension lat has changed and we need to explicitly specify that this will happen. xarray helpfully tells us that we need to specify the kwarg `exclude_dims`." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## `exclude_dims`" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "\n", - "```\n", - "exclude_dims : set, optional\n", - " Core dimensions on the inputs to exclude from alignment and\n", - " broadcasting entirely. Any input coordinates along these dimensions\n", - " will be dropped. Each excluded dimension must also appear in\n", - " ``input_core_dims`` for at least one argument. Only dimensions listed\n", - " here are allowed to change size between input and output objects.\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "ExecuteTime": { - "end_time": "2020-01-15T14:46:02.187012Z", - "start_time": "2020-01-15T14:46:02.105563Z" - }, - "tags": [ - "raises-exception" - ] - }, - "outputs": [], - "source": [ - "xr.apply_ufunc(\n", - " interp1d_np, # first the function\n", - " air.isel(time=0, lon=0), # now arguments in the order expected by 'interp1_np'\n", - " air.lat,\n", - " newlat,\n", - " exclude_dims=set((\"lat\",)), # dimensions allowed to change size. Must be set!\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Core dimensions\n", - "\n" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Core dimensions are central to using `apply_ufunc`. In our case, our function expects to receive a 1D vector along `lat` — this is the dimension that is \"core\" to the function's functionality. Multiple core dimensions are possible. `apply_ufunc` needs to know which dimensions of each variable are core dimensions.\n", - "\n", - " input_core_dims : Sequence[Sequence], optional\n", - " List of the same length as ``args`` giving the list of core dimensions\n", - " on each input argument that should not be broadcast. By default, we\n", - " assume there are no core dimensions on any input arguments.\n", - "\n", - " For example, ``input_core_dims=[[], ['time']]`` indicates that all\n", - " dimensions on the first argument and all dimensions other than 'time'\n", - " on the second argument should be broadcast.\n", - "\n", - " Core dimensions are automatically moved to the last axes of input\n", - " variables before applying ``func``, which facilitates using NumPy style\n", - " generalized ufuncs [2]_.\n", - " \n", - " output_core_dims : List[tuple], optional\n", - " List of the same length as the number of output arguments from\n", - " ``func``, giving the list of core dimensions on each output that were\n", - " not broadcast on the inputs. By default, we assume that ``func``\n", - " outputs exactly one array, with axes corresponding to each broadcast\n", - " dimension.\n", - "\n", - " Core dimensions are assumed to appear as the last dimensions of each\n", - " output in the provided order.\n", - " \n", - "Next we specify `\"lat\"` as `input_core_dims` on both `air` and `air.lat`" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "ExecuteTime": { - "end_time": "2020-01-15T14:46:05.031672Z", - "start_time": "2020-01-15T14:46:04.947588Z" - }, - "tags": [ - "raises-exception" - ] - }, - "outputs": [], - "source": [ - "xr.apply_ufunc(\n", - " interp1d_np, # first the function\n", - " air.isel(time=0, lon=0), # now arguments in the order expected by 'interp1_np'\n", - " air.lat,\n", - " newlat,\n", - " input_core_dims=[[\"lat\"], [\"lat\"], []],\n", - " exclude_dims=set((\"lat\",)), # dimensions allowed to change size. Must be set!\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "xarray is telling us that it expected to receive back a numpy array with 0 dimensions but instead received an array with 1 dimension corresponding to `newlat`. We can fix this by specifying `output_core_dims`" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "ExecuteTime": { - "end_time": "2020-01-15T14:46:09.325218Z", - "start_time": "2020-01-15T14:46:09.303020Z" - } - }, - "outputs": [], - "source": [ - "xr.apply_ufunc(\n", - " interp1d_np, # first the function\n", - " air.isel(time=0, lon=0), # now arguments in the order expected by 'interp1_np'\n", - " air.lat,\n", - " newlat,\n", - " input_core_dims=[[\"lat\"], [\"lat\"], []], # list with one entry per arg\n", - " output_core_dims=[[\"lat\"]],\n", - " exclude_dims=set((\"lat\",)), # dimensions allowed to change size. Must be set!\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finally we get some output! Let's check that this is right\n", - "\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "ExecuteTime": { - "end_time": "2020-01-15T14:46:11.295440Z", - "start_time": "2020-01-15T14:46:11.226553Z" - } - }, - "outputs": [], - "source": [ - "interped = xr.apply_ufunc(\n", - " interp1d_np, # first the function\n", - " air.isel(time=0, lon=0), # now arguments in the order expected by 'interp1_np'\n", - " air.lat,\n", - " newlat,\n", - " input_core_dims=[[\"lat\"], [\"lat\"], []], # list with one entry per arg\n", - " output_core_dims=[[\"lat\"]],\n", - " exclude_dims=set((\"lat\",)), # dimensions allowed to change size. Must be set!\n", - ")\n", - "interped[\"lat\"] = newlat # need to add this manually\n", - "xr.testing.assert_allclose(expected.isel(time=0, lon=0), interped)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "No errors are raised so it is right!" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Vectorization with `np.vectorize`" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now our function currently only works on one vector of data which is not so useful given our 3D dataset.\n", - "Let's try passing the whole dataset. We add a `print` statement so we can see what our function receives." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "ExecuteTime": { - "end_time": "2020-01-15T14:46:13.808646Z", - "start_time": "2020-01-15T14:46:13.680098Z" - }, - "tags": [ - "raises-exception" - ] - }, - "outputs": [], - "source": [ - "def interp1d_np(data, x, xi):\n", - " print(f\"data: {data.shape} | x: {x.shape} | xi: {xi.shape}\")\n", - " return np.interp(xi, x, data)\n", - "\n", - "\n", - "interped = xr.apply_ufunc(\n", - " interp1d_np, # first the function\n", - " air.isel(\n", - " lon=slice(3), time=slice(4)\n", - " ), # now arguments in the order expected by 'interp1_np'\n", - " air.lat,\n", - " newlat,\n", - " input_core_dims=[[\"lat\"], [\"lat\"], []], # list with one entry per arg\n", - " output_core_dims=[[\"lat\"]],\n", - " exclude_dims=set((\"lat\",)), # dimensions allowed to change size. Must be set!\n", - ")\n", - "interped[\"lat\"] = newlat # need to add this manually\n", - "xr.testing.assert_allclose(expected.isel(time=0, lon=0), interped)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "That's a hard-to-interpret error but our `print` call helpfully printed the shapes of the input data: \n", - "\n", - " data: (4, 3, 25) | x: (25,) | xi: (100,)\n", - "\n", - "We see that `air` has been passed as a 3D numpy array which is not what `np.interp` expects. Instead we want loop over all combinations of `lon` and `time`; and apply our function to each corresponding vector of data along `lat`.\n", - "`apply_ufunc` makes this easy by specifying `vectorize=True`:\n", - "\n", - " vectorize : bool, optional\n", - " If True, then assume ``func`` only takes arrays defined over core\n", - " dimensions as input and vectorize it automatically with\n", - " :py:func:`numpy.vectorize`. This option exists for convenience, but is\n", - " almost always slower than supplying a pre-vectorized function.\n", - " \n", - "Also see the documentation for `np.vectorize`: https://numpy.org/doc/stable/reference/generated/numpy.vectorize.html. Most importantly\n", - "\n", - " The vectorize function is provided primarily for convenience, not for performance. \n", - " The implementation is essentially a for loop." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "ExecuteTime": { - "end_time": "2020-01-15T14:46:26.633233Z", - "start_time": "2020-01-15T14:46:26.515209Z" - }, - "tags": [ - "raises-exception" - ] - }, - "outputs": [], - "source": [ - "def interp1d_np(data, x, xi):\n", - " print(f\"data: {data.shape} | x: {x.shape} | xi: {xi.shape}\")\n", - " return np.interp(xi, x, data)\n", - "\n", - "\n", - "interped = xr.apply_ufunc(\n", - " interp1d_np, # first the function\n", - " air, # now arguments in the order expected by 'interp1_np'\n", - " air.lat, # as above\n", - " newlat, # as above\n", - " input_core_dims=[[\"lat\"], [\"lat\"], []], # list with one entry per arg\n", - " output_core_dims=[[\"lat\"]], # returned data has one dimension\n", - " exclude_dims=set((\"lat\",)), # dimensions allowed to change size. Must be set!\n", - " vectorize=True, # loop over non-core dims\n", - ")\n", - "interped[\"lat\"] = newlat # need to add this manually\n", - "xr.testing.assert_allclose(expected, interped)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This unfortunately is another cryptic error from numpy. \n", - "\n", - "Notice that `newlat` is not an xarray object. Let's add a dimension name `new_lat` and modify the call. Note this cannot be `lat` because xarray expects dimensions to be the same size (or broadcastable) among all inputs. `output_core_dims` needs to be modified appropriately. We'll manually rename `new_lat` back to `lat` for easy checking." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "ExecuteTime": { - "end_time": "2020-01-15T14:46:30.026663Z", - "start_time": "2020-01-15T14:46:29.893267Z" - } - }, - "outputs": [], - "source": [ - "def interp1d_np(data, x, xi):\n", - " print(f\"data: {data.shape} | x: {x.shape} | xi: {xi.shape}\")\n", - " return np.interp(xi, x, data)\n", - "\n", - "\n", - "interped = xr.apply_ufunc(\n", - " interp1d_np, # first the function\n", - " air, # now arguments in the order expected by 'interp1_np'\n", - " air.lat, # as above\n", - " newlat, # as above\n", - " input_core_dims=[[\"lat\"], [\"lat\"], [\"new_lat\"]], # list with one entry per arg\n", - " output_core_dims=[[\"new_lat\"]], # returned data has one dimension\n", - " exclude_dims=set((\"lat\",)), # dimensions allowed to change size. Must be a set!\n", - " vectorize=True, # loop over non-core dims\n", - ")\n", - "interped = interped.rename({\"new_lat\": \"lat\"})\n", - "interped[\"lat\"] = newlat # need to add this manually\n", - "xr.testing.assert_allclose(\n", - " expected.transpose(*interped.dims), interped # order of dims is different\n", - ")\n", - "interped" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Notice that the printed input shapes are all 1D and correspond to one vector along the `lat` dimension.\n", - "\n", - "The result is now an xarray object with coordinate values copied over from `data`. This is why `apply_ufunc` is so convenient; it takes care of a lot of boilerplate necessary to apply functions that consume and produce numpy arrays to xarray objects.\n", - "\n", - "One final point: `lat` is now the *last* dimension in `interped`. This is a \"property\" of core dimensions: they are moved to the end before being sent to `interp1d_np` as was noted in the docstring for `input_core_dims`\n", - "\n", - " Core dimensions are automatically moved to the last axes of input\n", - " variables before applying ``func``, which facilitates using NumPy style\n", - " generalized ufuncs [2]_." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Parallelization with dask\n", - "\n" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "So far our function can only handle numpy arrays. A real benefit of `apply_ufunc` is the ability to easily parallelize over dask chunks _when needed_. \n", - "\n", - "We want to apply this function in a vectorized fashion over each chunk of the dask array. This is possible using dask's `blockwise`, `map_blocks`, or `apply_gufunc`. Xarray's `apply_ufunc` wraps dask's `apply_gufunc` and asking it to map the function over chunks using `apply_gufunc` is as simple as specifying `dask=\"parallelized\"`. With this level of flexibility we need to provide dask with some extra information: \n", - " 1. `output_dtypes`: dtypes of all returned objects, and \n", - " 2. `output_sizes`: lengths of any new dimensions. \n", - " \n", - "Here we need to specify `output_dtypes` since `apply_ufunc` can infer the size of the new dimension `new_lat` from the argument corresponding to the third element in `input_core_dims`. Here I choose the chunk sizes to illustrate that `np.vectorize` is still applied so that our function receives 1D vectors even though the blocks are 3D." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "ExecuteTime": { - "end_time": "2020-01-15T14:48:42.469341Z", - "start_time": "2020-01-15T14:48:42.344209Z" - } - }, - "outputs": [], - "source": [ - "def interp1d_np(data, x, xi):\n", - " print(f\"data: {data.shape} | x: {x.shape} | xi: {xi.shape}\")\n", - " return np.interp(xi, x, data)\n", - "\n", - "\n", - "interped = xr.apply_ufunc(\n", - " interp1d_np, # first the function\n", - " air.chunk(\n", - " {\"time\": 2, \"lon\": 2}\n", - " ), # now arguments in the order expected by 'interp1_np'\n", - " air.lat, # as above\n", - " newlat, # as above\n", - " input_core_dims=[[\"lat\"], [\"lat\"], [\"new_lat\"]], # list with one entry per arg\n", - " output_core_dims=[[\"new_lat\"]], # returned data has one dimension\n", - " exclude_dims=set((\"lat\",)), # dimensions allowed to change size. Must be a set!\n", - " vectorize=True, # loop over non-core dims\n", - " dask=\"parallelized\",\n", - " output_dtypes=[air.dtype], # one per output\n", - ").rename({\"new_lat\": \"lat\"})\n", - "interped[\"lat\"] = newlat # need to add this manually\n", - "xr.testing.assert_allclose(expected.transpose(*interped.dims), interped)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Yay! our function is receiving 1D vectors, so we've successfully parallelized applying a 1D function over a block. If you have a distributed dashboard up, you should see computes happening as equality is checked.\n", - "\n" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Vectorizing over shared dimensions only with `vectorize=\"shared\"`" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "`vectorize=True` calls our function once for every combination of the non-core dimensions, here 12 times for 4 `time` and 3 `lon` values. That is slow for large arrays, but often not necessary: many functions can process a whole stack of vectors at once, as long as the other inputs are the same for all of them.\n", - "\n", - "Let's interpolate to different latitudes at every time step. The new latitudes now vary along `time`, but not along `lon`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "newlat_t = xr.DataArray(\n", - " np.linspace(20, 70, 100) + np.arange(4)[:, np.newaxis],\n", - " dims=(\"time\", \"new_lat\"),\n", - " coords={\"time\": air.time},\n", - ")\n", - "expected_t = air.interp(lat=newlat_t).drop_vars(\"lat\")\n", - "newlat_t" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "With `vectorize=\"shared\"`, `apply_ufunc` only loops over the non-core dimensions along which more than one input varies, here `time`. All other non-core dimensions, here `lon`, are passed to our function in bulk. So our function now has to handle extra leading dimensions, following numpy's broadcasting rules. `np.interp` only supports 1D data, so we use scipy's `make_interp_spline` with `k=1` (linear) instead, which interpolates all vectors along the last axis at once:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from scipy.interpolate import make_interp_spline\n", - "\n", - "\n", - "def interp1d_batched(data, x, xi):\n", - " print(f\"data: {data.shape} | x: {x.shape} | xi: {xi.shape}\")\n", - " # xi does not vary along lon, so it has size 1 along that dimension\n", - " xi = xi.reshape(xi.shape[-1])\n", - " return make_interp_spline(x, data, k=1, axis=-1)(xi)\n", - "\n", - "\n", - "interped_t = xr.apply_ufunc(\n", - " interp1d_batched,\n", - " air,\n", - " air.lat,\n", - " newlat_t,\n", - " input_core_dims=[[\"lat\"], [\"lat\"], [\"new_lat\"]],\n", - " output_core_dims=[[\"new_lat\"]],\n", - " exclude_dims={\"lat\"},\n", - " vectorize=\"shared\", # only loop over dimensions shared by several inputs\n", - ")\n", - "xr.testing.assert_allclose(expected_t.transpose(*interped_t.dims), interped_t)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Our function is now called only 4 times, once per `time`, and receives all `lon` at once: `data` has shape `(lon, lat)`. `xi` keeps a size-1 axis for `lon`, because non-core dimensions follow numpy's broadcasting rules, so we drop it before interpolating. This also works with `dask=\"parallelized\"`, where the loop runs within each block.\n", - "\n", - "If none of the non-core dimensions are shared, e.g. for `newlat` from above, the function is called only once with all data." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## High performance vectorization: gufuncs, numba & guvectorize\n", - "\n" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "`np.vectorize` is a very convenient function but is unfortunately slow. It is only marginally faster than writing a for loop in Python and looping. A common way to get around this is to write a base interpolation function that can handle nD arrays in a compiled language like Fortran and then pass that to `apply_ufunc`.\n", - "\n", - "Another option is to use the numba package which provides a very convenient `guvectorize` decorator: https://numba.readthedocs.io/en/stable/user/vectorize.html#the-guvectorize-decorator\n", - "\n", - "Any decorated function gets compiled and will loop over any non-core dimension in parallel when necessary. We need to specify some extra information:\n", - "\n", - " 1. Our function cannot return a variable any more. Instead it must receive a variable (the last argument) whose contents the function will modify. So we change from `def interp1d_np(data, x, xi)` to `def interp1d_np_gufunc(data, x, xi, out)`. Our computed results must be assigned to `out`. All values of `out` must be assigned explicitly.\n", - " \n", - " 2. `guvectorize` needs to know the dtypes of the input and output. This is specified in string form as the first argument. Each element of the tuple corresponds to each argument of the function. In this case, we specify `float64` for all inputs and outputs: `\"(float64[:], float64[:], float64[:], float64[:])\"` corresponding to `data, x, xi, out`\n", - " \n", - " 3. Now we need to tell numba the size of the dimensions the function takes as inputs and returns as output i.e. core dimensions. This is done in symbolic form i.e. `data` and `x` are vectors of the same length, say `n`; `xi` and the output `out` have a different length, say `m`. So the second argument is (again as a string)\n", - " `\"(n), (n), (m) -> (m).\"` corresponding again to `data, x, xi, out`\n", - "" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "ExecuteTime": { - "end_time": "2020-01-15T14:48:45.267633Z", - "start_time": "2020-01-15T14:48:44.943939Z" - } - }, - "outputs": [], - "source": [ - "from numba import float64, guvectorize\n", - "\n", - "\n", - "@guvectorize(\"(float64[:], float64[:], float64[:], float64[:])\", \"(n), (n), (m) -> (m)\")\n", - "def interp1d_np_gufunc(data, x, xi, out):\n", - " # numba can only print simple values, so we print the lengths of the vectors\n", - " print(\"data:\", data.shape[0], \"| x:\", x.shape[0], \"| xi:\", xi.shape[0])\n", - " out[:] = np.interp(xi, x, data)\n", - " # gufuncs don't return data\n", - " # instead you assign to a the last arg\n", - " # return np.interp(xi, x, data)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We'll keep the `print` statement temporarily to make sure that `guvectorize` acts like we want it to." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "ExecuteTime": { - "end_time": "2020-01-15T14:48:54.755405Z", - "start_time": "2020-01-15T14:48:54.634724Z" - } - }, - "outputs": [], - "source": [ - "interped = xr.apply_ufunc(\n", - " interp1d_np_gufunc, # first the function\n", - " air.chunk(\n", - " {\"time\": 2, \"lon\": 2}\n", - " ), # now arguments in the order expected by 'interp1_np'\n", - " air.lat, # as above\n", - " newlat, # as above\n", - " input_core_dims=[[\"lat\"], [\"lat\"], [\"new_lat\"]], # list with one entry per arg\n", - " output_core_dims=[[\"new_lat\"]], # returned data has one dimension\n", - " exclude_dims=set((\"lat\",)), # dimensions allowed to change size. Must be a set!\n", - " # vectorize=True, # not needed since numba takes care of vectorizing\n", - " dask=\"parallelized\",\n", - " output_dtypes=[air.dtype], # one per output\n", - ").rename({\"new_lat\": \"lat\"})\n", - "interped[\"lat\"] = newlat # need to add this manually\n", - "xr.testing.assert_allclose(expected.transpose(*interped.dims), interped)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Yay! Our function is receiving 1D vectors and is working automatically with dask arrays. Finally let's comment out the print line and wrap everything up in a nice reusable function" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "ExecuteTime": { - "end_time": "2020-01-15T14:49:28.667528Z", - "start_time": "2020-01-15T14:49:28.103914Z" - } - }, - "outputs": [], - "source": [ - "from numba import float64, guvectorize\n", - "\n", - "\n", - "@guvectorize(\n", - " \"(float64[:], float64[:], float64[:], float64[:])\",\n", - " \"(n), (n), (m) -> (m)\",\n", - " nopython=True,\n", - ")\n", - "def interp1d_np_gufunc(data, x, xi, out):\n", - " out[:] = np.interp(xi, x, data)\n", - "\n", - "\n", - "def xr_interp(data, dim, newdim):\n", - " interped = xr.apply_ufunc(\n", - " interp1d_np_gufunc, # first the function\n", - " data, # now arguments in the order expected by 'interp1_np'\n", - " data[dim], # as above\n", - " newdim, # as above\n", - " input_core_dims=[[dim], [dim], [\"__newdim__\"]], # list with one entry per arg\n", - " output_core_dims=[[\"__newdim__\"]], # returned data has one dimension\n", - " exclude_dims=set((dim,)), # dimensions allowed to change size. Must be a set!\n", - " # vectorize=True, # not needed since numba takes care of vectorizing\n", - " dask=\"parallelized\",\n", - " output_dtypes=[\n", - " data.dtype\n", - " ], # one per output; could also be float or np.dtype(\"float64\")\n", - " ).rename({\"__newdim__\": dim})\n", - " interped[dim] = newdim # need to add this manually\n", - "\n", - " return interped\n", - "\n", - "\n", - "xr.testing.assert_allclose(\n", - " expected.transpose(*interped.dims),\n", - " xr_interp(air.chunk({\"time\": 2, \"lon\": 2}), \"lat\", newlat),\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This technique is generalizable to any 1D function." - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.8.10" - }, - "org": null, - "toc": { - "base_numbering": 1, - "nav_menu": {}, - "number_sections": false, - "sideBar": true, - "skip_h1_title": false, - "title_cell": "Table of Contents", - "title_sidebar": "Contents", - "toc_cell": false, - "toc_position": {}, - "toc_section_display": true, - "toc_window_display": true - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/doc/gallery.rst b/doc/gallery.rst index 9e0cb22da6b..c625c91461d 100644 --- a/doc/gallery.rst +++ b/doc/gallery.rst @@ -24,7 +24,6 @@ Notebook Examples examples/visualization_gallery examples/ROMS_ocean_model examples/ERA5-GRIB-example - examples/apply_ufunc_vectorize_1d examples/blank_template diff --git a/doc/gallery.yml b/doc/gallery.yml index 0701e393389..cd6cf46524d 100644 --- a/doc/gallery.yml +++ b/doc/gallery.yml @@ -23,11 +23,11 @@ notebooks-examples: path: examples/ERA5-GRIB-example.html thumbnail: _static/thumbnails/ERA5-GRIB-example.png +external-examples: - title: Applying unvectorized functions with apply_ufunc - path: examples/apply_ufunc_vectorize_1d.html + path: https://tutorial.xarray.dev/advanced/apply_ufunc/example-interp.html thumbnail: _static/logos/Xarray_Logo_RGB_Final.svg -external-examples: - title: Managing raster data with rioxarray path: https://vylor.github.io/rioxarray/stable/examples/examples.html thumbnail: _static/logos/Xarray_Logo_RGB_Final.svg diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 6f040596dd4..a8034651b4b 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -13,8 +13,8 @@ v2026.09.1 (unreleased) New Features ~~~~~~~~~~~~ -- :py:func:`apply_ufunc` supports ``vectorize="shared"``, which only loops over - the broadcast dimensions along which more than one input varies and passes +- :py:func:`apply_ufunc` accepts dimension names for ``vectorize``, e.g. + ``vectorize=["time"]``, to only loop over these broadcast dimensions and pass all other broadcast dimensions to ``func`` in bulk. This is much faster than ``vectorize=True`` if ``func`` can broadcast its inputs, for numpy and dask arrays alike (:pull:`11681`). @@ -349,9 +349,10 @@ Bug Fixes Documentation ~~~~~~~~~~~~~ - Add examples from basic to advanced usage to the :py:func:`apply_ufunc` - docstring, and a section on ``vectorize="shared"`` to the - :doc:`apply_ufunc example notebook `, - which is executed again in the documentation build (:pull:`11681`). + docstring. The ``apply_ufunc`` example notebook was removed in favor of the + `apply_ufunc interpolation tutorial + `_ + (:pull:`11681`). By `Michael Niklas `_. - Add hidden intersphinx inventory entries for methods, properties and attributes at the path where they are defined, e.g. @@ -6578,7 +6579,7 @@ Documentation - Added examples for :py:meth:`DataArray.quantile`, :py:meth:`Dataset.quantile` and ``GroupBy.quantile``. (:pull:`3576`) By `Justus Magin `_. -- Add new :doc:`example notebook ` example notebook demonstrating +- Add new example notebook demonstrating vectorization of a 1D function using :py:func:`apply_ufunc` , dask and numba. By `Deepak Cherian `_. - Added example for :py:func:`~xarray.map_blocks`. (:pull:`3667`) diff --git a/xarray/computation/apply_ufunc.py b/xarray/computation/apply_ufunc.py index 4688f946029..e1b2e1d07c1 100644 --- a/xarray/computation/apply_ufunc.py +++ b/xarray/computation/apply_ufunc.py @@ -701,28 +701,34 @@ def _vectorize(func, signature, output_dtypes, exclude_dims): return func -def _vectorize_shared[T]( +def _vectorize_dims[T]( func: Callable[..., T], signature: _UFuncSignature, + loop_axes: Sequence[int], + nbroadcast: int, output_dtypes: Sequence | None, output_core_sizes: Mapping[Hashable, int], ) -> Callable[..., T]: - """Vectorize ``func`` over the loop dimensions shared by several inputs. + """Vectorize ``func`` over some of the broadcast dimensions. Like :py:class:`numpy.vectorize` with a signature, but ``func`` is only - looped over the loop (non-core) dimensions along which more than one input - varies, i.e. has a size other than 1. All other loop dimensions are passed - to ``func`` in bulk. + looped over the given broadcast (non-core) dimensions. All other broadcast + dimensions are passed to ``func`` in bulk. Parameters ---------- func : callable - Function to vectorize. It must accept inputs with leading loop + Function to vectorize. It must accept inputs with leading broadcast dimensions that broadcast against each other following numpy's rules, - and return outputs with the broadcast loop dimensions followed by the + and return outputs with the broadcast dimensions followed by the output core dimensions. signature : _UFuncSignature Core dimensions signature for the operation. + loop_axes : sequence of int + Positions of the broadcast dimensions to loop over, in loop order. + nbroadcast : int + Number of broadcast dimensions. Inputs may omit leading broadcast + dimensions, which then have size 1. output_dtypes : sequence of dtype or None Dtypes of the outputs. If None, they are taken from the results of ``func``, so they are required if the loop is empty. @@ -746,24 +752,21 @@ def wrapper(*args: Any) -> T: ] broadcast_shape = np.broadcast_shapes(*loop_shapes) nloop = len(broadcast_shape) - # loop dimensions are aligned to the right, as in numpy broadcasting + # broadcast dimensions are aligned to the right, as in numpy broadcasting offsets = [nloop - len(shape) for shape in loop_shapes] - loop_axes = [ - axis - for axis in range(nloop) - if sum( - axis >= offset and shape[axis - offset] != 1 - for shape, offset in zip(loop_shapes, offsets, strict=True) - ) - > 1 + axes = [ + axis - (nbroadcast - nloop) + for axis in loop_axes + # dimensions missing on all inputs have size 1 + if axis >= nbroadcast - nloop ] - if not loop_axes: + if not axes: return func(*args) outputs = None output_key: list[int | slice] = [slice(None)] * nloop - for index in np.ndindex(*(broadcast_shape[axis] for axis in loop_axes)): - loop_index = dict(zip(loop_axes, index, strict=True)) + for index in np.ndindex(*(broadcast_shape[axis] for axis in axes)): + loop_index = dict(zip(axes, index, strict=True)) sliced_args = [] for arg, shape, offset in zip(args, loop_shapes, offsets, strict=True): # select the loop index, or the only element along a size-1 axis @@ -841,7 +844,7 @@ def apply_variable_ufunc( exclude_dims=frozenset(), dask="forbidden", output_dtypes=None, - vectorize: bool | Literal["shared"] = False, + vectorize: bool | tuple[Hashable, ...] = False, keep_attrs="override", dask_gufunc_kwargs=None, ) -> Variable | tuple[Variable, ...]: @@ -924,11 +927,17 @@ def apply_variable_ufunc( f"dimension '{key}' in 'output_core_dims' needs corresponding (dim, size) in 'output_sizes'" ) - if vectorize == "shared": - numpy_func = _vectorize_shared( + if isinstance(vectorize, tuple): + numpy_func = _vectorize_dims( numpy_func, signature, - output_dtypes, + loop_axes=[ + broadcast_dims.index(d) + for d in vectorize + if d in broadcast_dims + ], + nbroadcast=len(broadcast_dims), + output_dtypes=output_dtypes, output_core_sizes={**dim_sizes, **output_sizes}, ) @@ -950,12 +959,16 @@ def func(*arrays): raise ValueError( f"unknown setting for chunked array handling in apply_ufunc: {dask}" ) - elif vectorize == "shared": + elif isinstance(vectorize, tuple): output_sizes = (dask_gufunc_kwargs or {}).get("output_sizes", {}) - func = _vectorize_shared( + func = _vectorize_dims( func, signature, - output_dtypes, + loop_axes=[ + broadcast_dims.index(d) for d in vectorize if d in broadcast_dims + ], + nbroadcast=len(broadcast_dims), + output_dtypes=output_dtypes, output_core_sizes={**dim_sizes, **output_sizes}, ) elif vectorize: @@ -1045,7 +1058,7 @@ def apply_ufunc( input_core_dims: Sequence[Sequence] | None = None, output_core_dims: Sequence[Sequence] | None = ((),), exclude_dims: AbstractSet = frozenset(), - vectorize: bool | Literal["shared"] = False, + vectorize: bool | str | Iterable[Hashable] = False, join: JoinOptions = "exact", dataset_join: str = "exact", dataset_fill_value: object = _NO_FILL_VALUE, @@ -1103,22 +1116,22 @@ def apply_ufunc( will be dropped. Each excluded dimension must also appear in ``input_core_dims`` for at least one argument. Only dimensions listed here are allowed to change size between input and output objects. - vectorize : bool or "shared", optional + vectorize : bool, str or iterable of Hashable, optional If True, then assume ``func`` only takes arrays defined over core dimensions as input and vectorize it automatically with :py:func:`numpy.vectorize`. This option exists for convenience, but is almost always slower than supplying a pre-vectorized function. - If "shared", only loop over the broadcast dimensions along which more - than one input varies, and pass all other broadcast dimensions to - ``func`` in bulk. ``func`` then has to accept inputs with leading + If one or more dimension names, only loop over these broadcast + dimensions, in the given order, and pass all other broadcast dimensions + to ``func`` in bulk. ``func`` then has to accept inputs with leading broadcast dimensions that broadcast against each other following - numpy's rules, e.g. an input of shape ``(t, r, z)`` together with an - input of shape ``(1, z)``, and return outputs with the broadcast - dimensions followed by the output core dimensions. This is much faster - than ``vectorize=True`` if some broadcast dimensions only appear on a - single input. If the loop is empty, ``output_dtypes`` and the sizes of - all output core dimensions (from the inputs or ``output_sizes`` in + numpy's rules, e.g. an input of shape ``(r, z)`` together with an input + of shape ``(1, z)``, and return outputs with the broadcast dimensions + followed by the output core dimensions. This is much faster than + ``vectorize=True`` if ``func`` can handle some broadcast dimensions + itself. If a loop dimension has size 0, ``output_dtypes`` and the sizes + of all output core dimensions (from the inputs or ``output_sizes`` in ``dask_gufunc_kwargs``) are required. join : {"outer", "inner", "left", "right", "exact"}, default: "exact" Method for joining the indexes of the passed objects along each @@ -1349,12 +1362,12 @@ def apply_ufunc( ... ) ... - **Vectorizing over shared dimensions only.** ``vectorize=True`` calls + **Vectorizing over some dimensions only.** ``vectorize=True`` calls ``func`` once for every combination of the broadcast dimensions. If - ``func`` can handle some of them itself, ``vectorize="shared"`` only loops - over the dimensions along which more than one input varies and passes all - others in bulk. Here, ``np.quantile`` computes a single quantile per call - for all ``x`` at once, but the quantile differs along ``time``: + ``func`` can handle some of them itself, pass the dimensions to loop over + instead, all others are passed in bulk. Here, ``np.quantile`` computes a + single quantile per call for all ``x`` at once, but the quantile differs + along ``time``: >>> data = xr.DataArray( ... np.arange(24.0).reshape(2, 3, 4), dims=("time", "x", "sample") @@ -1365,7 +1378,7 @@ def apply_ufunc( ... return np.quantile(a, q.item(), axis=-1) ... >>> xr.apply_ufunc( - ... quantile, data, q, input_core_dims=[["sample"], []], vectorize="shared" + ... quantile, data, q, input_core_dims=[["sample"], []], vectorize=["time"] ... ) Size: 48B array([[ 0.75, 4.75, 8.75], @@ -1401,9 +1414,22 @@ def apply_ufunc( signature = _UFuncSignature(input_core_dims, output_core_dims) - if vectorize not in (True, False, "shared"): - raise ValueError(f"vectorize must be a bool or 'shared', got {vectorize!r}") - if vectorize != "shared": + if isinstance(vectorize, str): + vectorize = (vectorize,) + if isinstance(vectorize, Iterable): + vectorize = tuple(vectorize) + if core_dims := set(vectorize) & signature.all_core_dims: + raise ValueError( + f"vectorize can only loop over broadcast dimensions, but " + f"{core_dims} are core dimensions" + ) + all_dims = set().union(*(getattr(arg, "dims", ()) for arg in args)) + if missing_dims := set(vectorize) - all_dims: + raise ValueError( + f"vectorize got dimensions {missing_dims} that are not present " + f"on any of the arguments" + ) + else: vectorize = bool(vectorize) if exclude_dims: diff --git a/xarray/core/missing.py b/xarray/core/missing.py index 9a82bff0058..debb57e3edf 100644 --- a/xarray/core/missing.py +++ b/xarray/core/missing.py @@ -769,8 +769,8 @@ def interpolate_variable( dask_gufunc_kwargs=dict(output_sizes=output_sizes, allow_rechunk=True), output_dtypes=[dtype], # only loop over the dimensions along which the destination coordinates vary, - # np.vectorize would loop over all of them (GH10683) - vectorize="shared", + # np.vectorize would also loop over all other broadcast dimensions (GH10683) + vectorize=[d for d in var.dims if d in vectorize_dims], keep_attrs=True, ) return result @@ -808,9 +808,9 @@ def _interpnd( the other half are destination coordinates. `data` has the loop dimensions first, followed by the dimensions we - interpolate along. ``apply_ufunc(..., vectorize="shared")`` only loops over - the loop dimensions along which the destination coordinates vary, e.g. ``t`` - for ``da[t, r, z].interp(z=target[t])``, so the destination coordinates have + interpolate along. ``apply_ufunc`` only loops over the loop dimensions along + which the destination coordinates vary, e.g. ``t`` for + ``da[t, r, z].interp(z=target[t])``, so the destination coordinates have size 1 along all remaining loop dimensions. """ n_x = len(coords) // 2 diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py index 49491546c2b..a728e2fa629 100644 --- a/xarray/tests/test_computation.py +++ b/xarray/tests/test_computation.py @@ -1645,7 +1645,19 @@ def test_vectorize_exclude_dims_dask() -> None: assert_identical(expected, actual) -def test_vectorize_shared(use_dask: bool) -> None: +@pytest.mark.parametrize( + "vectorize, expected_shapes", + [ + (["t"], [((3, 1, 5), (1, 2, 5))] * 4), + ("t", [((3, 1, 5), (1, 2, 5))] * 4), + (["t", "r"], [((1, 5), (2, 5))] * 12), + ], +) +def test_vectorize_dims( + use_dask: bool, + vectorize: str | list[str], + expected_shapes: list[tuple[tuple[int, ...], tuple[int, ...]]], +) -> None: rng = np.random.default_rng(0) a = xr.DataArray(rng.random((4, 3, 5)), dims=("t", "r", "z")) b = xr.DataArray(rng.random((4, 2, 5)), dims=("t", "s", "z")) @@ -1662,18 +1674,30 @@ def dot(x: np.ndarray, y: np.ndarray) -> np.ndarray: a, b, input_core_dims=[["z"], ["z"]], - vectorize="shared", + vectorize=vectorize, dask="parallelized", output_dtypes=[float], ).compute() expected = (a * b).sum("z").transpose("t", "r", "s").compute() xr.testing.assert_allclose(actual, expected) - # only "t" is looped over, "r" and "s" are passed in bulk - assert len(shapes) == 4 - assert set(shapes) == {((3, 1, 5), (1, 2, 5))} + # only the given dimensions are looped over, all others are passed in bulk + assert shapes == expected_shapes + + +def test_vectorize_dims_loop_order() -> None: + a = xr.DataArray(np.arange(6.0).reshape(2, 3), dims=("x", "y")) + calls = [] + def func(x: np.ndarray) -> np.ndarray: + calls.append(x.item()) + return x + + actual = apply_ufunc(func, a, vectorize=["y", "x"]) + assert_identical(actual, a) + assert calls == [0.0, 3.0, 1.0, 4.0, 2.0, 5.0] -def test_vectorize_shared_without_loop() -> None: + +def test_vectorize_dims_missing_on_some_inputs() -> None: a = xr.DataArray(np.arange(6.0).reshape(2, 3), dims=("x", "z")) b = xr.DataArray(np.arange(3.0), dims="z") shapes = [] @@ -1682,12 +1706,23 @@ def dot(x: np.ndarray, y: np.ndarray) -> np.ndarray: shapes.append((x.shape, y.shape)) return (x * y).sum(axis=-1) - actual = apply_ufunc(dot, a, b, input_core_dims=[["z"], ["z"]], vectorize="shared") + actual = apply_ufunc(dot, a, b, input_core_dims=[["z"], ["z"]], vectorize=["x"]) assert_identical(actual, xr.DataArray([5.0, 14.0], dims="x")) - assert shapes == [((2, 3), (3,))] + assert shapes == [((3,), (3,))] * 2 + + +def test_vectorize_dims_dataset() -> None: + # the loop dimension is missing on one of the data variables + ds = xr.Dataset( + {"a": (("x", "z"), np.arange(6.0).reshape(2, 3)), "b": ("z", np.arange(3.0))} + ) + actual = apply_ufunc( + np.sum, ds, input_core_dims=[["z"]], kwargs={"axis": -1}, vectorize=["x"] + ) + assert_identical(actual, ds.sum("z")) -def test_vectorize_shared_multiple_outputs() -> None: +def test_vectorize_dims_multiple_outputs() -> None: a = xr.DataArray(np.arange(12.0).reshape(2, 2, 3), dims=("x", "y", "z")) b = xr.DataArray(np.ones((2, 3)), dims=("x", "z")) @@ -1700,13 +1735,13 @@ def func(x: np.ndarray, y: np.ndarray) -> tuple[np.ndarray, np.ndarray]: b, input_core_dims=[["z"], ["z"]], output_core_dims=[[], []], - vectorize="shared", + vectorize=["x"], ) assert_identical(actual_sum, a.sum("z")) assert_identical(actual_max, a.max("z")) -def test_vectorize_shared_empty_loop() -> None: +def test_vectorize_dims_empty_loop() -> None: a = xr.DataArray(np.zeros((0, 3, 2)), dims=("x", "y", "z")) b = xr.DataArray(np.zeros((0, 2)), dims=("x", "z")) @@ -1718,20 +1753,22 @@ def func(x: np.ndarray, y: np.ndarray) -> np.ndarray: a, b, input_core_dims=[["z"], ["z"]], - vectorize="shared", + vectorize=["x"], output_dtypes=[float], ) assert actual.sizes == {"x": 0, "y": 3} assert actual.dtype == float with pytest.raises(ValueError, match=r"empty loop dimension"): - apply_ufunc(func, a, b, input_core_dims=[["z"], ["z"]], vectorize="shared") + apply_ufunc(func, a, b, input_core_dims=[["z"], ["z"]], vectorize=["x"]) -def test_vectorize_invalid() -> None: - a = xr.DataArray([1.0], dims="x") - with pytest.raises(ValueError, match=r"vectorize must be a bool or 'shared'"): - apply_ufunc(np.sum, a, vectorize="all") # type: ignore[arg-type] +def test_vectorize_dims_invalid() -> None: + a = xr.DataArray([[1.0]], dims=("x", "y")) + with pytest.raises(ValueError, match=r"only loop over broadcast dimensions"): + apply_ufunc(np.sum, a, input_core_dims=[["y"]], vectorize=["y"]) + with pytest.raises(ValueError, match=r"not present on any of the arguments"): + apply_ufunc(np.sum, a, input_core_dims=[["y"]], vectorize=["z"]) def test_corr_only_dataarray() -> None: From 2a27f9b79875adc3c83496dc99c393c36374c898 Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Tue, 20 Oct 2026 22:52:56 +0200 Subject: [PATCH 12/18] Restore whats-new entry of a previous release Co-Authored-By: Claude Opus 5.5 --- doc/whats-new.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index a8034651b4b..a4458925866 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -6579,7 +6579,7 @@ Documentation - Added examples for :py:meth:`DataArray.quantile`, :py:meth:`Dataset.quantile` and ``GroupBy.quantile``. (:pull:`3576`) By `Justus Magin `_. -- Add new example notebook demonstrating +- Add new :doc:`example notebook ` example notebook demonstrating vectorization of a 1D function using :py:func:`apply_ufunc` , dask and numba. By `Deepak Cherian `_. - Added example for :py:func:`~xarray.map_blocks`. (:pull:`3667`) From effa06c3cf4627110b39f4626cb022849c3444ec Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Tue, 20 Oct 2026 22:56:43 +0200 Subject: [PATCH 13/18] Shorten apply_ufunc docstring, unlink removed notebook in whats-new Co-Authored-By: Claude Opus 5.5 --- doc/whats-new.rst | 2 +- xarray/computation/apply_ufunc.py | 59 ++++++------------------------- 2 files changed, 12 insertions(+), 49 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index a4458925866..3c6ab0acc93 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -6579,7 +6579,7 @@ Documentation - Added examples for :py:meth:`DataArray.quantile`, :py:meth:`Dataset.quantile` and ``GroupBy.quantile``. (:pull:`3576`) By `Justus Magin `_. -- Add new :doc:`example notebook ` example notebook demonstrating +- Add new ``apply_ufunc_vectorize_1d`` example notebook demonstrating vectorization of a 1D function using :py:func:`apply_ufunc` , dask and numba. By `Deepak Cherian `_. - Added example for :py:func:`~xarray.map_blocks`. (:pull:`3667`) diff --git a/xarray/computation/apply_ufunc.py b/xarray/computation/apply_ufunc.py index e1b2e1d07c1..13e476e0b4c 100644 --- a/xarray/computation/apply_ufunc.py +++ b/xarray/computation/apply_ufunc.py @@ -1122,17 +1122,10 @@ def apply_ufunc( :py:func:`numpy.vectorize`. This option exists for convenience, but is almost always slower than supplying a pre-vectorized function. - If one or more dimension names, only loop over these broadcast - dimensions, in the given order, and pass all other broadcast dimensions - to ``func`` in bulk. ``func`` then has to accept inputs with leading - broadcast dimensions that broadcast against each other following - numpy's rules, e.g. an input of shape ``(r, z)`` together with an input - of shape ``(1, z)``, and return outputs with the broadcast dimensions - followed by the output core dimensions. This is much faster than - ``vectorize=True`` if ``func`` can handle some broadcast dimensions - itself. If a loop dimension has size 0, ``output_dtypes`` and the sizes - of all output core dimensions (from the inputs or ``output_sizes`` in - ``dask_gufunc_kwargs``) are required. + If dimension names, only loop over these broadcast dimensions and pass + all others to ``func`` in bulk, so ``func`` must broadcast its inputs + following numpy's rules. Loop dimensions of size 0 require + ``output_dtypes`` and known sizes of all output core dimensions. join : {"outer", "inner", "left", "right", "exact"}, default: "exact" Method for joining the indexes of the passed objects along each dimension, and the variables of Dataset objects with mismatched @@ -1222,8 +1215,6 @@ def apply_ufunc( Examples -------- - The examples below go from basic to advanced usage. - **Element-wise functions.** Calculate the vector magnitude of two arguments: >>> def magnitude(a, b): @@ -1254,10 +1245,8 @@ def apply_ufunc( Coordinates: * x (x) float64 24B 0.1 0.2 0.3 - **Core dimensions.** Functions that operate along a dimension, e.g. a - reduction, need ``input_core_dims``. Core dimensions are moved to the end - of the arrays passed to ``func``, so here we reduce over ``axis=-1``, like - ``.mean("y")``: + **Core dimensions** are moved to the end of the arrays passed to ``func``. + Reduce over ``y`` like ``.mean("y")``: >>> da = xr.DataArray( ... np.arange(6.0).reshape(2, 3), coords={"x": [10, 20], "y": [0, 1, 2]} @@ -1278,9 +1267,8 @@ def apply_ufunc( ... return xr.apply_ufunc(_inner, a, b, input_core_dims=[[dim], [dim]]) ... - **Output core dimensions.** If ``func`` returns a dimension, list it in - ``output_core_dims``. If its size changes, it also has to be listed in - ``exclude_dims``, and its coordinates are dropped: + **Output core dimensions** that change size must be listed in + ``exclude_dims``: >>> xr.apply_ufunc( ... np.diff, @@ -1313,8 +1301,7 @@ def apply_ufunc( ... return result ... - **Multiple outputs.** A function returning a tuple needs one entry in - ``output_core_dims`` per output: + **Multiple outputs** need one entry in ``output_core_dims`` each: >>> quotient, remainder = xr.apply_ufunc( ... np.divmod, da, 4, output_core_dims=[[], []] @@ -1327,23 +1314,6 @@ def apply_ufunc( * x (x) int64 16B 10 20 * y (y) int64 24B 0 1 2 - **Dask arrays.** With ``dask="parallelized"``, ``func`` is applied to each - block of a chunked array. Core dimensions must not be split into several - chunks: - - >>> xr.apply_ufunc( - ... np.mean, - ... da.chunk(x=1), - ... input_core_dims=[["y"]], - ... kwargs={"axis": -1}, - ... dask="parallelized", - ... output_dtypes=[float], - ... ).compute() - Size: 16B - array([1., 4.]) - Coordinates: - * x (x) int64 16B 10 20 - **Vectorizing functions.** If your function is not vectorized but can be applied only to core dimensions, you can use ``vectorize=True`` to turn into a vectorized function. This wraps :py:func:`numpy.vectorize`, so the @@ -1362,12 +1332,8 @@ def apply_ufunc( ... ) ... - **Vectorizing over some dimensions only.** ``vectorize=True`` calls - ``func`` once for every combination of the broadcast dimensions. If - ``func`` can handle some of them itself, pass the dimensions to loop over - instead, all others are passed in bulk. Here, ``np.quantile`` computes a - single quantile per call for all ``x`` at once, but the quantile differs - along ``time``: + **Vectorizing over some dimensions only.** ``np.quantile`` handles all + ``x`` at once, but takes a single quantile, so only loop over ``time``: >>> data = xr.DataArray( ... np.arange(24.0).reshape(2, 3, 4), dims=("time", "x", "sample") @@ -1385,9 +1351,6 @@ def apply_ufunc( [14.25, 18.25, 22.25]]) Dimensions without coordinates: time, x - This calls ``quantile`` twice, once per ``time``, while ``vectorize=True`` - would call it six times, once per ``time`` and ``x``. - Most of NumPy's builtin functions already broadcast their inputs appropriately for use in ``apply_ufunc``. You may find helper functions such as :py:func:`numpy.broadcast_arrays` helpful in writing your function. ``apply_ufunc`` also From cd602af21d2901137f14e02ed9b170c9a6b3e8c8 Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Tue, 20 Oct 2026 23:34:10 +0200 Subject: [PATCH 14/18] Keep vectorize_dims ordered and pass it directly Co-Authored-By: Claude Opus 5.5 --- xarray/core/missing.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/xarray/core/missing.py b/xarray/core/missing.py index debb57e3edf..4fe869463ff 100644 --- a/xarray/core/missing.py +++ b/xarray/core/missing.py @@ -723,7 +723,10 @@ def interpolate_variable( # we are explicitly interpolating. These are the dimensions along which we vectorize. # `q` is the only input core dimensions, and changes size (disappears) # so it is in exclude_dims. - vectorize_dims = (result_dims - all_in_core_dims) & set(var.dims) + # ordered like var.dims, since it also determines the loop order + vectorize_dims = [ + d for d in var.dims if d in result_dims and d not in all_in_core_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) @@ -770,7 +773,7 @@ def interpolate_variable( output_dtypes=[dtype], # only loop over the dimensions along which the destination coordinates vary, # np.vectorize would also loop over all other broadcast dimensions (GH10683) - vectorize=[d for d in var.dims if d in vectorize_dims], + vectorize=vectorize_dims, keep_attrs=True, ) return result From 7f84169517d6d68dcac3f08e69f551c443ce038b Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Tue, 20 Oct 2026 23:34:42 +0200 Subject: [PATCH 15/18] Clarify comment on vectorize Co-authored-by: Deepak Cherian Co-Authored-By: Claude Opus 5.5 --- xarray/core/missing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/core/missing.py b/xarray/core/missing.py index 4fe869463ff..7af1af27046 100644 --- a/xarray/core/missing.py +++ b/xarray/core/missing.py @@ -772,7 +772,7 @@ def interpolate_variable( dask_gufunc_kwargs=dict(output_sizes=output_sizes, allow_rechunk=True), output_dtypes=[dtype], # only loop over the dimensions along which the destination coordinates vary, - # np.vectorize would also loop over all other broadcast dimensions (GH10683) + # vectorize=True would also loop over all other broadcast dimensions (GH10683) vectorize=vectorize_dims, keep_attrs=True, ) From 27ce6b465798e8d639c4324d24d03c058544f5a9 Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Tue, 20 Oct 2026 23:44:44 +0200 Subject: [PATCH 16/18] Test vectorize=[] Co-Authored-By: Claude Opus 5.5 --- xarray/tests/test_computation.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py index a728e2fa629..3baefb92d9d 100644 --- a/xarray/tests/test_computation.py +++ b/xarray/tests/test_computation.py @@ -1684,6 +1684,21 @@ def dot(x: np.ndarray, y: np.ndarray) -> np.ndarray: assert shapes == expected_shapes +def test_vectorize_dims_empty() -> None: + # no loop dimensions, same as vectorize=False + a = xr.DataArray(np.arange(6.0).reshape(2, 3), dims=("x", "z")) + b = xr.DataArray(np.arange(3.0), dims="z") + shapes = [] + + def dot(x: np.ndarray, y: np.ndarray) -> np.ndarray: + shapes.append((x.shape, y.shape)) + return (x * y).sum(axis=-1) + + actual = apply_ufunc(dot, a, b, input_core_dims=[["z"], ["z"]], vectorize=[]) + assert_identical(actual, xr.DataArray([5.0, 14.0], dims="x")) + assert shapes == [((2, 3), (3,))] + + def test_vectorize_dims_loop_order() -> None: a = xr.DataArray(np.arange(6.0).reshape(2, 3), dims=("x", "y")) calls = [] From fff68f58eb9ed2c39f0e897d61d8d8095850d06a Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Tue, 20 Oct 2026 23:47:01 +0200 Subject: [PATCH 17/18] Use an interpolation example for vectorize=[dims] Co-Authored-By: Claude Opus 5.5 --- xarray/computation/apply_ufunc.py | 42 +++++++++++++++++++++---------- 1 file changed, 29 insertions(+), 13 deletions(-) diff --git a/xarray/computation/apply_ufunc.py b/xarray/computation/apply_ufunc.py index 13e476e0b4c..f0ba7e0bef3 100644 --- a/xarray/computation/apply_ufunc.py +++ b/xarray/computation/apply_ufunc.py @@ -1332,24 +1332,40 @@ def apply_ufunc( ... ) ... - **Vectorizing over some dimensions only.** ``np.quantile`` handles all - ``x`` at once, but takes a single quantile, so only loop over ``time``: + **Vectorizing over some dimensions only.** ``make_interp_spline`` + interpolates all ``x`` at once, but only to a single set of new points. + These vary along ``time``, so only loop over ``time``: - >>> data = xr.DataArray( - ... np.arange(24.0).reshape(2, 3, 4), dims=("time", "x", "sample") + >>> from scipy.interpolate import make_interp_spline + >>> da = xr.DataArray( + ... np.arange(24.0).reshape(2, 3, 4), + ... dims=("time", "x", "z"), + ... coords={"z": [0, 1, 2, 3]}, ... ) - >>> q = xr.DataArray([0.25, 0.75], dims="time") - >>> def quantile(a, q): - ... # a has shape (x, sample), q has size 1 along x - ... return np.quantile(a, q.item(), axis=-1) + >>> new_z = xr.DataArray([[0.5, 1.5], [1.0, 2.5]], dims=("time", "new_z")) + >>> def interp(data, z, new_z): + ... # data has shape (x, z), new_z has size 1 along x + ... return make_interp_spline(z, data, k=1, axis=-1)(new_z.reshape(-1)) ... >>> xr.apply_ufunc( - ... quantile, data, q, input_core_dims=[["sample"], []], vectorize=["time"] + ... interp, + ... da, + ... da.z, + ... new_z, + ... input_core_dims=[["z"], ["z"], ["new_z"]], + ... output_core_dims=[["new_z"]], + ... exclude_dims={"z"}, + ... vectorize=["time"], ... ) - Size: 48B - array([[ 0.75, 4.75, 8.75], - [14.25, 18.25, 22.25]]) - Dimensions without coordinates: time, x + Size: 96B + array([[[ 0.5, 1.5], + [ 4.5, 5.5], + [ 8.5, 9.5]], + + [[13. , 14.5], + [17. , 18.5], + [21. , 22.5]]]) + Dimensions without coordinates: time, x, new_z Most of NumPy's builtin functions already broadcast their inputs appropriately for use in ``apply_ufunc``. You may find helper functions such as From dda83a0fc9429d3c1eaa3c095d3e436be2846928 Mon Sep 17 00:00:00 2001 From: Michael Niklas Date: Thu, 22 Oct 2026 09:24:11 +0200 Subject: [PATCH 18/18] Localize each block for chunked destination coordinates `_localize` skips chunked destination coordinates, so without localizing per block scipy casts the whole block to float for integer data. Add a peakmem benchmark for this case. Co-authored-by: Claude --- asv_bench/benchmarks/interp.py | 32 ++++++++++++++++++++++++++++++++ xarray/core/missing.py | 20 ++++++++++++++++++-- 2 files changed, 50 insertions(+), 2 deletions(-) diff --git a/asv_bench/benchmarks/interp.py b/asv_bench/benchmarks/interp.py index 52b521698f5..a524e259ddd 100644 --- a/asv_bench/benchmarks/interp.py +++ b/asv_bench/benchmarks/interp.py @@ -1,8 +1,11 @@ +from typing import Any + import numpy as np import pandas as pd from asv_runner.benchmarks.mark import parameterize import xarray as xr +from xarray.core.types import InterpOptions from . import randn, requires_dask @@ -91,3 +94,32 @@ def setup(self, *args, **kwargs): requires_dask() super().setup(**kwargs) self.da = self.da.chunk({"t": 25, "r": 25}) + + +class InterpolationDaskChunkedCoords: + # scipy casts integer data to float, which copies the whole block unless + # it is localized to the range of the (chunked) destination coordinates + def setup(self, *args: Any, **kwargs: Any) -> None: + requires_dask() + import dask.array + + ny, nx = 16, 2_000_000 + chunks = (8, nx) # 122 MiB blocks, generated lazily + self.data = { + "int": dask.array.random.randint(0, 100, (ny, nx), chunks=chunks), + "float": dask.array.random.random((ny, nx), chunks=chunks), + } + self.x = np.arange(nx) + self.new_x = xr.DataArray( + dask.array.linspace(1000, 1500, 400, chunks=100), dims="z" + ) + + @parameterize({"method": ["linear", "nearest"], "dtype": ["int", "float"]}) + def peakmem_interpolation_chunked_coords( + self, method: InterpOptions, dtype: str + ) -> None: + import dask + + da = xr.DataArray(self.data[dtype], dims=("y", "x"), coords={"x": self.x}) + with dask.config.set(scheduler="synchronous"): + da.interp(x=self.new_x, method=method).compute() diff --git a/xarray/core/missing.py b/xarray/core/missing.py index 7af1af27046..487735ce961 100644 --- a/xarray/core/missing.py +++ b/xarray/core/missing.py @@ -698,13 +698,20 @@ def interpolate_variable( else: func, kwargs = _get_interpolator_nd(method, **kwargs) - if kwargs.get("method") in ["linear", "nearest"]: + localize = kwargs.get("method") in ["linear", "nearest"] + if localize: # simple speed up for the local interpolation var, indexes_coords = _localize(var, dict(indexes_coords)) in_coords, result_coords = _floatize_x( *(list(_) for _ in zip(*indexes_coords.values(), strict=True)) ) + # `_localize` skips chunked destination coordinates, so localize each block + # instead: scipy casts integer data to float, which would otherwise copy + # the whole block + localize_blocks = localize and any( + is_chunked_array(coord._data) for coord in result_coords + ) # input coordinates along which we are interpolation are core dimensions # the corresponding output coordinates may or may not have the same name, @@ -767,6 +774,7 @@ def interpolate_variable( # on the method and the coordinate dtypes, so cast the result to # the dtype we promise dask dtype=dtype, + localize=localize_blocks, ), # TODO: deprecate and have the user rechunk themselves dask_gufunc_kwargs=dict(output_sizes=output_sizes, allow_rechunk=True), @@ -801,9 +809,10 @@ def _interpnd( data: np.ndarray, *coords: np.ndarray, interp_func: Interpolator | InterpCallable, - interp_kwargs, + interp_kwargs: dict[str, Any], result_coord_core_dims: list[tuple[Hashable, ...]], dtype: np.dtype, + localize: bool = False, ) -> np.ndarray: """ Core nD array interpolation routine. @@ -839,6 +848,13 @@ def _interpnd( ) var = Variable([f"dim_{dim}" for dim in range(ndim)], data, fastpath=True) + if localize: + indexes_coords = { + _x.dims[0]: (_x, _new_x) for _x, _new_x in zip(x, new_x, strict=True) + } + var, indexes_coords = _localize(var, indexes_coords) + x, new_x = (list(_) for _ in zip(*indexes_coords.values(), strict=True)) + if len(x) == 1: # TODO: narrow interp_func to interpolator here return astype(