使用Xarray创建新变量后如何保留原始维度？

核心原因 #

1. 用ds.variables['uwnd']取出的变量虽为DataArray，但第三方库（如windpowerlib）可能会自动转换为NumPy数组处理，丢失xarray的维度/坐标元数据。
2. 直接调用NumPy函数处理DataArray时，部分场景会返回纯NumPy数组，而非带维度的DataArray。

修复步骤 #

1. 保留DataArray包装进行计算

避免单独取出变量后脱离数据集操作，直接基于原数据集的变量进行计算，xarray会自动保留维度信息：

# 直接基于数据集计算风速，自动保留维度
ws = (ds['uwnd'] ** 2 + ds['vwnd'] ** 2) ** 0.5

2. 处理第三方库输出（重新包装为DataArray）

windpowerlib返回的是NumPy数组，需要手动包装成带维度和坐标的DataArray：

# 先获取NumPy数组结果
ws100_np = windpowerlib.wind_speed.hellman(ws.values, 10, 100)
# 用原风速变量的维度和坐标创建新DataArray
ws100 = xr.DataArray(
    ws100_np,
    dims=ws.dims,
    coords=ws.coords,
    name='ws100'
)

3. 用xarray工具包装NumPy风向计算

使用xr.apply_ufunc将NumPy的向量操作包装为保留维度的xarray操作：

wdir = xr.apply_ufunc(
    lambda u, v: np.mod(180 + np.rad2deg(np.arctan2(u, v)), 360),
    ds['uwnd'], ds['vwnd'],
    output_dtypes=[float]
)
wdir = wdir.rename('wdir')

4. 将新变量插入原数据集

用assign方法一次性添加所有新变量，自动继承原数据集的维度和结构：

ds = ds.assign(
    ws=ws,
    ws100=ws100,
    wdir=wdir
)

简化版完整代码 #

# 计算风速并添加到数据集
ds['ws'] = (ds['uwnd']**2 + ds['vwnd']**2)**0.5

# 计算100米风速并添加
ws100_np = windpowerlib.wind_speed.hellman(ds['ws'].values, 10, 100)
ds['ws100'] = xr.DataArray(ws100_np, dims=ds['ws'].dims, coords=ds['ws'].coords)

# 计算风向并添加
ds['wdir'] = xr.apply_ufunc(
    lambda u, v: np.mod(180 + np.rad2deg(np.arctan2(u, v)), 360),
    ds['uwnd'], ds['vwnd'],
    output_dtypes=[float]
)

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