如何使用自定义/外部函数沿指定维度归约xarray.Dataset

报错原因 #

• 首先xarray的reduce方法执行归约时，会默认向传入的自定义函数传递axis参数，但sklearn的mean_squared_error没有定义该参数，直接调用就会触发参数不匹配的错误。
• 其次Dataset的reduce默认会对每个变量单独执行归约逻辑，而你需要同时用到obs和sim两个变量计算MSE，单变量处理无法得到预期结果。

解决方法 #

方案1：用xarray原生运算实现（推荐，性能最优） #

不需要引入外部依赖，直接基于xarray的内置运算完成归约，代码最简洁：

# 先算残差平方，再沿time维度取均值
mse = ((original["obs"] - original["sim"]) ** 2).mean(dim="time")
# 转成指定格式的Dataset
result = mse.to_dataset(name="mean_squared_error")

方案2：用apply_ufunc适配外部函数 #

如果你必须使用sklearn的mean_squared_error或者其他自定义外部函数，可以用xarray的apply_ufunc方法适配，它可以灵活处理多输入、多维度的归约场景：

from sklearn.metrics import mean_squared_error
import xarray as xr

result = xr.apply_ufunc(
    mean_squared_error,
    original["obs"],  # 第一个输入参数：观测值
    original["sim"],  # 第二个输入参数：模拟值
    input_core_dims=[["time"], ["time"]],  # 两个输入的核心计算维度都是time，归约时保留该维度传入函数
    output_core_dims=[[]],  # 输出结果不需要保留time维度
    vectorize=True,  # 自动向量化遍历其他维度（这里就是station_id维度）
    kwargs={"multioutput": "raw_values"}  # 传给mean_squared_error的参数，保证每个站点单独返回结果
).to_dataset(name="mean_squared_error")

两种方案运行后都可以得到你预期的输出格式。

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