遍历 xarray。 DataArray 第一维及其坐标

Question

假设我有以下 DataArray

arr = xarray.DataArray(np.arange(6).reshape(2,3),
                        dims=['A', 'B'],
                        coords=dict(A=['a0', 'a1'], 
                                    B=['b0', 'b1', 'b2']))

我想遍历第一个维度，做下面的事情（当然我想做比打印更复杂的事情）

for coor in arr.A.values:
    print(coor, arr.sel(A=coor).values)

并得到

a0 [0 1 2]
a1 [3 4 5]

我是 xarray 的新手，所以我想知道是否有更自然的方法来实现这一点，比如

for coor, sub_arr in arr.some_method():
    print(coor, sub_arr)

Answer 1

您可以简单地迭代 DataArray - 迭代器的每个元素本身都是一个 DataArray，第一个坐标只有一个值：

for a in arr:
    print(a.A.item(), a.values)

打印

a0 [0 1 2]
a1 [3 4 5]

注意使用.item()方法访问零维数组的标量值a.A。

要遍历第二个维度，您只需转置数据即可：

for b in arr.T: # or arr.transpose()
    print(b.B.item(), b.values)

打印

b0 [0 3]
b1 [1 4]
b2 [2 5]

对于多维数据，您可以使用省略号将要迭代的维度移动到第一位：

for x in arr.transpose("B", ...):
    # x has one less dimension than arr, and x.B is a scalar
    do_stuff_with(x)

关于重塑和重组数据的 documentation 有更多详细信息。

Answer 2

这是一个老问题，但我发现当您想迭代第一个维度以外的某个维度时，使用 groupby 比使用 transpose 更清晰，对我来说更直观：

for coor, sub_arr in arr.groupby('A'):
    print(coor)
    print(sub_arr)

a0
<xarray.DataArray (B: 3)>
array([0, 1, 2])
Coordinates:
  * B        (B) <U2 'b0' 'b1' 'b2'
    A        <U2 'a0'
a1
<xarray.DataArray (B: 3)>
array([3, 4, 5])
Coordinates:
  * B        (B) <U2 'b0' 'b1' 'b2'
    A        <U2 'a1'

此外，旧版本的 xarray 似乎无法正确处理省略号（请参阅 mgunyho 的回答），但 groupby 仍然可以正常工作。

遍历 xarray。 DataArray 第一维及其坐标

Iterating over xarray. DataArray first dimension and its coordinates

python

python-xarray