如何在xarray中为指定维度和坐标添加新值？

数据集定义 #

用户的Xarray DataSet由以下代码生成：

import pandas as pd
import numpy as np
import xarray as xr

df = pd.DataFrame({
    'var': np.random.rand(16),
    'country': 4*['UK'] + 4*['US'] + 4*['FR'] + 4*['DE'],
    'time_delta': 4*list(pd.timedelta_range(
        start='30D',
        periods=4,
        freq='30D'
    )),
})

ds = df.set_index(['country','time_delta']).to_xarray()

需求 #

为维度time_delta的新坐标0D设置变量var的取值为0，同时保留其他现有维度（此处为country）。

Pandas实现方式 #

用户已通过Pandas完成该需求，步骤如下：

# 1. 长表转宽表
df_wide = df.pivot(
    index='country',
    columns='time_delta'
).droplevel(0,axis=1)

# 2. 设置新值
df_wide['0D'] = 0

# 3. 宽表转长表
df_new = df_wide.melt(ignore_index=False).rename(
    columns={'value': 'var'}
).reset_index()

ds_new = df_new.set_index(['country','time_delta']).to_xarray()

失败的Xarray尝试 #

用户尝试直接使用loc赋值报错：

ds['var'].loc[{'time_delta': '0D'}] = 0

报错信息：KeyError: "not all values found in index 'time_delta'"

Xarray通用解法 #

要在Xarray中直接实现需求，需先扩展维度坐标再赋值，以下是几种通用方法：

方法1：通过xr.concat合并新数据集 #

# 创建包含新坐标的数据集，变量值设为0，匹配原数据集的`country`维度
new_coord = pd.Timedelta('0D')
new_ds = xr.Dataset(
    {'var': (['country', 'time_delta'], [[0] for _ in ds.country])},
    coords={
        'country': ds.country,
        'time_delta': [new_coord]
    }
)

# 合并原数据集与新数据集，自动对齐维度后排序
ds_new = xr.concat([ds, new_ds], dim='time_delta').sortby('time_delta')

方法2：扩展维度后合并变量 #

# 创建对应新坐标的变量数组，维度与`country`一致
new_var = xr.DataArray(
    np.zeros_like(ds.country),
    coords={'country': ds.country},
    dims=['country']
)
# 扩展`time_delta`维度并指定新坐标
new_var = new_var.expand_dims(time_delta=[pd.Timedelta('0D')])

# 合并到原数据集
ds_new = xr.merge([ds, new_var]).sortby('time_delta')

方法3：基于现有数据复制修改后合并 #

# 复制原数据的一个切片，替换值为0并修改坐标
new_entry = ds['var'].isel(time_delta=0).copy()
new_entry.data = np.zeros_like(new_entry.data)
new_entry = new_entry.assign_coords(time_delta=pd.Timedelta('0D'))

# 合并到原数据集
ds_new = xr.concat([ds, new_entry.to_dataset(name='var')], dim='time_delta').sortby('time_delta')

以上方法均无需在Pandas与Xarray间来回转换，即可保留原country维度，同时为新的time_delta坐标0D设置var值为0。

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