2

我正在尝试使用 xarrayapply_ufunc来包装 numpy 的gradient函数，以便沿一维获取渐变。但是，apply_ufunc返回的数组与np.gradient直接使用返回的数组形状不同：

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

我希望预期和实际相同，但它们却不同：

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

（expected并且actual也不仅仅是彼此的转置版本。）我很困惑为什么 - 我的理解apply_ufunc是核心尺寸被移到最后，所以axis=-1应该总是提供给 ufunc？

4

1 回答 1

1

xr.apply_ufunc将 移动input_core_dims到最后一个位置。计算梯度的维度被移动到最后一个位置，因此与结果相比，结果形状将被转置np.gradient。

问题是在您的脚本中，坐标没有在apply_ufunc. 我认为您需要input_core_dim为所有输入传递 s ；在你的情况下，那些 forda和coord_vals。更改[[dim]]为[[dim], []]将正确计算，即

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

顺便说一句，我认为当 input_core_dims 与输入的预期不匹配时，xarray 应该引发错误。我将在 Github 上提出一个问题。

于 2018-08-05T02:02:32.407 回答