通过xarray.Dataset的groupby往返操作时的坐标异常及修复方法

问题原因 #

你遇到的情况是因为分组键da基于原数据集所有坐标创建，维度为(z, y)。groupby过程中xarray会自动将所有坐标与分组键维度对齐——原坐标x仅含z维度，因此被广播至(z, y)维度匹配分组键，后续unstack操作保留了这个广播后的形态，导致与原数据集不一致。

解决方案 #

方法1：用标量作为分组键（推荐）

若需求是将整个数据集作为单个分组，无需创建带坐标的DataArray，直接用标量常量作为分组键即可，不会触发坐标广播：

ds_before = xr.Dataset(coords={'x': ('z', [0]), 'y': ('y', [1, 2])})
key, value = next(iter(ds_before.groupby(True)))
ds_after = value.unstack()

此时ds_after的坐标x会保持原有的(z)维度，与ds_before完全一致。

方法2：手动还原被广播的坐标

如果必须使用原分组键da，可在unstack后手动将坐标恢复至原始状态：

# 执行原有分组逻辑
da = xr.DataArray(True, coords=ds_before.coords)
key, value = next(iter(ds_before.groupby(da)))
ds_after = value.unstack()

# 还原单个坐标x
ds_after['x'] = ds_before['x']

# 或批量还原所有坐标（更通用）
for coord_name in ds_before.coords:
    if coord_name in ds_after.coords:
        ds_after[coord_name] = ds_before[coord_name]

额外说明 #

xarray的groupby默认会让所有变量、坐标与分组键维度对齐，以此保证分组后数据的一致性。当分组键包含额外维度时，非维度坐标会被广播填充，这是预期行为，可通过上述方法规避或还原。

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