遍历 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.values, a.values)

印刷

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

注意这里的aAvalues是一个零维的 numpy 数组，使用.item()方法访问底层标量。

要迭代第二个维度，您只需转置数据：

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

印刷

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

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

for x in arr.transpose("B", ...):
    ...

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

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 仍然可以正常工作。