Disparity between result of numpy gradient applied directly and applied using xarray.apply_ufunc

I'm trying to use xarray's apply_ufunc to wrap numpy's gradient function, in order to take gradients along one dimension. However, apply_ufunc is returning an array with a different shape to the one which using np.gradient directly returns:

import xarray as xr
import numpy as np

def wrapped_gradient(da, coord):
    """Finds the gradient along a given dimension of a dataarray."""

    dims_of_coord = da.coords[coord].dims
    if len(dims_of_coord) == 1:
        dim = dims_of_coord[0]
    else:
        raise ValueError('Coordinate ' + coord + ' has multiple dimensions: ' + str(dims_of_coord))

    coord_vals = da.coords[coord].values

    return xr.apply_ufunc(np.gradient, da, coord_vals, kwargs={'axis': -1},
                          input_core_dims=[[dim]], output_core_dims=[[dim]],
                          output_dtypes=[da.dtype])



# Test it out by comparing with applying np.gradient directly:
orig = xr.DataArray(np.random.randn(4, 3), coords={'x': [5, 7, 9, 11]}, dims=('x', 'y'))

expected = np.gradient(orig.values, np.array([5, 7, 9, 11]), axis=0)

actual = wrapped_gradient(orig, 'x').values

I want expected and actual to be the same, but instead they are different:

print(expected.shape)
> (4,3)
print(actual.shape)
> (3,4)

(expected and actual are also not just transposed versions of each other.) I'm confused as to why - my understanding of apply_ufunc was that the core dimensions are moved to the end, so that axis=-1 should always be supplied to the ufunc?

Solution

xr.apply_ufunc moves the input_core_dims to the last position. The dimension that gradient was computed along are moved to the last position, and therefore the resultant shape would be transposed compared with the result by np.gradient.

The problem is that in your script the coordinate is not considered in the apply_ufunc. I think you need to pass input_core_dims for all the inputs; in your case, those for da and coord_vals. Changing [[dim]] to [[dim], []] will compute correctly, i.e.,

return xr.apply_ufunc(np.gradient, da, coord_vals, kwargs={'axis': -1},
                      input_core_dims=[[dim], []], output_core_dims=[[dim]],
                      output_dtypes=[da.dtype])

BTW, I think xarray should raise an Error when input_core_dims does not match those expected for the inputs. I will raise an issue on Github.