在只能接受标量的函数上使用 xr.apply_unfunc，输入为多维 DataArrays

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

我还有两个二维的 xarray 数据数组：

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的计算开销比这个例子中展示的要大，那就可能会有问题，因为建议的解决方案会对da1和da2中的每一对值调用两次f。在这种情况下，就需要想出一个聪明的方法来存储中间结果，不过我想这要看具体的f是什么了。