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

ぃ、小莉子 提交于 2020-03-16 07:09:09

问题


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?


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



来源:https://stackoverflow.com/questions/51680659/disparity-between-result-of-numpy-gradient-applied-directly-and-applied-using-xa

易学教程内所有资源均来自网络或用户发布的内容,如有违反法律规定的内容欢迎反馈
该文章没有解决你所遇到的问题?点击提问,说说你的问题,让更多的人一起探讨吧!