np.interp : An end-to-end example

np.interp : An end-to-end example#

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.

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)

Learning goals#

Our goal is to use apply_ufunc with a general function so that we can reuse our code to apply to different xarray datasets or along different dimensions. Specifically, this example will illustrate

Specifying core dimensions with input_core_dims
Handling core dimensions of the output with output_core_dims
Handling core dimensions that change size using exclude_dims
Automatic vectorizing or looping over dimensions that are not core dimensions using vectorize=True
Automatically parallelization with dask arrays using dask="parallelized"
High-performance vectorization with numba and vectorize=False.

It puts together all the concepts covered earlier.

Tip

We’ll reduce the length of error messages using in this tutorial using %xmode minimal See the ipython documentation for details.

`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: 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. The data returned was:

array([296.29    , 296.195455, 296.100909, 296.006364, 295.911818, 296.048485,
       296.218182, 296.387879, 296.557576, 296.672727, 296.769697, 296.866667,
       296.963636, 296.757576, 296.369697, 295.981818, 295.593939, 295.204848,
       294.814545, 294.424242, 294.033939, 293.727273, 293.56    , 293.392727,
       293.225455, 292.924242, 292.221212, 291.518182, 290.815152, 290.130303,
       289.572727, 289.015152, 288.457576, 287.9     , 287.560606, 287.221212,
       286.881818, 286.542424, 286.09697 , 285.636364, 285.175758, 284.715152,
       284.270909, 283.832121, 283.393333, 282.954545, 282.367273, 281.690909,
       281.014545, 280.338182, 279.806061, 279.418182, 279.030303, 278.642424,
       278.299091, 278.03    , 277.760909, 277.491818, 277.254242, 277.111212,
       276.968182, 276.825152, 276.675758, 276.481818, 276.287879, 276.093939,
       275.9     , 275.630909, 275.361818, 275.092727, 274.823636, 274.558788,
       274.294545, 274.030303, 273.766061, 273.409091, 273.021212, 272.633333,
       272.245455, 272.463636, 273.045455, 273.627273, 274.209091, 273.530303,
       271.590909, 269.651515, 267.712121, 265.      , 261.      , 257.      ,
       253.      , 249.624242, 248.121212, 246.618182, 245.115152, 243.721212,
       243.090909, 242.460606, 241.830303, 241.2     ])

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: 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: applied function returned data with an unexpected number of dimensions. Received 1 dimension(s) but expected 0 dimensions with names (), from:

array([296.29    , 296.195455, 296.100909, 296.006364, 295.911818, 296.048485,
       296.218182, 296.387879, 296.557576, 296.672727, 296.769697, 296.866667,
       296.963636, 296.757576, 296.369697, 295.981818, 295.593939, 295.204848,
       294.814545, 294.424242, 294.033939, 293.727273, 293.56    , 293.392727,
       293.225455, 292.924242, 292.221212, 291.518182, 290.815152, 290.130303,
       289.572727, 289.015152, 288.457576, 287.9     , 287.560606, 287.221212,
       286.881818, 286.542424, 286.09697 , 285.636364, 285.175758, 284.715152,
       284.270909, 283.832121, 283.393333, 282.954545, 282.367273, 281.690909,
       281.014545, 280.338182, 279.806061, 279.418182, 279.030303, 278.642424,
       278.299091, 278.03    , 277.760909, 277.491818, 277.254242, 277.111212,
       276.968182, 276.825152, 276.675758, 276.481818, 276.287879, 276.093939,
       275.9     , 275.630909, 275.361818, 275.092727, 274.823636, 274.558788,
       274.294545, 274.030303, 273.766061, 273.409091, 273.021212, 272.633333,
       272.245455, 272.463636, 273.045455, 273.627273, 274.209091, 273.530303,
       271.590909, 269.651515, 267.712121, 265.      , 261.      , 257.      ,
       253.      , 249.624242, 248.121212, 246.618182, 245.115152, 243.721212,
       243.090909, 242.460606, 241.830303, 241.2     ])

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)> Size: 800B
296.3 296.2 296.1 296.0 295.9 296.0 ... 245.1 243.7 243.1 242.5 241.8 241.2
Coordinates:
    lon      float32 4B 200.0
    time     datetime64[ns] 8B 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!

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:

output_dtypes: dtypes of all returned objects, and
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. Any decorated function gets compiled and will loop over any non-core dimension in parallel when necessary.

We need to specify some extra information:

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

Summary#

This technique is generalizable to any 1D function that can be compiled by Numba.

np.interp : An end-to-end example

Contents

np.interp : An end-to-end example#

Learning goals#

Load data#

`apply_ufunc`#

`exclude_dims`#

Core dimensions#

Automatic vectorization with `np.vectorize`#

Parallelization with dask#

High performance vectorization: gufuncs, numba & guvectorize#

Summary#

np.interp : An end-to-end example

Contents

np.interp : An end-to-end example#

Learning goals#

Load data#

apply_ufunc#

exclude_dims#

Core dimensions#

Automatic vectorization with np.vectorize#

Parallelization with dask#

High performance vectorization: gufuncs, numba & guvectorize#

Summary#

`apply_ufunc`#

`exclude_dims`#

Automatic vectorization with `np.vectorize`#