结合互补的DataArrays

时间:2018-04-30 06:26:20

标签: pandas python-xarray xarray

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

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,但到目前为止我所尝试的似乎都没有。 结果应为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)

1 个答案:

答案 0 :(得分: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不同的算法,该算法为每个参数分配完整的输出数组。因此对于大型阵列来说会慢得多。