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

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 ...

但是这个很贵，有没有更好的办法？基本上在每个 'time' 切片上，我想要一个对角线数组，其中填充了原始数据中的值。似乎我应该能够以某种方式构建原始数据的视图来完成此操作，但我不知道该怎么做。

Answer 1

您不需要 for-loop:

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

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

Reindex data variable by different dimensions

python

numpy

xarray