直接应用和使用xarray.apply_ufunc应用的numpy渐变结果之间的差异 [英] Disparity between result of numpy gradient applied directly and applied using xarray.apply_ufunc

问题描述

我试图使用xarray的apply_ufunc包装numpy的gradient函数，以便沿一维进行渐变.但是，apply_ufunc返回的数组形状与使用np.gradient直接返回的数组不同:

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和actual也不只是彼此的换位版本.)我对为什么感到困惑-我对apply_ufunc的理解是核心尺寸移到了最后，所以axis=-1应该始终提供给ufunc吗?

(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?

推荐答案

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

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.

问题是在您的脚本中，apply_ufunc中未考虑坐标. 我认为您需要为所有输入传递input_core_dim；就您而言，是da和coord_vals的那些. 将[[dim]]更改为[[dim], []]可以正确计算，即

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])

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

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.

这篇关于直接应用和使用xarray.apply_ufunc应用的numpy渐变结果之间的差异的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！