xarray教程之GroupBy: split-apply-combine

xarray 支持与 pandas 相同API的 “group by” 操作，以实现split-apply-combine策略：

• 将数据分成多个独立的组。
• 将一些功能应用于每个组。
• 将您的组重新组合到一个数据对象中。

分组操作对 Dataset 和 DataArray 对象均起作用。 尽管最近已实现了对多维变量分组的支持，但大多数示例着重于按单个一维变量分组。 请注意，对于一维数据，通常依靠 pandas对同一管道的实现会更快。

分割(Split)

让我们创建一个简单的 dataset 事例：

In [1]: ds = xr.Dataset(
   ...:     {"foo": (("x", "y"), np.random.rand(4, 3))},
   ...:     coords={"x": [10, 20, 30, 40], "letters": ("x", list("abba"))},
   ...: )
   ...: 

In [2]: arr = ds["foo"]

In [3]: ds
Out[3]: 
<xarray.Dataset>
Dimensions:  (x: 4, y: 3)
Coordinates:
  * x        (x) int64 10 20 30 40
    letters  (x) <U1 'a' 'b' 'b' 'a'
Dimensions without coordinates: y
Data variables:
    foo      (x, y) float64 0.127 0.9667 0.2605 0.8972 ... 0.543 0.373 0.448

如果我们根据变量名称或数据集中的坐标进行分组（也可以直接使用DataArray），则会返回GroupBy对象：

In [4]: ds.groupby('letters')
Out[4]: 
DatasetGroupBy, grouped over 'letters' 
2 groups with labels 'a', 'b'.

该对象的工作方式与pandas 的GroupBy对象非常相似，可以使用groups属性查看组索引：

In [5]: ds.groupby('letters').groups
Out[5]: {'a': [0, 3], 'b': [1, 2]}

也可以遍历(label, group)对中的组：

In [6]: list(ds.groupby('letters'))
Out[6]: 
[('a', <xarray.Dataset>
  Dimensions:  (x: 2, y: 3)
  Coordinates:
    * x        (x) int64 10 40
      letters  (x) <U1 'a' 'a'
  Dimensions without coordinates: y
  Data variables:
      foo      (x, y) float64 0.127 0.9667 0.2605 0.543 0.373 0.448),
 ('b', <xarray.Dataset>
  Dimensions:  (x: 2, y: 3)
  Coordinates:
    * x        (x) int64 20 30
      letters  (x) <U1 'b' 'b'
  Dimensions without coordinates: y
  Data variables:
      foo      (x, y) float64 0.8972 0.3767 0.3362 0.4514 0.8403 0.1231)]

就像在 pandas 中一样，创建GroupBy对象很低消耗：在访问特定值之前，它实际上不会拆分数据。

分箱(Binning)

有时，我们不想使用所有唯一值来确定组，而是想将数据“分箱(bin)”为更粗糙的组。 您始终可以创建自定义坐标，但是xarray通过groupby_bins()方法简化了此过程。

In [7]: x_bins = [0,25,50]

In [8]: ds.groupby_bins('x', x_bins).groups
Out[8]: 
{Interval(0, 25, closed='right'): [0, 1],
 Interval(25, 50, closed='right'): [2, 3]}

分箱是通过pandas.cut()实现的，其文档详细说明了分箱的分配方式。 从上面的示例中可以看出，默认情况下，会将标识精确的箱子范围的字符串用来标识箱子。 要取代此行为，可以显式指定箱子标签。 在这里，我们选择浮点数标签来标识箱子中心：

In [9]: x_bin_labels = [12.5,37.5]

In [10]: ds.groupby_bins('x', x_bins, labels=x_bin_labels).groups
Out[10]: {12.5: [0, 1], 37.5: [2, 3]}

应用(Apply)

要将功能函数应用于每个组，可以使用灵活的map()方法。结果对象将沿着组轴自动重新串联在一起：

In [11]: def standardize(x):
   ....:     return (x - x.mean()) / x.std()
   ....: 

In [12]: arr.groupby('letters').map(standardize)
Out[12]: 
<xarray.DataArray 'foo' (x: 4, y: 3)>
array([[-1.23 ,  1.937, -0.726],
       [ 1.42 , -0.46 , -0.607],
       [-0.191,  1.214, -1.376],
       [ 0.339, -0.302, -0.019]])
Coordinates:
  * x        (x) int64 10 20 30 40
    letters  (x) <U1 'a' 'b' 'b' 'a'
Dimensions without coordinates: y

GroupBy对象还具有reduce()方法和诸如mean()之类的方法作为应用聚合函数的快捷方式：

In [13]: arr.groupby('letters').mean(dim='x')
Out[13]: 
<xarray.DataArray 'foo' (letters: 2, y: 3)>
array([[0.335, 0.67 , 0.354],
       [0.674, 0.609, 0.23 ]])
Coordinates:
  * letters  (letters) object 'a' 'b'
Dimensions without coordinates: y

因此，使用groupby也是聚合除提供的维度之外的所有维度的便捷捷径：

In [14]: ds.groupby('x').std(...)
Out[14]: 
<xarray.Dataset>
Dimensions:  (x: 4)
Coordinates:
  * x        (x) int64 10 20 30 40
    letters  (x) <U1 'a' 'b' 'b' 'a'
Data variables:
    foo      (x) float64 0.3684 0.2554 0.2931 0.06957

