迭代 xarray。DataArray 第一维及其坐标

Question

假设我有以下内容 DataArray

arr = xarrray.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

这是一个老问题，但我发现groupby当你想迭代第一个维度以外的某个维度时，使用比使用转置更干净，对我来说更直观：

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

Answer 2

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

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

印刷

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

注意这里a.A.values是一个零维的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", ...):
    ...

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