如何使用新维度将xarray.Dataset的多个变量合并为一个
问题描述
我有一个包含多个变量(如t1、t2、t3等)的xarray.Dataset,希望通过新增维度将这些变量合并为一个名为t的变量。以下是我实现需求的最小工作示例(MWE),但我认为xarray已有现成函数或更简洁的实现方式:
import xarray as xr import numpy as np ds = xr.Dataset({"t1": (("y", "x"), np.random.rand(6).reshape(2, 3)), "t2": (("y", "x"), np.random.rand(6).reshape(2, 3)), "t3": (("y", "x"), np.random.rand(6).reshape(2, 3)), }, coords={"y": [0, 1], "x": [10, 20, 30]},) t_list = [] for i, v in enumerate(ds.variables.keys()): if v.startswith("t"): t_list.append(ds[v].expand_dims("α").assign_coords(α=[i])) ds = ds.drop(v) ds["t"] = xr.concat(t_list, dim="α")
简洁实现方案
xarray确实有更简洁的内置方法可以完成这个需求,无需手动循环,以下是两种常用写法:
方案一:筛选变量后直接concat
先筛选所有以t开头的变量,用xr.concat直接合并并新增维度,最后将合并后的变量加入原Dataset(可选择删除原变量):
import xarray as xr import numpy as np ds = xr.Dataset({"t1": (("y", "x"), np.random.rand(6).reshape(2, 3)), "t2": (("y", "x"), np.random.rand(6).reshape(2, 3)), "t3": (("y", "x"), np.random.rand(6).reshape(2, 3)), }, coords={"y": [0, 1], "x": [10, 20, 30]},) # 筛选目标变量 t_vars = [var for var in ds if var.startswith("t")] # 合并为新变量`t`,新增维度`α`,坐标可设为原变量名或索引 ds["t"] = xr.concat([ds[var] for var in t_vars], dim="α").assign_coords(α=t_vars) # 可选:移除原`t1`/`t2`/`t3`变量 ds = ds.drop(t_vars)
方案二:用xr.merge+expand_dims
先给每个目标变量新增维度,再通过合并转成单一变量:
# 为每个`t`开头的变量新增`α`维度,坐标设为变量名 expanded_t_vars = [ds[var].expand_dims(α=[var]) for var in ds if var.startswith("t")] # 合并后转成单一变量`t` ds["t"] = xr.merge(expanded_t_vars).to_array(dim="α").squeeze() # 可选:删除原变量 ds = ds.drop([var for var in ds if var.startswith("t")])
两种方案都能实现和原代码一致的效果,且更符合xarray的惯用写法,代码更简洁易读。
内容的提问来源于stack exchange,提问作者TomCho
相关产品推荐
相关产品推荐

