# Applying unvectorized functions#

Author Deepak Cherian (NCAR)

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 coveniently apply this function along a dimension of xarray objects that may or may not wrap dask arrays with a signature.

We will illustrate this using np.interp:

Signature: np.interp(x, xp, fp, left=None, right=None, period=None)
Docstring:
One-dimensional linear interpolation.

Returns the one-dimensional piecewise linear interpolant to a function
with given discrete data points (xp, fp), evaluated at x.


and write an xr_interp function with signature

xr_interp(xarray_object, dimension_name, new_coordinate_to_interpolate_to)


## Load data#

First lets load an example dataset

import xarray as xr
import numpy as np

xr.set_options(display_style="html")  # fancy HTML repr

air = (
xr.tutorial.load_dataset("air_temperature")
.air.sortby("lat")  # np.interp needs coordinate in ascending order
.isel(time=slice(4), lon=slice(3))
)  # choose a small subset for convenience
air

<xarray.DataArray 'air' (time: 4, lat: 25, lon: 3)>
array([[[296.29   , 296.79   , 297.1    ],
[295.9    , 296.19998, 296.79   ],
[296.6    , 296.19998, 296.4    ],
[297.     , 296.69998, 296.1    ],
[295.4    , 295.69998, 295.79   ],
[293.79   , 294.1    , 294.6    ],
[293.1    , 293.29   , 293.29   ],
[290.19998, 290.79   , 291.4    ],
[287.9    , 288.     , 288.29   ],
[286.5    , 286.5    , 285.69998],
[284.6    , 284.9    , 284.19998],
[282.79   , 283.19998, 282.6    ],
[280.     , 280.69998, 280.19998],
[278.4    , 279.     , 279.     ],
[277.29   , 277.4    , 277.79   ],
[276.69998, 277.4    , 277.69998],
[275.9    , 276.9    , 276.9    ],
[274.79   , 275.19998, 275.6    ],
[273.69998, 273.6    , 273.79   ],
[272.1    , 270.9    , 270.     ],
...
[293.     , 293.5    , 294.29   ],
[291.9    , 291.9    , 292.19998],
[289.19998, 289.4    , 289.9    ],
[286.6    , 287.1    , 287.9    ],
[284.79   , 284.79   , 285.4    ],
[282.79   , 282.     , 282.69998],
[281.19998, 280.19998, 280.6    ],
[279.5    , 278.69998, 278.6    ],
[278.     , 277.69998, 277.6    ],
[276.4    , 275.9    , 276.4    ],
[275.6    , 275.69998, 276.1    ],
[274.5    , 275.6    , 276.29   ],
[273.4    , 274.5    , 275.5    ],
[274.1    , 274.     , 273.5    ],
[273.29   , 272.6    , 271.5    ],
[272.79   , 272.4    , 271.9    ],
[267.69998, 266.29   , 264.4    ],
[256.6    , 254.7    , 252.09999],
[246.29999, 245.29999, 244.2    ],
[241.89   , 241.79999, 241.79999]]], dtype=float32)
Coordinates:
* lat      (lat) float32 15.0 17.5 20.0 22.5 25.0 ... 65.0 67.5 70.0 72.5 75.0
* lon      (lon) float32 200.0 202.5 205.0
* time     (time) datetime64[ns] 2013-01-01 ... 2013-01-01T18:00:00
Attributes:
long_name:     4xDaily Air temperature at sigma level 995
units:         degK
precision:     2
GRIB_id:       11
GRIB_name:     TMP
var_desc:      Air temperature
dataset:       NMC Reanalysis
level_desc:    Surface
statistic:     Individual Obs
parent_stat:   Other
actual_range:  [185.16 322.1 ]

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.

newlat = np.linspace(15, 75, 100)
air.interp(lat=newlat)

<xarray.DataArray 'air' (time: 4, lat: 100, lon: 3)>
array([[[296.29000854, 296.79000854, 297.1000061 ],
[296.19545954, 296.64697173, 297.02485518],
[296.10091053, 296.50393491, 296.94970426],
...,
[242.46059851, 243.46969695, 244.08181672],
[241.83029767, 242.98484846, 243.79090837],
[241.19999683, 242.49999997, 243.50000003]],

[[296.29000854, 297.19998169, 297.3999939 ],
[296.26818385, 297.07876957, 297.25211866],
[296.24635916, 296.95755744, 297.10424342],
...,
[242.82726354, 243.37878187, 243.63332714],
[242.46362716, 243.03938941, 243.36665899],
[242.09999079, 242.69999695, 243.09999084]],

[[296.3999939 , 296.29000854, 296.3999939 ],
[296.35150609, 296.34091556, 296.37333078],
[296.30301828, 296.39182258, 296.34666767],
...,
[243.4151408 , 243.26181628, 243.12423612],
[242.85756431, 242.73090659, 242.71211194],
[242.29998782, 242.19999689, 242.29998776]],

[[297.5       , 297.69998169, 297.5       ],
[297.37878788, 297.65150128, 297.40303179],
[297.25757575, 297.60302087, 297.30606357],
...,
[244.02817552, 243.49695752, 242.96362858],
[242.9590874 , 242.64847269, 242.38180817],
[241.88999927, 241.79998785, 241.79998776]]])
Coordinates:
* lon      (lon) float32 200.0 202.5 205.0
* time     (time) datetime64[ns] 2013-01-01 ... 2013-01-01T18:00:00
* lat      (lat) float64 15.0 15.61 16.21 16.82 ... 73.18 73.79 74.39 75.0
Attributes:
long_name:     4xDaily Air temperature at sigma level 995
units:         degK
precision:     2
GRIB_id:       11
GRIB_name:     TMP
var_desc:      Air temperature
dataset:       NMC Reanalysis
level_desc:    Surface
statistic:     Individual Obs
parent_stat:   Other
actual_range:  [185.16 322.1 ]

