按不同维度重新索引数据变量

Question

我有一个由time和id标注的数据集，但它也有lat和lon坐标。

数据变量由time和id标注，我想要做的是用time，lat和lon来标注它。例如：

import numpy
import xarray

ds = xarray.Dataset()
ds['data'] = (('time', 'id'), numpy.arange(0, 50).reshape((5, 10)))

ds.coords['time'] = (('time',), numpy.arange(0, 5))
ds.coords['id'] = (('id',), numpy.arange(0, 10))

ds.coords['lat'] = (('lat',), numpy.arange(10, 20))
ds.coords['lon'] = (('lon',), numpy.arange(20, 30))

print ds

结果：

<xarray.Dataset>
Dimensions:  (id: 10, lat: 10, lon: 10, time: 5)
Coordinates:
  * time     (time) int64 0 1 2 3 4
  * id       (id) int64 0 1 2 3 4 5 6 7 8 9
  * lat      (lat) int64 10 11 12 13 14 15 16 17 18 19
  * lon      (lon) int64 20 21 22 23 24 25 26 27 28 29
Data variables:
    data     (time, id) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 ...

我能弄清楚如何实现这一目标的唯一方法是迭代索引，构建一个具有正确形状和尺寸的新数据数组：

reshaped_array = numpy.ma.masked_all((5, 10, 10))
for t_idx in range(0, 5):
    for r_idx in range(0, 10):
        reshaped_array[t_idx, r_idx, r_idx] = ds['data'][t_idx, r_idx]

ds['data2'] = (('time', 'lat', 'lon'), reshaped_array)

print ds

结果：

<xarray.Dataset>
Dimensions:  (id: 10, lat: 10, lon: 10, time: 5)
Coordinates:
  * time     (time) int64 0 1 2 3 4
  * id       (id) int64 0 1 2 3 4 5 6 7 8 9
  * lat      (lat) int64 10 11 12 13 14 15 16 17 18 19
  * lon      (lon) int64 20 21 22 23 24 25 26 27 28 29
Data variables:
    data     (time, id) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 ...
    data2    (time, lat, lon) float64 0.0 nan nan nan nan nan nan nan nan ...

但这是非常昂贵的，有更好的方法吗？基本上在每个'时间'切片我想要一个填充了原始数据值的对角线数组。看起来我应该能够以某种方式构建一个原始数据的视图来实现这一点，但我不知道如何做到这一点。

Answer 1

你不需要for循环：

res = np.full((5, 10, 10), np.nan)
idx = np.arange(10)
res[:, idx, idx] = ds['data']
ds['data2'] = (('time', 'lat', 'lon'), res)