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 differen

Solution 1:

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.