结合互补的 DataArrays

Question

我有一个包含三个维度的 DataArray 列表。对于列表中的每个项目，其中两个维度是单个值，但所有项目的组合将产生完整的组合值。

import itertools
import numpy as np
import xarray as xr

ds = []
for vals_dim1, vals_dim2 in itertools.product(list(range(2)), list(range(3))):
    d = xr.DataArray(np.random.rand(1, 1, 4),
                     coords={'dim1': [vals_dim1], 'dim2': [vals_dim2], 'dim3': range(4)},
                     dims=['dim1', 'dim2', 'dim3'])
    ds.append(d)

然后我想结合这些免费的 DataArray，但到目前为止我尝试过的 none 似乎有效。结果应该是 DataArray，形状 |2x3x4|，尺寸 dim1: |2|, dim2: |3|, dim3: |4|。

以下不起作用：

# does not automatically infer dimensions and fails with
# "ValueError: conflicting sizes for dimension 'concat_dim': length 2 on 'concat_dim' and length 6 on <this-array>"
ds = xr.concat(ds, dim=['dim1', 'dim2'])

# will still try to insert a new `concat_dim` and fails with
# "ValueError: conflicting MultiIndex level name(s): 'dim1' (concat_dim), (dim1) 'dim2' (concat_dim), (dim2)"
import pandas as pd
dims = [[0] * 3 + [1] * 3, list(range(3)) * 2]
dims = pd.MultiIndex.from_arrays(dims, names=['dim1', 'dim2'])
ds = xr.concat(ds, dim=dims)

# fails with
# AttributeError: 'DataArray' object has no attribute 'data_vars'
ds = xr.auto_combine(ds)

Answer 1

不幸的是（正如您在此处发现的那样），您目前无法在 xarray 中同时沿多个维度连接。

有几种方法可以解决这个问题。性能最高的是stack()沿新维度的所有对象，然后在连接后unstack()：

>>> xr.concat([d.stack(z=['dim1', 'dim2']) for d in ds], 'z').unstack('z')
<xarray.DataArray (dim3: 4, dim1: 2, dim2: 3)>
array([[[0.300328, 0.544551, 0.751339],
        [0.612358, 0.937376, 0.67688 ]],

       [[0.065146, 0.85845 , 0.962857],
        [0.102126, 0.395406, 0.245373]],

       [[0.309324, 0.362568, 0.676552],
        [0.709206, 0.719578, 0.960803]],

       [[0.613187, 0.205054, 0.021796],
        [0.434595, 0.779576, 0.937855]]])
Coordinates:
  * dim3     (dim3) int64 0 1 2 3
  * dim1     (dim1) int64 0 1
  * dim2     (dim2) int64 0 1 2

（这里的z是一个占位符，实际上只是临时新维度的任意名称。）

另一种选择是使用 merge()。合并与 DataArray 对象一起使用有点笨拙（我们应该解决这个问题），但这会达到相同的结果：

>>> xr.merge([x.rename('z') for x in ds])['z'].rename(None)
<xarray.DataArray (dim1: 2, dim2: 3, dim3: 4)>
array([[[0.300328, 0.065146, 0.309324, 0.613187],
        [0.544551, 0.85845 , 0.362568, 0.205054],
        [0.751339, 0.962857, 0.676552, 0.021796]],

       [[0.612358, 0.102126, 0.709206, 0.434595],
        [0.937376, 0.395406, 0.719578, 0.779576],
        [0.67688 , 0.245373, 0.960803, 0.937855]]])
Coordinates:
  * dim1     (dim1) int64 0 1
  * dim2     (dim2) int64 0 1 2
  * dim3     (dim3) int64 0 1 2 3

(z这里也是占位符名称。)

请注意，merge 使用与 concat 不同的算法，后者为每个参数分配完整的输出数组。所以对于大数组会慢很多。

结合互补的 DataArrays

combine complementary DataArrays

pandas

python-xarray

xarray