在只能采用标量且输入为多维 DataArray 的函数上使用 xr.apply_unfunc

Question

假设我有一些函数

f

只能接受标量值：

def f(a, b):
    # Ensure inputs are scalars
    assert(np.isscalar(a) and np.isscalar(b))
    
    result1 = a + b
    result2 = a * b
    return result1, result2

我还有两个 2D xarray DataArray：

da1 = xr.DataArray(np.random.randn(3, 3), dims=('x', 'y'), name='a')
da2 = xr.DataArray(np.random.randn(3, 3), dims=('x', 'y'), name='b')

我想将 f() 一起应用于 da1 和 da2 的每个索引。本质上，我想这样做：

result1_da = xr.zeros_like(da1)
result2_da = xr.zeros_like(da2)
for xi in da1['x']:
    for yi in da2['y']:
        result1, result2 = f(da1.sel(x = xi, y = yi).item(), 
                             da2.sel(x = xi, y = yi).item())
        result1_da.loc[dict(x=xi, y=yi)] = result1
        result2_da.loc[dict(x=xi, y=yi)] = result2

但是没有循环。我想我应该能够使用

xr.apply_unfunc

来做到这一点。但我无法完全让它发挥作用。如果我这样做：

xr.apply_ufunc(f, da1, da2)

我收到断言错误（未传入标量）

assert(np.isscalar(a) and np.isscalar(b) )
AssertionError

我也尝试过使用

input_core_dims

和其他

xr.apply_unfunc

参数，但我无法让任何东西发挥作用。

Answer 1

据我了解 xarray.apply_ufunc

的 描述，此方法默认假设给定函数

f

可以作用于 NumPy 数组。 如果它不能 - 就像你的函数

f

- 我们必须设置

vectorize=True

。 但是，这不能直接起作用，因为您的函数

f

返回两个数字的序列，而不是单个数字，因此返回值不能放置在新数组中。 因此，我建议如下：

result1 = xr.apply_ufunc(lambda a, b: f(a, b)[0], da1, da2, vectorize=True)
result2 = xr.apply_ufunc(lambda a, b: f(a, b)[1], da1, da2, vectorize=True)

这会为您的 for 循环提供相同的结果。

如果函数

f

的计算成本比本示例中显示的要高，则建议的解决方案可能会为

f

和

da1

中的每对值对

da2

进行两次调用。 在这种情况下，必须想出一种聪明的方法来存储中间结果，但我想这取决于实际的

f

。