用于多维坐标的 xarray 数据集 swap_dims答案

【问题标题】：xarray Dataset swap_dims for multi-dimensional coord用于多维坐标的 xarray 数据集 swap_dims
【发布时间】：2018-12-31 07:35:01
【问题描述】：

我想用一个多维坐标交换两个坐标维度，这样我就可以按 time.month 执行 groupby 并减去另一个数据集。

import xarray as xr

ds = xr.Dataset()

# DataArray indexed by 'init_time' and an offset, 'tau'
ds['tst'] = xr.DataArray(
    [[0, 1, 2], [3, 4, 5]],
    dims=('init_time', 'tau'),
    coords={
        'init_time': pd.date_range('2017-01-01', periods=2),
        'tau': pd.to_timedelta([1, 2, 3], unit='days')})

# multidimensional coordinate 'time'
ds.coords['time'] = ds['init_time'] + ds['tau']

ds.swap_dims({('init_time', 'tau'): 'time'})

ds

有点像这样的结果：

clim = pd.Series([2], index=[1]).rename_axis('month')
df = ds.to_dataframe().reset_index()
df['month'] = df['time'].dt.month
df = (
    pd.DataFrame(
        df.set_index(['init_time', 'tau', 'time', 'month'])['tst']
        - clim))

df

【问题讨论】：

标签： python-xarray

【解决方案1】：

这样做的问题是交换尺寸会导致索引中的重复值。理想情况下，您可以在多维坐标上进行分组。您目前可以执行此操作，但功能不完整（例如，您不能执行ds.groupby('time.month').mean(dim='time')）。这看起来可能正在进行中（请参阅#324、#2525）。

现在，我认为您有两个选择。你可以在 pandas 中做到这一点：

df = ds.to_dataframe().reset_index()
monthly_mean = (
    df
    .groupby([df.other_dims, df.time.dt.month])
    .mean()[['tst']]
    .to_xarray())

clim = xr.DataArray([2], dims=['month'], coords=[[1]])

anom = monthly_mean.rename({'time': 'month'}) - clim

或者，您可以通过堆叠 init_time 和 tau 将其保存在 xarray 中：

In [35]: stacked = ds.stack(obs=('init_time', 'tau'))

In [36]: stacked.coords['obs_num'] = ('obs', ), np.arange(len(stacked.obs))

In [37]: stacked.coords['time'] = ('obs', ), stacked.init_time + stacked.tau

In [38]: swapped = stacked.swap_dims({'obs': 'obs_num'})

In [39]: swapped
Out[39]:
<xarray.Dataset>
Dimensions:  (obs_num: 150)
Coordinates:
    time     (obs_num) datetime64[ns] 2017-01-01 2017-01-02 ... 2017-02-03
    obs      (obs_num) object (Timestamp('2017-01-01 00:00:00', freq='D'), Timedelta('0 days 00:00:00')) ... (Timestamp('2017-01-30 00:00:00', freq='D'), Timedelta('4 days 00:00:00'))
  * obs_num  (obs_num) int64 0 1 2 3 4 5 6 7 ... 142 143 144 145 146 147 148 149
Data variables:
    tst      (obs_num) int64 0 1 2 3 4 5 6 7 ... 142 143 144 145 146 147 148 149

In [47]: swapped.groupby(swapped.time.dt.month).mean(dim='obs_num')
Out[47]:
<xarray.Dataset>
Dimensions:  (month: 2)
Coordinates:
  * month    (month) int64 1 2
Data variables:
    tst      (month) float64 71.56 145.0

【讨论】：