Xarray如何对同一维度下不同坐标的两行数据执行求和运算

问题根源 #

你遇到的求和不符合预期（通常返回全NaN）的核心原因是：如果你选择坐标时保留了x维度（比如传入列表sel(x=['a'])，或者手动关闭了squeeze），那么M.sel(x='a')和N.sel(x='b')的x维度坐标分别只有a和b，没有重叠值，xarray按坐标自动对齐后相加就会得到全空的结果。而直接取.values相加时绕开了xarray的坐标对齐逻辑，所以能拿到数值，但多维度场景下会因为维度顺序不匹配报错。

解决方案 #

xarray原生支持这类需求，不需要手动处理numpy维度适配，下面是两种常用场景的写法：

场景1：不需要在结果中保留x维度（最常用） #

直接选择单个坐标值，xarray默认会挤压掉长度为1的x维度，此时两个数组的运算只会对其余维度按名称自动对齐，完全不需要手动调整维度顺序：

import xarray as xr
import numpy as np

# 简单场景验证
M = xr.DataArray([1, 2], dims="x",coords={'x':['a','b']})
N = xr.DataArray([3, 4], dims="x",coords={'x':['a','b']})
res_simple = M.sel(x='a') + N.sel(x='b')
print(res_simple) # 输出为5，和.values相加结果一致

# 多维度场景验证
m = np.arange(3*2*4).reshape(3,2,4)
n = np.arange(4*2*3).reshape(4,2,3)
M = xr.DataArray(m, dims=['z1',"x","z2"],coords={'x':['a','b']})
N = xr.DataArray(n, dims=["z2",'x','z1'],coords={'x':['a','b']})
res_multi = M.sel(x='a') + N.sel(x='b')
print(res_multi.shape) # 自动对齐后维度为(3, 4)，无广播错误
print(res_multi.values) # 直接返回对齐后的numpy数组，不需要手动调整维度

场景2：需要在结果中保留x维度 #

如果后续还要和包含x维度的数组做运算，选择坐标时手动关闭squeeze保留x维度，再将其中一个数组的x维度重命名避免对齐冲突即可：

M_a = M.sel(x='a', squeeze=False)
N_b = N.sel(x='b', squeeze=False).rename_dims({'x':'x_target'})
res_with_x = M_a + N_b
print(res_with_x.dims) # 输出为('z1', 'x', 'z2', 'x_target')，保留所有维度

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