xarray处理ragged array时如何获取仅含单轨迹对应观测的groupby分组

问题原因 #

你遇到的分组结果冗余是因为groupby默认仅对分组字段所属的维度做过滤：

• 按轨迹维度traj下的id分组时，仅会过滤traj维度，不会自动筛选匹配该轨迹的观测，所以obs维度保留全部值
• 按观测维度obs下的ids分组时，仅会过滤obs维度，不会自动筛选匹配的轨迹，所以traj维度保留全部值

原生实现方案 #

方案1：广播轨迹级变量到观测维度（推荐） #

如果内存允许，你可以先把所有轨迹级变量按ids的映射关系广播到观测维度，后续分组直接得到符合要求的结果：

# 把所有轨迹级变量广播到obs维度，自动匹配对应轨迹的属性
ds = ds.assign(
    id = ds['id'].isel(traj=ds['ids']),
    count = ds['count'].isel(traj=ds['ids']),
    buoy_diameter = ds['buoy_diameter'].isel(traj=ds['ids'])
).drop_dims('traj') # 可选：移除冗余的traj维度

# 直接按id分组即可，每个分组仅包含对应轨迹的观测和属性
gr = ds.groupby('id')

该方案后续不需要额外处理分组结果，gr.map()可以直接对单条轨迹执行自定义函数。

方案2：分组后批量过滤（内存友好） #

如果轨迹级变量较多、内存压力大，可以直接在groupby后通过map批量执行过滤逻辑，不需要手动遍历分组：

def filter_single_trajectory(group):
    # 取当前分组对应的轨迹索引
    traj_idx = group['ids'].values[0]
    # 过滤traj维度仅保留对应轨迹
    return group.sel(traj=traj_idx)

# 直接得到处理后的分组结果
processed_groups = ds.groupby('ids').map(filter_single_trajectory)

两种方案都是xarray原生支持的操作，不需要引入额外依赖。

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