如何用xr.apply_ufunc处理仅支持标量的函数与多维DataArray

要解决这个问题，核心是让xr.apply_ufunc逐元素调用你的标量函数，而不是把整个数组传入。关键是正确设置vectorize和维度相关参数，以下是具体实现：

解决方案代码 #

import numpy as np
import xarray as xr

def f(a, b):
    assert(np.isscalar(a) and np.isscalar(b))
    result1 = a + b
    result2 = a * b
    return result1, result2

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

# 正确调用apply_ufunc实现逐元素处理
result1_da, result2_da = xr.apply_ufunc(
    f,
    da1,
    da2,
    vectorize=True,
    input_core_dims=[[], []],
    output_core_dims=[[], []],
    output_dtypes=[float, float]
)

参数解释 #

• vectorize=True：启用numpy的向量化机制，自动将标量函数f转换成可处理数组的函数，逐元素调用f，满足函数的标量输入要求。
• input_core_dims=[[], []]：指定每个输入参数不需要保留任何核心维度，即把每个位置的元素单独传入f，而不是传入整个二维数组。
• output_core_dims=[[], []]：因为f返回两个标量结果，该参数告诉xarray将每个输出映射回原DataArray的维度结构。
• output_dtypes=[float, float]：明确两个输出结果的数据类型，帮助xarray正确创建输出的DataArray对象。

结果验证 #

可以通过和手动循环的结果对比，确认实现正确性：

# 手动循环生成结果
manual_result1 = xr.zeros_like(da1)
manual_result2 = xr.zeros_like(da2)
for xi in da1['x']:
    for yi in da2['y']:
        r1, r2 = f(da1.sel(x=xi, y=yi).item(), da2.sel(x=xi, y=yi).item())
        manual_result1.loc[dict(x=xi, y=yi)] = r1
        manual_result2.loc[dict(x=xi, y=yi)] = r2

# 检查结果是否一致
print(np.allclose(result1_da, manual_result1))  # 输出True
print(np.allclose(result2_da, manual_result2))  # 输出True

内容的提问来源于stack exchange，提问作者hm8