直接应用 numpy 梯度结果与使用 xarray.apply_ufunc 应用的结果之间的差异

我正在尝试使用 xarray 的apply_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 and actual也不只是彼此的转置版本。）我很困惑为什么 - 我的理解apply_ufunc是核心尺寸移到最后，这样axis=-1应该总是提供给ufunc？

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

问题是在您的脚本中没有考虑坐标apply_ufunc。 我认为你需要通过input_core_dims 代表所有输入；在你的情况下，那些da and coord_vals。 改变[[dim]] to [[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 上提出问题。