Let’s define a function that works with one vector of data along lat at a time.

def interp1d_np(data, x, xi):
return np.interp(xi, x, data)

interped = interp1d_np(air.isel(time=0, lon=0), air.lat, newlat)
expected = air.interp(lat=newlat)

# no errors are raised if values are equal to within floating point precision
np.testing.assert_allclose(expected.isel(time=0, lon=0).values, interped)


No errors are raised so our interpolation is working.

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.

## apply_ufunc#

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.


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.

xr.apply_ufunc(
interp1d_np,  # first the function
air.isel(time=0, lon=0),  # now arguments in the order expected by 'interp1_np'
air.lat,
newlat,
)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[4], line 1
----> 1 xr.apply_ufunc(
2     interp1d_np,  # first the function
3     air.isel(time=0, lon=0),  # now arguments in the order expected by 'interp1_np'
4     air.lat,
5     newlat,
6 )

File /usr/share/miniconda3/envs/xarray-tutorial/lib/python3.10/site-packages/xarray/core/computation.py:1204, in apply_ufunc(func, input_core_dims, output_core_dims, exclude_dims, vectorize, join, dataset_join, dataset_fill_value, keep_attrs, kwargs, dask, output_dtypes, output_sizes, meta, dask_gufunc_kwargs, *args)
1202 # feed DataArray apply_variable_ufunc through apply_dataarray_vfunc
1203 elif any(isinstance(a, DataArray) for a in args):
-> 1204     return apply_dataarray_vfunc(
1205         variables_vfunc,
1206         *args,
1207         signature=signature,
1208         join=join,
1209         exclude_dims=exclude_dims,
1210         keep_attrs=keep_attrs,
1211     )
1212 # feed Variables directly through apply_variable_ufunc
1213 elif any(isinstance(a, Variable) for a in args):

File /usr/share/miniconda3/envs/xarray-tutorial/lib/python3.10/site-packages/xarray/core/computation.py:315, in apply_dataarray_vfunc(func, signature, join, exclude_dims, keep_attrs, *args)
310 result_coords, result_indexes = build_output_coords_and_indexes(
311     args, signature, exclude_dims, combine_attrs=keep_attrs
312 )
314 data_vars = [getattr(a, "variable", a) for a in args]
--> 315 result_var = func(*data_vars)
317 out: tuple[DataArray, ...] | DataArray
318 if signature.num_outputs > 1:

File /usr/share/miniconda3/envs/xarray-tutorial/lib/python3.10/site-packages/xarray/core/computation.py:805, in apply_variable_ufunc(func, signature, exclude_dims, dask, output_dtypes, vectorize, keep_attrs, dask_gufunc_kwargs, *args)
803 for dim, new_size in var.sizes.items():
804     if dim in dim_sizes and new_size != dim_sizes[dim]:
--> 805         raise ValueError(
806             "size of dimension {!r} on inputs was unexpectedly "
807             "changed by applied function from {} to {}. Only "
808             "dimensions specified in exclude_dims with "
809             "xarray.apply_ufunc are allowed to change size.".format(
810                 dim, dim_sizes[dim], new_size
811             )
812         )
814 var.attrs = attrs
815 output.append(var)

ValueError: size of dimension 'lat' on inputs was unexpectedly changed by applied function from 25 to 100. Only dimensions specified in exclude_dims with xarray.apply_ufunc are allowed to change size.


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.

## exclude_dims#

exclude_dims : 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.

xr.apply_ufunc(
interp1d_np,  # first the function
air.isel(time=0, lon=0),  # now arguments in the order expected by 'interp1_np'
air.lat,
newlat,
exclude_dims=set(("lat",)),  # dimensions allowed to change size. Must be set!
)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[5], line 1
----> 1 xr.apply_ufunc(
2     interp1d_np,  # first the function
3     air.isel(time=0, lon=0),  # now arguments in the order expected by 'interp1_np'
4     air.lat,
5     newlat,
6     exclude_dims=set(("lat",)),  # dimensions allowed to change size. Must be set!
7 )

