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

【问题标题】：Iterating over xarray. DataArray first dimension and its coordinates遍历 xarray。 DataArray 第一维及其坐标
【发布时间】：2022-05-02 16:05:09
【问题描述】：

假设我有以下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)

【问题讨论】：

标签： python python-xarray

【解决方案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 有更多详细信息。

【讨论】：

这很酷，但仍然有点丑陋和不直观。很遗憾，xarray 没有更自然的方式来做到这一点。
执行此操作时出现错误：“ValueError：在数组维度（'time'、'ens_member'、'station'、'lead_time'）中找不到省略号”。我的 xarray 版本可能有点旧？
对我来说，遍历数组沿第一个维度进行切片是有道理的，我不知道它还能做什么。似乎在version 0.14.1 中添加了对使用省略号进行转置的支持。
哦，是的，第一个维度迭代还可以，虽然我仍然觉得有点难看，因为 xarray 通常比这更抽象，即大多数时候你不应该真正需要关心哪个维度在底层数组中的顺序是什么；当然，它在代码中的可读性较差，因为这种排序通常不明显。转置版本具有抽象优势，但有点丑陋，因为您需要处理所有维度，而不仅仅是您关心的维度。使用 groupby 感觉更好，尤其是在省略号位在转置中不起作用的旧版本中。

【解决方案2】：

这是一个老问题，但我发现当您想要迭代除第一个维度之外的某个维度时，使用 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 仍然可以正常工作。

【讨论】：