# xarray.apply_ufunc

### xarray.apply_ufunc(func, \*args, input_core_dims=None, output_core_dims=((), ), exclude_dims=frozenset({}), vectorize=False, join='exact', dataset_join='exact', dataset_fill_value=<no-fill-value>, keep_attrs=None, kwargs=None, dask='forbidden', output_dtypes=None, output_sizes=None, meta=None, dask_gufunc_kwargs=None, on_missing_core_dim='raise')[[source]](https://github.com/pydata/xarray/blob/main/xarray/computation/apply_ufunc.py#L896-L1305)

Apply a vectorized function for unlabeled arrays on xarray objects.

The function will be mapped over the data variable(s) of the input
arguments using xarray’s standard rules for labeled computation, including
alignment, broadcasting, looping over GroupBy/Dataset variables, and
merging of coordinates.

* **Parameters:**
  * **func** ([`callable()`](https://docs.python.org/3/builtins/functions.html#callable)) – Function to call like `func(*args, **kwargs)` on unlabeled arrays
    (`.data`) that returns an array or tuple of arrays. If multiple
    arguments with non-matching dimensions are supplied, this function is
    expected to vectorize (broadcast) over axes of positional arguments in
    the style of NumPy universal functions <sup>[1](#id3)</sup> (if this is not the case,
    set `vectorize=True`). If this function returns multiple outputs, you
    must set `output_core_dims` as well.
  * **\*args** ([`DataTree`](xarray.DataTree.html.md#xarray.DataTree), [`Dataset`](xarray.Dataset.html.md#xarray.Dataset), [`DataArray`](xarray.DataArray.html.md#xarray.DataArray), [`DataArrayGroupBy`](xarray.core.groupby.DataArrayGroupBy.html.md#xarray.core.groupby.DataArrayGroupBy), [`DatasetGroupBy`](xarray.core.groupby.DatasetGroupBy.html.md#xarray.core.groupby.DatasetGroupBy),         Variable, [`numpy.ndarray`](https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html#numpy.ndarray), [`dask.array.Array`](https://docs.dask.org/en/latest/generated/dask.array.Array.html#dask.array.Array) or [scalar](../user-guide/terminology.html.md#term-scalar)) – Mix of labeled and/or unlabeled arrays to which to apply the function.
  * **input_core_dims** ([sequence](https://docs.python.org/3/glossary.html#term-sequence) of [sequence](https://docs.python.org/3/glossary.html#term-sequence), *optional*) – List of the same length as `args` giving the list of core dimensions
    on each input argument that should not be broadcast. By default, we
    assume there are no core dimensions on any input arguments.

    For example, `input_core_dims=[[], ['time']]` indicates that all
    dimensions on the first argument and all dimensions other than ‘time’
    on the second argument should be broadcast.

    Core dimensions are automatically moved to the last axes of input
    variables before applying `func`, which facilitates using NumPy style
    generalized ufuncs <sup>[2](#id4)</sup>.
  * **output_core_dims** ([`list`](https://docs.python.org/3/builtins/stdtypes.html#list) of [`tuple`](https://docs.python.org/3/builtins/stdtypes.html#tuple), *optional*) – List of the same length as the number of output arguments from
    `func`, giving the list of core dimensions on each output that were
    not broadcast on the inputs. By default, we assume that `func`
    outputs exactly one array, with axes corresponding to each broadcast
    dimension.

    Core dimensions are assumed to appear as the last dimensions of each
    output in the provided order.
  * **exclude_dims** ([`set`](https://docs.python.org/3/builtins/stdtypes.html#set), *optional*) – Core dimensions on the inputs to exclude from alignment and
    broadcasting entirely. Any input coordinates along these dimensions
    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`](https://docs.python.org/3/builtins/functions.html#bool), *optional*) – If True, then assume `func` only takes arrays defined over core
    dimensions as input and vectorize it automatically with
    `numpy.vectorize()`. This option exists for convenience, but is
    almost always slower than supplying a pre-vectorized function.
  * **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
    data variables:
    - ‘outer’: use the union of object indexes
    - ‘inner’: use the intersection of object indexes
    - ‘left’: use indexes from the first object with each dimension
    - ‘right’: use indexes from the last object with each dimension
    - ‘exact’: raise ValueError instead of aligning when indexes to be
      aligned are not equal
  * **dataset_join** (`{"outer", "inner", "left", "right", "exact"}`, *default*: `"exact"`) – Method for joining variables of Dataset objects with mismatched
    data variables.
    - ‘outer’: take variables from both Dataset objects
    - ‘inner’: take only overlapped variables
    - ‘left’: take only variables from the first object
    - ‘right’: take only variables from the last object
    - ‘exact’: data variables on all Dataset objects must match exactly
  * **dataset_fill_value** (**optional**) – Value used in place of missing variables on Dataset inputs when the
    datasets do not share the exact same `data_vars`. Required if
    `dataset_join not in {'inner', 'exact'}`, otherwise ignored.
  * **keep_attrs** (`{"drop", "identical", "no_conflicts", "drop_conflicts", "override"}` or [`bool`](https://docs.python.org/3/builtins/functions.html#bool), *optional*) –
    - ‘drop’ or False: empty attrs on returned xarray object.
    - ‘identical’: all attrs must be the same on every object.
    - ‘no_conflicts’: attrs from all objects are combined, any that have the same name must also have the same value.
    - ‘drop_conflicts’: attrs from all objects are combined, any that have the same name but different values are dropped.
    - ‘override’ or True: skip comparing and copy attrs from the first object to the result.
  * **kwargs** ([`dict`](https://docs.python.org/3/builtins/stdtypes.html#dict), *optional*) – Optional keyword arguments passed directly on to call `func`.
  * **dask** (`{"forbidden", "allowed", "parallelized"}`, *default*: `"forbidden"`) – How to handle applying to objects containing lazy data in the form of
    dask arrays:
    - ‘forbidden’ (default): raise an error if a dask array is encountered.
    - ‘allowed’: pass dask arrays directly on to `func`. Prefer this option if
      `func` natively supports dask arrays.
    - ‘parallelized’: automatically parallelize `func` if any of the
      inputs are a dask array by using `dask.array.apply_gufunc()`. Multiple output
      arguments are supported. Only use this option if `func` does not natively
      support dask arrays (e.g. converts them to numpy arrays).
  * **dask_gufunc_kwargs** ([`dict`](https://docs.python.org/3/builtins/stdtypes.html#dict), *optional*) – Optional keyword arguments passed to `dask.array.apply_gufunc()` if
    dask=’parallelized’. Possible keywords are `output_sizes`, `allow_rechunk`
    and `meta`.
  * **output_dtypes** ([`list`](https://docs.python.org/3/builtins/stdtypes.html#list) of [`dtype`](https://numpy.org/doc/stable/reference/generated/numpy.dtype.html#numpy.dtype), *optional*) – Optional list of output dtypes. Only used if `dask='parallelized'` or
    `vectorize=True`.
  * **output_sizes** ([`dict`](https://docs.python.org/3/builtins/stdtypes.html#dict), *optional*) – Optional mapping from dimension names to sizes for outputs. Only used
    if dask=’parallelized’ and new dimensions (not found on inputs) appear
    on outputs. `output_sizes` should be given in the `dask_gufunc_kwargs`
    parameter. It will be removed as direct parameter in a future version.
  * **meta** (**optional**) – Size-0 object representing the type of array wrapped by dask array. Passed on to
    `dask.array.apply_gufunc()`. `meta` should be given in the
    `dask_gufunc_kwargs` parameter . It will be removed as direct parameter
    a future version.
  * **on_missing_core_dim** (`{"raise", "copy", "drop"}`, *default*: `"raise"`) – How to handle missing core dimensions on input variables.
* **Returns:**
  * `Single value` or [`tuple`](https://docs.python.org/3/builtins/stdtypes.html#tuple) of [`DataTree`](xarray.DataTree.html.md#xarray.DataTree), [`Dataset`](xarray.Dataset.html.md#xarray.Dataset), [`DataArray`](xarray.DataArray.html.md#xarray.DataArray), `Variable,`
  * [`dask.array.Array`](https://docs.dask.org/en/latest/generated/dask.array.Array.html#dask.array.Array) or [`numpy.ndarray`](https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html#numpy.ndarray), `the first type on that list` to `appear on`
  * `an input. For DataTree inputs`, `func` is applied to `the datasets at each`
  * `node` and `the results are rebuilt into trees with the same structure.`

### Notes

This function is designed for the more common case where `func` can work on numpy
arrays. If `func` needs to manipulate a whole xarray object subset to each block
it is possible to use [`xarray.map_blocks()`](xarray.map_blocks.html.md#xarray.map_blocks).

Note that due to the overhead [`xarray.map_blocks()`](xarray.map_blocks.html.md#xarray.map_blocks) is considerably slower than `apply_ufunc`.

[Parallelize custom functions with apply_ufunc and map_blocks](../user-guide/dask.html.md#dask-automatic-parallelization)
: User guide describing [`apply_ufunc()`](#xarray.apply_ufunc) and [`map_blocks()`](xarray.map_blocks.html.md#xarray.map_blocks).

[apply_ufunc](https://tutorial.xarray.dev/advanced/apply_ufunc/apply_ufunc.html)
: Advanced Tutorial on applying numpy function using [`apply_ufunc()`](#xarray.apply_ufunc)

### Examples

<div class="try_examples_button_container"><button class="try_examples_button" onclick="window.tryExamplesShowIframe('736ebcb6-d3c8-4aac-b48b-7dda7f57e3bd','4277a9c7-69e8-4cb3-8801-4645a874f8f7','776a67d2-dd3e-44b9-9b4b-8f9e438d6036','../lite/tree/../notebooks/index.html?path=231575ae_88e7_4689_9296_9159a58aa9fe.ipynb','None')">Try it in your browser!</button></div>

Calculate the vector magnitude of two arguments:

```pycon
>>> def magnitude(a, b):
...     func = lambda x, y: np.sqrt(x**2 + y**2)
...     return xr.apply_ufunc(func, a, b)
...
```

You can now apply `magnitude()` to [`DataArray`](xarray.DataArray.html.md#xarray.DataArray) and [`Dataset`](xarray.Dataset.html.md#xarray.Dataset)
objects, with automatically preserved dimensions and coordinates, e.g.,

```pycon
>>> array = xr.DataArray([1, 2, 3], coords=[("x", [0.1, 0.2, 0.3])])
>>> magnitude(array, -array)
<xarray.DataArray (x: 3)> Size: 24B
array([1.41421356, 2.82842712, 4.24264069])
Coordinates:
  * x        (x) float64 24B 0.1 0.2 0.3
```

Plain scalars, numpy arrays and a mix of these with xarray objects is also
supported:

```pycon
>>> magnitude(3, 4)
np.float64(5.0)
>>> magnitude(3, np.array([0, 4]))
array([3., 5.])
>>> magnitude(array, 0)
<xarray.DataArray (x: 3)> Size: 24B
array([1., 2., 3.])
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:

```pycon
>>> 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}
...     )
...
```

Inner product over a specific dimension (like [`dot()`](xarray.dot.html.md#xarray.dot)):

```pycon
>>> def _inner(x, y):
...     result = np.matmul(x[..., np.newaxis, :], y[..., :, np.newaxis])
...     return result[..., 0, 0]
...
>>> def inner_product(a, b, dim):
...     return apply_ufunc(_inner, a, b, input_core_dims=[[dim], [dim]])
...
```

Stack objects along a new dimension (like [`concat()`](xarray.concat.html.md#xarray.concat)):

```pycon
>>> def stack(objects, dim, new_coord):
...     # note: this version does not stack coordinates
...     func = lambda *x: np.stack(x, axis=-1)
...     result = apply_ufunc(
...         func,
...         *objects,
...         output_core_dims=[[dim]],
...         join="outer",
...         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 `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:

```pycon
>>> import scipy.stats
>>> def earth_mover_distance(first_samples, second_samples, dim="ensemble"):
...     return apply_ufunc(
...         scipy.stats.wasserstein_distance,
...         first_samples,
...         second_samples,
...         input_core_dims=[[dim], [dim]],
...         vectorize=True,
...     )
...
```

Most of NumPy’s builtin functions already broadcast their inputs
appropriately for use in `apply_ufunc`. You may find helper functions such as
[`numpy.broadcast_arrays()`](https://numpy.org/doc/stable/reference/generated/numpy.broadcast_arrays.html#numpy.broadcast_arrays) helpful in writing your function. `apply_ufunc` also
works well with [`numba.vectorize()`](https://numba.readthedocs.io/en/stable/reference/jit-compilation.html#numba.vectorize) and [`numba.guvectorize()`](https://numba.readthedocs.io/en/stable/reference/jit-compilation.html#numba.guvectorize).

<div id="776a67d2-dd3e-44b9-9b4b-8f9e438d6036" class="try_examples_outer_iframe  hidden"><div class="try_examples_button_container"><button class="try_examples_button" onclick="window.tryExamplesHideIframe('736ebcb6-d3c8-4aac-b48b-7dda7f57e3bd','776a67d2-dd3e-44b9-9b4b-8f9e438d6036')">Go Back</button><button class="try_examples_button" onclick="window.openInNewTab('736ebcb6-d3c8-4aac-b48b-7dda7f57e3bd','776a67d2-dd3e-44b9-9b4b-8f9e438d6036')">Open In Tab</button></div><div id="4277a9c7-69e8-4cb3-8801-4645a874f8f7" class="jupyterlite_sphinx_iframe_container"></div></div><script>document.addEventListener("DOMContentLoaded", function() {window.loadTryExamplesConfig("../try_examples.json");});</script>

#### SEE ALSO
[`numpy.broadcast_arrays`](https://numpy.org/doc/stable/reference/generated/numpy.broadcast_arrays.html#numpy.broadcast_arrays), [`numba.vectorize`](https://numba.readthedocs.io/en/stable/reference/jit-compilation.html#numba.vectorize), [`numba.guvectorize`](https://numba.readthedocs.io/en/stable/reference/jit-compilation.html#numba.guvectorize), `dask.array.apply_gufunc`, [`xarray.map_blocks`](xarray.map_blocks.html.md#xarray.map_blocks)

### References

* <a id='id3'>**[1]**</a> [https://numpy.org/doc/stable/reference/ufuncs.html](https://numpy.org/doc/stable/reference/ufuncs.html)
* <a id='id4'>**[2]**</a> [https://numpy.org/doc/stable/reference/c-api/generalized-ufuncs.html](https://numpy.org/doc/stable/reference/c-api/generalized-ufuncs.html)