File /usr/share/miniconda3/envs/xarray-tutorial/lib/python3.10/site-packages/xarray/core/computation.py:1120, in apply_ufunc(func, input_core_dims, output_core_dims, exclude_dims, vectorize, join, dataset_join, dataset_fill_value, keep_attrs, kwargs, dask, output_dtypes, output_sizes, meta, dask_gufunc_kwargs, *args)
1116         raise TypeError(
1117             f"Expected exclude_dims to be a 'set'. Received '{type(exclude_dims).__name__}' instead."
1118         )
1119     if not exclude_dims <= signature.all_core_dims:
-> 1120         raise ValueError(
1121             f"each dimension in exclude_dims must also be a "
1122             f"core dimension in the function signature. "
1123             f"Please make {(exclude_dims - signature.all_core_dims)} a core dimension"
1124         )
1126 # handle dask_gufunc_kwargs
1127 if dask == "parallelized":

ValueError: each dimension in exclude_dims must also be a core dimension in the function signature. Please make {'lat'} a core dimension


## Core dimensions#

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.

input_core_dims : Sequence[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 [2]_.

output_core_dims : List[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.


Next we specify "lat" as input_core_dims on both air and air.lat

xr.apply_ufunc(
interp1d_np,  # first the function
air.isel(time=0, lon=0),  # now arguments in the order expected by 'interp1_np'
air.lat,
newlat,
input_core_dims=[["lat"], ["lat"], []],
exclude_dims=set(("lat",)),  # dimensions allowed to change size. Must be set!
)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[6], line 1
----> 1 xr.apply_ufunc(
2     interp1d_np,  # first the function
3     air.isel(time=0, lon=0),  # now arguments in the order expected by 'interp1_np'
4     air.lat,
5     newlat,
6     input_core_dims=[["lat"], ["lat"], []],
7     exclude_dims=set(("lat",)),  # dimensions allowed to change size. Must be set!
8 )

File /usr/share/miniconda3/envs/xarray-tutorial/lib/python3.10/site-packages/xarray/core/computation.py:1204, in apply_ufunc(func, input_core_dims, output_core_dims, exclude_dims, vectorize, join, dataset_join, dataset_fill_value, keep_attrs, kwargs, dask, output_dtypes, output_sizes, meta, dask_gufunc_kwargs, *args)
1202 # feed DataArray apply_variable_ufunc through apply_dataarray_vfunc
1203 elif any(isinstance(a, DataArray) for a in args):
-> 1204     return apply_dataarray_vfunc(
1205         variables_vfunc,
1206         *args,
1207         signature=signature,
1208         join=join,
1209         exclude_dims=exclude_dims,
1210         keep_attrs=keep_attrs,
1211     )
1212 # feed Variables directly through apply_variable_ufunc
1213 elif any(isinstance(a, Variable) for a in args):

File /usr/share/miniconda3/envs/xarray-tutorial/lib/python3.10/site-packages/xarray/core/computation.py:315, in apply_dataarray_vfunc(func, signature, join, exclude_dims, keep_attrs, *args)
310 result_coords, result_indexes = build_output_coords_and_indexes(
311     args, signature, exclude_dims, combine_attrs=keep_attrs
312 )
314 data_vars = [getattr(a, "variable", a) for a in args]
--> 315 result_var = func(*data_vars)
317 out: tuple[DataArray, ...] | DataArray
318 if signature.num_outputs > 1:

File /usr/share/miniconda3/envs/xarray-tutorial/lib/python3.10/site-packages/xarray/core/computation.py:796, in apply_variable_ufunc(func, signature, exclude_dims, dask, output_dtypes, vectorize, keep_attrs, dask_gufunc_kwargs, *args)
794 data = as_compatible_data(data)
795 if data.ndim != len(dims):
--> 796     raise ValueError(
797         "applied function returned data with unexpected "
798         f"number of dimensions. Received {data.ndim} dimension(s) but "
799         f"expected {len(dims)} dimensions with names: {dims!r}"
800     )
802 var = Variable(dims, data, fastpath=True)
803 for dim, new_size in var.sizes.items():

ValueError: applied function returned data with unexpected number of dimensions. Received 1 dimension(s) but expected 0 dimensions with names: ()


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

xr.apply_ufunc(
interp1d_np,  # first the function
air.isel(time=0, lon=0),  # now arguments in the order expected by 'interp1_np'
air.lat,
newlat,
input_core_dims=[["lat"], ["lat"], []],  # list with one entry per arg
output_core_dims=[["lat"]],
exclude_dims=set(("lat",)),  # dimensions allowed to change size. Must be set!
)

<xarray.DataArray (lat: 100)>
array([296.29000854, 296.19545954, 296.10091053, 296.00636153,
295.91181252, 296.04848133, 296.21818126, 296.38788119,
296.55758112, 296.67273227, 296.76970048, 296.8666687 ,
296.96363692, 296.75757483, 296.36969457, 295.9818143 ,
295.59393403, 295.20484416, 294.81454468, 294.4242452 ,
294.03394572, 293.72728105, 293.56000773, 293.39273441,
293.22546109, 292.92424705, 292.22121083, 291.5181746 ,
290.81513838, 290.13028509, 289.57271229, 289.01513949,
288.45756669, 287.8999939 , 287.56060144, 287.22120898,
286.88181652, 286.54242406, 286.09697099, 285.63636641,
285.17576183, 284.71515725, 284.27091564, 283.83212835,
283.39334106, 282.95455378, 282.36727998, 281.69091427,
281.01454856, 280.33818285, 279.80605987, 279.4181796 ,
279.03029933, 278.64241906, 278.29908614, 278.02999878,
277.76091142, 277.49182406, 277.25424934, 277.11121253,
276.96817571, 276.8251389 , 276.67573964, 276.4818032 ,
276.28786677, 276.09393033, 275.8999939 , 275.63090654,
275.36181918, 275.09273182, 274.82364446, 274.55879073,
274.29454179, 274.03029286, 273.76604392, 273.40907704,
273.02120417, 272.6333313 , 272.24545843, 272.46364154,
273.04545824, 273.62727495, 274.20909165, 273.53030303,
271.59090909, 269.65151515, 267.71212121, 265.        ,
261.        , 257.        , 253.        , 249.62424168,
248.12120842, 246.61817516, 245.1151419 , 243.72120019,
243.09089938, 242.46059857, 241.83029776, 241.19999695])
Coordinates:
lon      float32 200.0
time     datetime64[ns] 2013-01-01
Dimensions without coordinates: lat

Finally we get some output! Let’s check that this is right

interped = xr.apply_ufunc(
interp1d_np,  # first the function
air.isel(time=0, lon=0),  # now arguments in the order expected by 'interp1_np'
air.lat,
newlat,
input_core_dims=[["lat"], ["lat"], []],  # list with one entry per arg
output_core_dims=[["lat"]],
exclude_dims=set(("lat",)),  # dimensions allowed to change size. Must be set!
)
interped["lat"] = newlat  # need to add this manually
xr.testing.assert_allclose(expected.isel(time=0, lon=0), interped)


No errors are raised so it is right!

## Vectorization with np.vectorize#

Now our function currently only works on one vector of data which is not so useful given our 3D dataset. Let’s try passing the whole dataset. We add a print statement so we can see what our function receives.

def interp1d_np(data, x, xi):
print(f"data: {data.shape} | x: {x.shape} | xi: {xi.shape}")
return np.interp(xi, x, data)

interped = xr.apply_ufunc(
interp1d_np,  # first the function
air.isel(lon=slice(3), time=slice(4)),  # now arguments in the order expected by 'interp1_np'
air.lat,
newlat,
input_core_dims=[["lat"], ["lat"], []],  # list with one entry per arg
output_core_dims=[["lat"]],
exclude_dims=set(("lat",)),  # dimensions allowed to change size. Must be set!
)
interped["lat"] = newlat  # need to add this manually
xr.testing.assert_allclose(expected.isel(time=0, lon=0), interped)

data: (4, 3, 25) | x: (25,) | xi: (100,)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[9], line 6
2     print(f"data: {data.shape} | x: {x.shape} | xi: {xi.shape}")
3     return np.interp(xi, x, data)
----> 6 interped = xr.apply_ufunc(
7     interp1d_np,  # first the function
8     air.isel(lon=slice(3), time=slice(4)),  # now arguments in the order expected by 'interp1_np'
9     air.lat,
10     newlat,
11     input_core_dims=[["lat"], ["lat"], []],  # list with one entry per arg
12     output_core_dims=[["lat"]],
13     exclude_dims=set(("lat",)),  # dimensions allowed to change size. Must be set!
14 )
15 interped["lat"] = newlat  # need to add this manually
16 xr.testing.assert_allclose(expected.isel(time=0, lon=0), interped)

File /usr/share/miniconda3/envs/xarray-tutorial/lib/python3.10/site-packages/xarray/core/computation.py:1204, in apply_ufunc(func, input_core_dims, output_core_dims, exclude_dims, vectorize, join, dataset_join, dataset_fill_value, keep_attrs, kwargs, dask, output_dtypes, output_sizes, meta, dask_gufunc_kwargs, *args)
1202 # feed DataArray apply_variable_ufunc through apply_dataarray_vfunc
1203 elif any(isinstance(a, DataArray) for a in args):
-> 1204     return apply_dataarray_vfunc(
1205         variables_vfunc,
1206         *args,
1207         signature=signature,
1208         join=join,
1209         exclude_dims=exclude_dims,
1210         keep_attrs=keep_attrs,
1211     )
1212 # feed Variables directly through apply_variable_ufunc
1213 elif any(isinstance(a, Variable) for a in args):

File /usr/share/miniconda3/envs/xarray-tutorial/lib/python3.10/site-packages/xarray/core/computation.py:315, in apply_dataarray_vfunc(func, signature, join, exclude_dims, keep_attrs, *args)
310 result_coords, result_indexes = build_output_coords_and_indexes(
311     args, signature, exclude_dims, combine_attrs=keep_attrs
312 )
314 data_vars = [getattr(a, "variable", a) for a in args]
--> 315 result_var = func(*data_vars)
317 out: tuple[DataArray, ...] | DataArray
318 if signature.num_outputs > 1:

File /usr/share/miniconda3/envs/xarray-tutorial/lib/python3.10/site-packages/xarray/core/computation.py:771, in apply_variable_ufunc(func, signature, exclude_dims, dask, output_dtypes, vectorize, keep_attrs, dask_gufunc_kwargs, *args)
766     if vectorize:
767         func = _vectorize(
768             func, signature, output_dtypes=output_dtypes, exclude_dims=exclude_dims
769         )
--> 771 result_data = func(*input_data)
773 if signature.num_outputs == 1:
774     result_data = (result_data,)

Cell In[9], line 3, in interp1d_np(data, x, xi)
1 def interp1d_np(data, x, xi):
2     print(f"data: {data.shape} | x: {x.shape} | xi: {xi.shape}")
----> 3     return np.interp(xi, x, data)

File <__array_function__ internals>:180, in interp(*args, **kwargs)

File /usr/share/miniconda3/envs/xarray-tutorial/lib/python3.10/site-packages/numpy/lib/function_base.py:1594, in interp(x, xp, fp, left, right, period)
1591     xp = np.concatenate((xp[-1:]-period, xp, xp[0:1]+period))
1592     fp = np.concatenate((fp[-1:], fp, fp[0:1]))
-> 1594 return interp_func(x, xp, fp, left, right)

ValueError: object too deep for desired array


That’s a hard-to-interpret error but our print call helpfully printed the shapes of the input data:

data: (10, 53, 25) | x: (25,) | xi: (100,)


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. apply_ufunc makes this easy by specifying vectorize=True:

vectorize : bool, 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.
Using this option requires NumPy version 1.12 or newer.


Also see the documentation for np.vectorize: https://numpy.org/doc/stable/reference/generated/numpy.vectorize.html. Most importantly

The vectorize function is provided primarily for convenience, not for performance.
The implementation is essentially a for loop.

def interp1d_np(data, x, xi):
print(f"data: {data.shape} | x: {x.shape} | xi: {xi.shape}")
return np.interp(xi, x, data)

interped = xr.apply_ufunc(
interp1d_np,  # first the function
air,  # now arguments in the order expected by 'interp1_np'
air.lat,  # as above
newlat,  # as above
input_core_dims=[["lat"], ["lat"], []],  # list with one entry per arg
output_core_dims=[["lat"]],  # returned data has one dimension
exclude_dims=set(("lat",)),  # dimensions allowed to change size. Must be set!
vectorize=True,  # loop over non-core dims
)
interped["lat"] = newlat  # need to add this manually
xr.testing.assert_allclose(expected, interped)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[10], line 6
2     print(f"data: {data.shape} | x: {x.shape} | xi: {xi.shape}")
3     return np.interp(xi, x, data)
----> 6 interped = xr.apply_ufunc(
7     interp1d_np,  # first the function
8     air,  # now arguments in the order expected by 'interp1_np'
9     air.lat,  # as above
10     newlat,  # as above
11     input_core_dims=[["lat"], ["lat"], []],  # list with one entry per arg
12     output_core_dims=[["lat"]],  # returned data has one dimension
13     exclude_dims=set(("lat",)),  # dimensions allowed to change size. Must be set!
14     vectorize=True,  # loop over non-core dims
15 )
16 interped["lat"] = newlat  # need to add this manually
17 xr.testing.assert_allclose(expected, interped)

File /usr/share/miniconda3/envs/xarray-tutorial/lib/python3.10/site-packages/xarray/core/computation.py:1204, in apply_ufunc(func, input_core_dims, output_core_dims, exclude_dims, vectorize, join, dataset_join, dataset_fill_value, keep_attrs, kwargs, dask, output_dtypes, output_sizes, meta, dask_gufunc_kwargs, *args)
1202 # feed DataArray apply_variable_ufunc through apply_dataarray_vfunc
1203 elif any(isinstance(a, DataArray) for a in args):
-> 1204     return apply_dataarray_vfunc(
1205         variables_vfunc,
1206         *args,
1207         signature=signature,
1208         join=join,
1209         exclude_dims=exclude_dims,
1210         keep_attrs=keep_attrs,
1211     )
1212 # feed Variables directly through apply_variable_ufunc
1213 elif any(isinstance(a, Variable) for a in args):

File /usr/share/miniconda3/envs/xarray-tutorial/lib/python3.10/site-packages/xarray/core/computation.py:315, in apply_dataarray_vfunc(func, signature, join, exclude_dims, keep_attrs, *args)
310 result_coords, result_indexes = build_output_coords_and_indexes(
311     args, signature, exclude_dims, combine_attrs=keep_attrs
312 )
314 data_vars = [getattr(a, "variable", a) for a in args]
--> 315 result_var = func(*data_vars)
317 out: tuple[DataArray, ...] | DataArray
318 if signature.num_outputs > 1:

File /usr/share/miniconda3/envs/xarray-tutorial/lib/python3.10/site-packages/xarray/core/computation.py:771, in apply_variable_ufunc(func, signature, exclude_dims, dask, output_dtypes, vectorize, keep_attrs, dask_gufunc_kwargs, *args)
766     if vectorize:
767         func = _vectorize(
768             func, signature, output_dtypes=output_dtypes, exclude_dims=exclude_dims
769         )
--> 771 result_data = func(*input_data)
773 if signature.num_outputs == 1:
774     result_data = (result_data,)

File /usr/share/miniconda3/envs/xarray-tutorial/lib/python3.10/site-packages/numpy/lib/function_base.py:2328, in vectorize.__call__(self, *args, **kwargs)
2325     vargs = [args[_i] for _i in inds]
2326     vargs.extend([kwargs[_n] for _n in names])
-> 2328 return self._vectorize_call(func=func, args=vargs)

File /usr/share/miniconda3/envs/xarray-tutorial/lib/python3.10/site-packages/numpy/lib/function_base.py:2402, in vectorize._vectorize_call(self, func, args)
2400 """Vectorized call to func over positional args."""
2401 if self.signature is not None:
-> 2402     res = self._vectorize_call_with_signature(func, args)
2403 elif not args:
2404     res = func()

File /usr/share/miniconda3/envs/xarray-tutorial/lib/python3.10/site-packages/numpy/lib/function_base.py:2430, in vectorize._vectorize_call_with_signature(self, func, args)
2425     raise TypeError('wrong number of positional arguments: '
2426                     'expected %r, got %r'
2427                     % (len(input_core_dims), len(args)))
2428 args = tuple(asanyarray(arg) for arg in args)
-> 2430 broadcast_shape, dim_sizes = _parse_input_dimensions(
2431     args, input_core_dims)
2432 input_shapes = _calculate_shapes(broadcast_shape, dim_sizes,
2433                                  input_core_dims)
2434 args = [np.broadcast_to(arg, shape, subok=True)
2435         for arg, shape in zip(args, input_shapes)]

File /usr/share/miniconda3/envs/xarray-tutorial/lib/python3.10/site-packages/numpy/lib/function_base.py:2090, in _parse_input_dimensions(args, input_core_dims)
2088     dummy_array = np.lib.stride_tricks.as_strided(0, arg.shape[:ndim])
2089     broadcast_args.append(dummy_array)
-> 2090 broadcast_shape = np.lib.stride_tricks._broadcast_shape(*broadcast_args)
2091 return broadcast_shape, dim_sizes

File /usr/share/miniconda3/envs/xarray-tutorial/lib/python3.10/site-packages/numpy/lib/stride_tricks.py:422, in _broadcast_shape(*args)
417 """Returns the shape of the arrays that would result from broadcasting the
418 supplied arrays against each other.
419 """
420 # use the old-iterator because np.nditer does not handle size 0 arrays
421 # consistently
--> 422 b = np.broadcast(*args[:32])
423 # unfortunately, it cannot handle 32 or more arguments directly
424 for pos in range(32, len(args), 31):
425     # ironically, np.broadcast does not properly handle np.broadcast
426     # objects (it treats them as scalars)
427     # use broadcasting to avoid allocating the full array

ValueError: shape mismatch: objects cannot be broadcast to a single shape.  Mismatch is between arg 0 with shape (4, 3) and arg 2 with shape (100,).


This unfortunately is another cryptic error from numpy.

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.

def interp1d_np(data, x, xi):
print(f"data: {data.shape} | x: {x.shape} | xi: {xi.shape}")
return np.interp(xi, x, data)

interped = xr.apply_ufunc(
interp1d_np,  # first the function
air,  # now arguments in the order expected by 'interp1_np'
air.lat,  # as above
newlat,  # as above
input_core_dims=[["lat"], ["lat"], ["new_lat"]],  # list with one entry per arg
output_core_dims=[["new_lat"]],  # returned data has one dimension
exclude_dims=set(("lat",)),  # dimensions allowed to change size. Must be a set!
vectorize=True,  # loop over non-core dims
)
interped = interped.rename({"new_lat": "lat"})
interped["lat"] = newlat  # need to add this manually
xr.testing.assert_allclose(
expected.transpose(*interped.dims), interped  # order of dims is different
)
interped

data: (25,) | x: (25,) | xi: (100,)
data: (25,) | x: (25,) | xi: (100,)
data: (25,) | x: (25,) | xi: (100,)
data: (25,) | x: (25,) | xi: (100,)
data: (25,) | x: (25,) | xi: (100,)
data: (25,) | x: (25,) | xi: (100,)
data: (25,) | x: (25,) | xi: (100,)
data: (25,) | x: (25,) | xi: (100,)
data: (25,) | x: (25,) | xi: (100,)
data: (25,) | x: (25,) | xi: (100,)
data: (25,) | x: (25,) | xi: (100,)
data: (25,) | x: (25,) | xi: (100,)

<xarray.DataArray (time: 4, lon: 3, lat: 100)>
array([[[296.29000854, 296.19545954, 296.10091053, ..., 242.46059857,
241.83029776, 241.19999695],
[296.79000854, 296.64697173, 296.50393492, ..., 243.46969697,
242.98484848, 242.5       ],
[297.1000061 , 297.02485518, 296.94970426, ..., 244.0818167 ,
243.79090835, 243.5       ]],

[[296.29000854, 296.26818385, 296.24635916, ..., 242.82726357,
242.46362721, 242.09999084],
[297.19998169, 297.07876957, 296.95755745, ..., 243.37878187,
243.03938941, 242.69999695],
[297.3999939 , 297.25211866, 297.10424342, ..., 243.63332714,
243.36665899, 243.09999084]],

[[296.3999939 , 296.35150609, 296.30301828, ..., 243.41514079,
242.85756429, 242.29998779],
[296.29000854, 296.34091556, 296.39182258, ..., 243.26181631,
242.73090663, 242.19999695],
[296.3999939 , 296.37333078, 296.34666767, ..., 243.12423614,
242.71211196, 242.29998779]],

[[297.5       , 297.37878788, 297.25757576, ..., 244.02817559,
242.95908749, 241.88999939],
[297.69998169, 297.65150128, 297.60302087, ..., 243.49695749,
242.64847264, 241.79998779],
[297.5       , 297.40303178, 297.30606357, ..., 242.9636286 ,
242.38180819, 241.79998779]]])
Coordinates:
* lon      (lon) float32 200.0 202.5 205.0
* time     (time) datetime64[ns] 2013-01-01 ... 2013-01-01T18:00:00
* lat      (lat) float64 15.0 15.61 16.21 16.82 ... 73.18 73.79 74.39 75.0

Notice that the printed input shapes are all 1D and correspond to one vector along the lat dimension.

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.

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

    Core dimensions are automatically moved to the last axes of input
variables before applying func, which facilitates using NumPy style
generalized ufuncs [2]_.


## Parallelization with dask#

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.

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:

1. output_dtypes: dtypes of all returned objects, and

2. output_sizes: lengths of any new dimensions.

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.

def interp1d_np(data, x, xi):
print(f"data: {data.shape} | x: {x.shape} | xi: {xi.shape}")
return np.interp(xi, x, data)

interped = xr.apply_ufunc(
interp1d_np,  # first the function
air.chunk({"time": 2, "lon": 2}),  # now arguments in the order expected by 'interp1_np'
air.lat,  # as above
newlat,  # as above
input_core_dims=[["lat"], ["lat"], ["new_lat"]],  # list with one entry per arg
output_core_dims=[["new_lat"]],  # returned data has one dimension
exclude_dims=set(("lat",)),  # dimensions allowed to change size. Must be a set!
vectorize=True,  # loop over non-core dims
dask="parallelized",
output_dtypes=[air.dtype],  # one per output
).rename({"new_lat": "lat"})
interped["lat"] = newlat  # need to add this manually
xr.testing.assert_allclose(expected.transpose(*interped.dims), interped)

data: (25,) | x: (25,) | xi: (100,)
data: (25,) | x: (25,) | xi: (100,)
data: (25,) | x: (25,) | xi: (100,)
data: (25,) | x: (25,) | xi: (100,)
data: (25,) | x: (25,) | xi: (100,)
data: (25,) | x: (25,) | xi: (100,)
data: (25,) | x: (25,) | xi: (100,)
data: (25,) | x: (25,) | xi: (100,)
data: (25,) | x: (25,) | xi: (100,)
data: (25,) | x: (25,) | xi: (100,)
data: (25,) | x: (25,) | xi: (100,)
data: (25,) | x: (25,) | xi: (100,)


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.

## High performance vectorization: gufuncs, numba & guvectorize#

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.

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

Any decorated function gets compiled and will loop over any non-core dimension in parallel when necessary. We need to specify some extra information:

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.

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

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), (m) -> (m)." corresponding again to data, x, xi, out

from numba import float64, guvectorize

@guvectorize("(float64[:], float64[:], float64[:], float64[:])", "(n), (n), (m) -> (m)")
def interp1d_np_gufunc(data, x, xi, out):
# numba doesn't really like this.
# seem to support fstrings so do it the old way
print("data: " + str(data.shape) + " | x:" + str(x.shape) + " | xi: " + str(xi.shape))
out[:] = np.interp(xi, x, data)
# gufuncs don't return data
# instead you assign to a the last arg
# return np.interp(xi, x, data)

/tmp/ipykernel_2777/409442179.py:4: NumbaWarning:
Compilation is falling back to object mode WITHOUT looplifting enabled because Function "interp1d_np_gufunc" failed type inference due to: No implementation of function Function(<class 'str'>) found for signature:

>>> str(UniTuple(int64 x 1))

There are 10 candidate implementations:
- Of which 10 did not match due to:
Overload of function 'str': File: <numerous>: Line N/A.
With argument(s): '(UniTuple(int64 x 1))':
No match.

During: resolving callee type: Function(<class 'str'>)
During: typing of call at /tmp/ipykernel_2777/409442179.py (8)

File "../../../../../../../tmp/ipykernel_2777/409442179.py", line 8:
<source missing, REPL/exec in use?>

@guvectorize("(float64[:], float64[:], float64[:], float64[:])", "(n), (n), (m) -> (m)")
/usr/share/miniconda3/envs/xarray-tutorial/lib/python3.10/site-packages/numba/core/object_mode_passes.py:151: NumbaWarning: Function "interp1d_np_gufunc" was compiled in object mode without forceobj=True.

File "../../../../../../../tmp/ipykernel_2777/409442179.py", line 4:
<source missing, REPL/exec in use?>

warnings.warn(errors.NumbaWarning(warn_msg,
/usr/share/miniconda3/envs/xarray-tutorial/lib/python3.10/site-packages/numba/core/object_mode_passes.py:161: NumbaDeprecationWarning:
Fall-back from the nopython compilation path to the object mode compilation path has been detected, this is deprecated behaviour.

For more information visit https://numba.readthedocs.io/en/stable/reference/deprecation.html#deprecation-of-object-mode-fall-back-behaviour-when-using-jit

File "../../../../../../../tmp/ipykernel_2777/409442179.py", line 4:
<source missing, REPL/exec in use?>

warnings.warn(errors.NumbaDeprecationWarning(msg,


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.

interped = xr.apply_ufunc(
interp1d_np_gufunc,  # first the function
air.chunk({"time": 2, "lon": 2}),  # now arguments in the order expected by 'interp1_np'
air.lat,  # as above
newlat,  # as above
input_core_dims=[["lat"], ["lat"], ["new_lat"]],  # list with one entry per arg
output_core_dims=[["new_lat"]],  # returned data has one dimension
exclude_dims=set(("lat",)),  # dimensions allowed to change size. Must be a set!
# vectorize=True,  # not needed since numba takes care of vectorizing
dask="parallelized",
output_dtypes=[air.dtype],  # one per output
).rename({"new_lat": "lat"})
interped["lat"] = newlat  # need to add this manually
xr.testing.assert_allclose(expected.transpose(*interped.dims), interped)

data: (25,) | x:(25,) | xi: (100,)
data: (25,) | x:(25,) | xi: (100,)
data: (25,) | x:(25,) | xi: (100,)
data: (25,) | x:(25,) | xi: (100,)
data: (25,) | x:(25,) | xi: (100,)
data: (25,) | x:(25,) | xi: (100,)
data: (25,) | x:(25,) | xi: (100,)
data: (25,) | x:(25,) | xi: (100,)
data: (25,) | x:(25,) | xi: (100,)
data: (25,) | x:(25,) | xi: (100,)
data: (25,) | x:(25,) | xi: (100,)
data: (25,) | x:(25,) | xi: (100,)


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

from numba import float64, guvectorize

@guvectorize(
"(float64[:], float64[:], float64[:], float64[:])",
"(n), (n), (m) -> (m)",
nopython=True,
)
def interp1d_np_gufunc(data, x, xi, out):
out[:] = np.interp(xi, x, data)

def xr_interp(data, dim, newdim):
interped = xr.apply_ufunc(
interp1d_np_gufunc,  # first the function
data,  # now arguments in the order expected by 'interp1_np'
data[dim],  # as above
newdim,  # as above
input_core_dims=[[dim], [dim], ["__newdim__"]],  # list with one entry per arg
output_core_dims=[["__newdim__"]],  # returned data has one dimension
exclude_dims=set((dim,)),  # dimensions allowed to change size. Must be a set!
# vectorize=True,  # not needed since numba takes care of vectorizing
dask="parallelized",
output_dtypes=[data.dtype],  # one per output; could also be float or np.dtype("float64")
).rename({"__newdim__": dim})
interped[dim] = newdim  # need to add this manually

return interped

xr.testing.assert_allclose(
expected.transpose(*interped.dims),
xr_interp(air.chunk({"time": 2, "lon": 2}), "lat", newlat),
)


This technique is generalizable to any 1D function.