在此处使用省略号（…）表示我们要缩减所有其他纬度

First与last

当前仅在groupby对象上有两种特殊的聚合操作：first与last。 这提供了沿分组维度的分组值的第一个或最后一个示例：

In [15]: ds.groupby('letters').first(...)
Out[15]: 
<xarray.Dataset>
Dimensions:  (letters: 2, y: 3)
Coordinates:
  * letters  (letters) object 'a' 'b'
Dimensions without coordinates: y
Data variables:
    foo      (letters, y) float64 0.127 0.9667 0.2605 0.8972 0.3767 0.3362

默认情况下，它们跳过缺少的值（使用skipna进行控制）。

分组算数(Grouped arithmetic)

GroupBy对象还支持一组有限的二元算术运算，提供了映射所有唯一标签的快捷方式。 对于(GroupBy, Dataset)和(GroupBy, DataArray)对，只要数据集或数据数组使用唯一的分组值作为其索引坐标之一，就支持二元算术运算。例如:

In [16]: alt = arr.groupby('letters').mean(...)

In [17]: alt
Out[17]: 
<xarray.DataArray 'foo' (letters: 2)>
array([0.453, 0.504])
Coordinates:
  * letters  (letters) object 'a' 'b'

In [18]: ds.groupby('letters') - alt
Out[18]: 
<xarray.Dataset>
Dimensions:  (x: 4, y: 3)
Coordinates:
  * x        (x) int64 10 20 30 40
    letters  (x) <U1 'a' 'b' 'b' 'a'
Dimensions without coordinates: y
Data variables:
    foo      (x, y) float64 -0.3261 0.5137 -0.1926 ... -0.08002 -0.005036

最后一行大致等同于：

results = []
for label, group in ds.groupby('letters'):
    results.append(group - alt.sel(x=label))
xr.concat(results, dim='x')

对维度进行分组时，可以使用squeeze参数控制是压缩维度还是在每个组上保留长度为1的纬度：

In [19]: next(iter(arr.groupby('x')))
Out[19]: 
(10, <xarray.DataArray 'foo' (y: 3)>
 array([0.127, 0.967, 0.26 ])
 Coordinates:
     x        int64 10
     letters  <U1 'a'
 Dimensions without coordinates: y)

In [20]: next(iter(arr.groupby('x', squeeze=False)))
Out[20]: 
(10, <xarray.DataArray 'foo' (x: 1, y: 3)>
 array([[0.127, 0.967, 0.26 ]])
 Coordinates:
   * x        (x) int64 10
     letters  (x) <U1 'a'
 Dimensions without coordinates: y)

尽管xarray会在您使用apply时尝试自动将维度转换(transpose)回其原始顺序，但是有时设置squeeze = False可以确保所有原始维度保持不变。

以后也可以使用Dataset或DataArray的 squeeze()方法显式地进行压缩纬度。

多维分组(Multidimensional Grouping)

许多数据集的多维坐标变量（例如经度）与逻辑网格尺寸（例如nx，ny）不同。 这些变量在CF conventions下有效。 Xarray支持对多维坐标变量进行groupby操作：

In [21]: da = xr.DataArray([[0,1],[2,3]],
   ....:     coords={'lon': (['ny','nx'], [[30,40],[40,50]] ),
   ....:             'lat': (['ny','nx'], [[10,10],[20,20]] ),},
   ....:     dims=['ny','nx'])
   ....: 

In [22]: da
Out[22]: 
<xarray.DataArray (ny: 2, nx: 2)>
array([[0, 1],
       [2, 3]])
Coordinates:
    lon      (ny, nx) int64 30 40 40 50
    lat      (ny, nx) int64 10 10 20 20
Dimensions without coordinates: ny, nx

In [23]: da.groupby('lon').sum(...)
Out[23]: 
<xarray.DataArray (lon: 3)>
array([0, 3, 3])
Coordinates:
  * lon      (lon) int64 30 40 50

In [24]: da.groupby('lon').map(lambda x: x - x.mean(), shortcut=False)
Out[24]: 
<xarray.DataArray (ny: 2, nx: 2)>
array([[ 0. , -0.5],
       [ 0.5,  0. ]])
Coordinates:
    lon      (ny, nx) int64 30 40 40 50
    lat      (ny, nx) int64 10 10 20 20
Dimensions without coordinates: ny, nx

因为多维组具有生成大量bin的能力，所以可能希望通过groupby_bins()进行粗略分类：

In [25]: da.groupby_bins('lon', [0,45,50]).sum()
Out[25]: 
<xarray.DataArray (lon_bins: 2)>
array([3, 3])
Coordinates:
  * lon_bins  (lon_bins) object (0, 45] (45, 50]

这些方法按 lon 值分组。 也可以通过 stacking 多个维度，应用到函数，然后 unstacking 结果来对网格中的每个像元进行分组，而不需管其值如何：

In [26]: stacked = da.stack(gridcell=['ny', 'nx'])

In [27]: stacked.groupby('gridcell').sum(...).unstack('gridcell')
Out[27]: 
<xarray.DataArray (ny: 2, nx: 2)>
array([[0, 1],
       [2, 3]])
Coordinates:
    lon      (ny, nx) int64 30 40 40 50
    lat      (ny, nx) int64 10 10 20 20
  * ny       (ny) int64 0 1
  * nx       (nx) int64 0 1