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

Question

我正在尝试使用 xarray 的 apply_ufunc 来包装 numpy 的 gradient 函数，以便沿一维获取梯度。但是，apply_ufunc 返回的数组与直接使用 np.gradient 的数组形状不同 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

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

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

（expected 和 actual 也不仅仅是彼此的转置版本。）我很困惑为什么 - 我对 apply_ufunc 的理解是核心维度是移动到最后，所以 axis=-1 应该总是提供给 ufunc?

Answer 1

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

问题是在您的脚本中 apply_ufunc 中没有考虑坐标。我认为您需要为所有输入传递 input_core_dims；在您的情况下，da 和 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])

顺便说一句，我认为 xarray 应该在 input_core_dims 与输入的预期不匹配时引发错误。我将在 Github.

上提出问题

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

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

python

numpy

numpy-ufunc

python-xarray