xarray 相当于 np.reshape

Question

我有一个 3d 数组 (10x10x3)，由于某种原因，它被保存为 2d xr.DataArray (100x3)。它看起来有点像这样：

data = xr.DataArray(np.random.randn(100, 3),
                    dims=('ct', 'x'),
                    coords={'ct': range(100)})

c = [x%10 for x in range(100)]
t = [1234+x//10 for x in range(100)]

c 和 t 是在 ct 中捆绑在一起的坐标。

过去我已经解决了分离二维的问题，如下：

t_x_c,x = data.shape
nc = 10
data = np.reshape(data.values,(t_x_c//nc,nc, x))

但这需要数据结构中的许多假设，这些假设在不久的将来可能不成立（例如，c 和 t 可能不像我的示例中那样规则）。

我已设法将 c 和 t 作为附加坐标分配给数组：

data2 = data.assign_coords(
    coords={"c": ("ct", c),
            "t": ("ct", t),
},)

但我想将它们提升到数组的维度。我该怎么做呢？

Answer 1

.set_index()您想要使用和方法的组合.unstack()。

让我们把它分开吧。

首先，我在“c”和“t”已经是坐标的阶段创建虚拟数组：

c, t = [arr.flatten() for arr in np.meshgrid(range(10), range(1234, 1234+10))]

da = xr.DataArray( 
    np.random.randn(100, 3), 
    dims=('ct', 'x'), 
    coords={ 
        'c': ('ct', c), 
        't': ('ct', t) 
    }
)

然后，使用set_index()创建MultiIndex组合“c”和“t”坐标：

>>> da.set_index(ct=("c", "t"))                                                                  
<xarray.DataArray (ct: 100, x: 3)>
[...]
Coordinates:
  * ct       (ct) MultiIndex
  - c        (ct) int64 0 1 2 3 4 5 6 7 8 9 0 1 2 ... 
  - t        (ct) int64 1234 1234 1234 1234 1234 ...
Dimensions without coordinates: x

然后，使用unstack()“ct”MultiIndex 的“c”和“t”级别作为维度：

>>> da.set_index(ct=("c", "t")).unstack("ct")
<xarray.DataArray (x: 3, c: 10, t: 10)>
Coordinates:
  * c        (c) int64 0 1 2 3 4 5 6 7 8 9
  * t        (t) int64 1234 1235 1236 1237 1238 1239 1240 1241 1242 1243
Dimensions without coordinates: x

>>> da.set_index(ct=("c", "t")).unstack("ct").dims
('x', 'c', 't')

然而，正如您所看到的，.unstack()将未堆叠的维度放在最后。所以你最终可能想要转置：

>>> da.set_index(ct=("c", "t")).unstack("ct").transpose("c", "t", "x").dims                      
('c', 't', 'x')