Previously we looked at
applying functions
on numpy arrays, and the concept of
core dimensions
.
We learned that functions commonly support specifying “core dimensions” through the
axis
keyword
argument.
However many functions exist, that implicitly have core dimensions, but do not provide an
axis
keyword
argument. Applying such functions to a nD array usually involves one or multiple loops over the other dimensions
— termed “loop dimensions” or “broadcast dimensions”.
A good example is numpy’s 1D interpolate function
numpy.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`.
This function expects 1D arrays as input, so there is one core dimension and we cannot easily apply
it to a nD array since there is no axis keyword argument.
Our goal here is to
Understand the difference between core dimensions and loop dimensions
Understand vectorization
Learn how to apply such functions without loops using apply_ufunc by providing the vectorize keyword argument.
Consider the example problem of interpolating a 2D array with dimensions space and time along the time dimension.
Which dimension is the core dimension, and which is the “loop dimension”?
Solution
time is the core dimension, and space is the loop dimension.
The pattern of looping over any number of “loop dimensions” and applying a function along “core dimensions”
is so common that numpy provides wrappers that automate these steps:
apply_ufunc provides an easy interface to numpy.vectorize through the keyword argument vectorize. Here we see how to use
that to automatically apply np.interp along a single axis of a nD array
We’ll reduce the length of error messages using %xmodeminimal See the ipython documentation for details.
%xmode minimal
importxarrayasxrimportnumpyasnpxr.set_options(display_expand_data=False)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
We’ll work with the apply_ufunc call from the section on handling dimensions that change size. See the “Handling Complex Output” section for how to get here.
This version only works with 1D vectors. We will expand that to work with inputs of any number of dimensions.
newlat=np.linspace(15,75,100)xr.apply_ufunc(np.interp,# first the functionnewlat,air.lat,air.isel(lon=0,time=0),# this version only works with 1D vectorsinput_core_dims=[["lat"],["lat"],["lat"]],output_core_dims=[["lat"]],exclude_dims={"lat"},
Our goal is to interpolate latitude at every longitude and time, such that we go from a dataset with dimensions (time:4,lat:25,lon:3) to (time:4,lat:100,lon:3).
If we blindly try passing air (a 3D DataArray), we get a hard-to-understand error
newlat=np.linspace(15,75,100)xr.apply_ufunc(np.interp,# first the functionnewlat,air.lat,air,input_core_dims=[["lat"],["lat"],["lat"]],output_core_dims=[["lat"]],exclude_dims={"lat"},
We will use a “wrapper” function debug_interp to examine what gets passed to numpy.interp.
Such wrapper functions are a great way to understand and debug apply_ufunc use cases.
defdebug_interp(xi,x,data):print(f"data: {data.shape} | x: {x.shape} | xi: {xi.shape}")returnnp.interp(xi,x,data)interped=xr.apply_ufunc(debug_interp,# first the functionnewlat,air.lat,air,input_core_dims=[["lat"],["lat"],["lat"]],output_core_dims=[["lat"]],exclude_dims={"lat"},# dimensions allowed to change size. Must be set!
That’s a hard-to-interpret error from NumPy but our print call helpfully printed the shapes of the input data:
data: (4, 3, 25) | x: (25,) | xi: (100,)
We see that apply_ufunc passes the full 3D array to interp1d_np which in turn passes that on to numpy.interp. But numpy.interp requires a 1D input, and thus the error.
Instead of passing the full 3D array 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 it easy to loop over the loop dimensions 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.
Warning
Also see the numpy documentation for numpy.vectorize. Most importantly
The vectorize function is provided primarily for convenience, not for performance.
The implementation is essentially a for loop.
air,input_core_dims=[["lat"],["lat"],["lat"]],output_core_dims=[["lat"]],exclude_dims={"lat"},# dimensions allowed to change size. Must be set!vectorize=True,interped
the printed input shapes are all 1D and correspond to one vector of size 25 along the lat dimension.
debug_interp was called 4x3 = 12 times which is the total number lat vectors since the size along time is 4, and the size along lon is 3.
The result interped is now an xarray object with coordinate values copied over from data.
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 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]_.
This is why apply_ufunc is so convenient; it takes care of a lot of code necessary to apply functions that consume and produce numpy arrays to xarray objects.
The vectorize keyword argument, when set to True, will use numpy.vectorize to apply the function by looping over the “loop dimensions” — dimensions that are not the core dimensions for the applied function.
