GroupBy: split-apply-combine

xarray supports “group by” operations with the same API as pandas toimplement the split-apply-combine strategy:

  • Split your data into multiple independent groups.

  • Apply some function to each group.

  • Combine your groups back into a single data object.

Group by operations work on both Dataset andDataArray objects. Most of the examples focus on grouping bya single one-dimensional variable, although support for groupingover a multi-dimensional variable has recently been implemented. Note that forone-dimensional data, it is usually faster to rely on pandas’ implementation ofthe same pipeline.

Split

Let’s create a simple example dataset:

  1. In [1]: ds = xr.Dataset({'foo': (('x', 'y'), np.random.rand(4, 3))},
  2.    ...:                 coords={'x': [10, 20, 30, 40],
  3.    ...:                         'letters': ('x', list('abba'))})
  4.    ...: 
  5.  
  6. In [2]: arr = ds['foo']
  7.  
  8. In [3]: ds
  9. Out[3]: 
  10. <xarray.Dataset>
  11. Dimensions:  (x: 4, y: 3)
  12. Coordinates:
  13.   * x        (x) int64 10 20 30 40
  14.     letters  (x) <U1 'a' 'b' 'b' 'a'
  15. Dimensions without coordinates: y
  16. Data variables:
  17.     foo      (x, y) float64 0.127 0.9667 0.2605 0.8972 ... 0.543 0.373 0.448

If we groupby the name of a variable or coordinate in a dataset (we can alsouse a DataArray directly), we get back a GroupBy object:

  1. In [4]: ds.groupby('letters')
  2. Out[4]: <xarray.core.groupby.DatasetGroupBy at 0x7f3425a45a20>

This object works very similarly to a pandas GroupBy object. You can viewthe group indices with the groups attribute:

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

You can also iterate over groups in (label, group) pairs:

  1. In [6]: list(ds.groupby('letters'))
  2. Out[6]: 
  3. [('a', <xarray.Dataset>
  4.   Dimensions:  (x: 2, y: 3)
  5.   Coordinates:
  6.     * x        (x) int64 10 40
  7.       letters  (x) <U1 'a' 'a'
  8.   Dimensions without coordinates: y
  9.   Data variables:
  10.       foo      (x, y) float64 0.127 0.9667 0.2605 0.543 0.373 0.448),
  11.  ('b', <xarray.Dataset>
  12.   Dimensions:  (x: 2, y: 3)
  13.   Coordinates:
  14.     * x        (x) int64 20 30
  15.       letters  (x) <U1 'b' 'b'
  16.   Dimensions without coordinates: y
  17.   Data variables:
  18.       foo      (x, y) float64 0.8972 0.3767 0.3362 0.4514 0.8403 0.1231)]

Just like in pandas, creating a GroupBy object is cheap: it does not actuallysplit the data until you access particular values.

Binning

Sometimes you don’t want to use all the unique values to determine the groupsbut instead want to “bin” the data into coarser groups. You could always createa customized coordinate, but xarray facilitates this via thegroupby_bins() method.

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

The binning is implemented via pandas.cut, whose documentation details howthe bins are assigned. As seen in the example above, by default, the bins arelabeled with strings using set notation to precisely identify the bin limits. Tooverride this behavior, you can specify the bin labels explicitly. Here wechoose float labels which identify the bin centers:

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

Apply

To apply a function to each group, you can use the flexibleapply() method. The resulting objects are automaticallyconcatenated back together along the group axis:

  1. In [11]: def standardize(x):
  2.    ....:     return (x - x.mean()) / x.std()
  3.    ....: 
  4.  
  5. In [12]: arr.groupby('letters').apply(standardize)
  6. Out[12]: 
  7. <xarray.DataArray 'foo' (x: 4, y: 3)>
  8. array([[-1.229778,  1.93741 , -0.726247],
  9.        [ 1.419796, -0.460192, -0.606579],
  10.        [-0.190642,  1.21398 , -1.376362],
  11.        [ 0.339417, -0.301806, -0.018995]])
  12. Coordinates:
  13.   * x        (x) int64 10 20 30 40
  14.     letters  (x) <U1 'a' 'b' 'b' 'a'
  15. Dimensions without coordinates: y

GroupBy objects also have a reduce() method andmethods like mean() as shortcuts for applying anaggregation function:

  1. In [13]: arr.groupby('letters').mean(dim='x')
  2. Out[13]: 
  3. <xarray.DataArray 'foo' (letters: 2, y: 3)>
  4. array([[0.334998, 0.669865, 0.354236],
  5.        [0.674306, 0.608502, 0.229662]])
  6. Coordinates:
  7.   * letters  (letters) object 'a' 'b'
  8. Dimensions without coordinates: y

Using a groupby is thus also a convenient shortcut for aggregating over alldimensions other than the provided one:

  1. In [14]: ds.groupby('x').std(xr.ALL_DIMS)
  2. Out[14]: 
  3. <xarray.Dataset>
  4. Dimensions:  (x: 4)
  5. Coordinates:
  6.   * x        (x) int64 10 20 30 40
  7.     letters  (x) <U1 'a' 'b' 'b' 'a'
  8. Data variables:
  9.     foo      (x) float64 0.3684 0.2554 0.2931 0.06957

First and last

There are two special aggregation operations that are currently only found ongroupby objects: first and last. These provide the first or last example ofvalues for group along the grouped dimension:

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

By default, they skip missing values (control this with skipna).

Grouped arithmetic

GroupBy objects also support a limited set of binary arithmetic operations, asa shortcut for mapping over all unique labels. Binary arithmetic is supportedfor (GroupBy, Dataset) and (GroupBy, DataArray) pairs, as long as thedataset or data array uses the unique grouped values as one of its indexcoordinates. For example:

  1. In [16]: alt = arr.groupby('letters').mean(xr.ALL_DIMS)
  2.  
  3. In [17]: alt
  4. Out[17]: 
  5. <xarray.DataArray 'foo' (letters: 2)>
  6. array([0.453033, 0.504157])
  7. Coordinates:
  8.   * letters  (letters) object 'a' 'b'
  9.  
  10. In [18]: ds.groupby('letters') - alt
  11. Out[18]: 
  12. <xarray.Dataset>
  13. Dimensions:  (x: 4, y: 3)
  14. Coordinates:
  15.   * x        (x) int64 10 20 30 40
  16.     letters  (x) <U1 'a' 'b' 'b' 'a'
  17. Dimensions without coordinates: y
  18. Data variables:
  19.     foo      (x, y) float64 -0.3261 0.5137 -0.1926 ... -0.08002 -0.005036

This last line is roughly equivalent to the following:

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

Squeezing

When grouping over a dimension, you can control whether the dimension issqueezed out or if it should remain with length one on each group by usingthe squeeze parameter:

  1. In [19]: next(iter(arr.groupby('x')))
  2. Out[19]: 
  3. (10, <xarray.DataArray 'foo' (y: 3)>
  4.  array([0.12697 , 0.966718, 0.260476])
  5.  Coordinates:
  6.      x        int64 10
  7.      letters  <U1 'a'
  8.  Dimensions without coordinates: y)
  1. In [20]: next(iter(arr.groupby('x', squeeze=False)))
  2. Out[20]: 
  3. (10, <xarray.DataArray 'foo' (x: 1, y: 3)>
  4.  array([[0.12697 , 0.966718, 0.260476]])
  5.  Coordinates:
  6.    * x        (x) int64 10
  7.      letters  (x) <U1 'a'
  8.  Dimensions without coordinates: y)

Although xarray will attempt to automaticallytranspose dimensions back into their original orderwhen you use apply, it is sometimes useful to set squeeze=False toguarantee that all original dimensions remain unchanged.

You can always squeeze explicitly later with the Dataset or DataArraysqueeze() methods.

Multidimensional Grouping

Many datasets have a multidimensional coordinate variable (e.g. longitude)which is different from the logical grid dimensions (e.g. nx, ny). Suchvariables are valid under the CF conventions. Xarray supports groupbyoperations over multidimensional coordinate variables:

  1. In [21]: da = xr.DataArray([[0,1],[2,3]],
  2.    ....:     coords={'lon': (['ny','nx'], [[30,40],[40,50]] ),
  3.    ....:             'lat': (['ny','nx'], [[10,10],[20,20]] ),},
  4.    ....:     dims=['ny','nx'])
  5.    ....: 
  6.  
  7. In [22]: da
  8. Out[22]: 
  9. <xarray.DataArray (ny: 2, nx: 2)>
  10. array([[0, 1],
  11.        [2, 3]])
  12. Coordinates:
  13.     lon      (ny, nx) int64 30 40 40 50
  14.     lat      (ny, nx) int64 10 10 20 20
  15. Dimensions without coordinates: ny, nx
  16.  
  17. In [23]: da.groupby('lon').sum(xr.ALL_DIMS)
  18. Out[23]: 
  19. <xarray.DataArray (lon: 3)>
  20. array([0, 3, 3])
  21. Coordinates:
  22.   * lon      (lon) int64 30 40 50
  23.  
  24. In [24]: da.groupby('lon').apply(lambda x: x - x.mean(), shortcut=False)
  25. Out[24]: 
  26. <xarray.DataArray (ny: 2, nx: 2)>
  27. array([[ 0. , -0.5],
  28.        [ 0.5,  0. ]])
  29. Coordinates:
  30.     lon      (ny, nx) int64 30 40 40 50
  31.     lat      (ny, nx) int64 10 10 20 20
  32. Dimensions without coordinates: ny, nx

Because multidimensional groups have the ability to generate a very largenumber of bins, coarse-binning via groupby_bins()may be desirable:

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

These methods group by lon values. It is also possible to groupby eachcell in a grid, regardless of value, by stacking multiple dimensions,applying your function, and then unstacking the result:

  1. In [26]: stacked = da.stack(gridcell=['ny', 'nx'])
  2.  
  3. In [27]: stacked.groupby('gridcell').sum().unstack('gridcell')
  4. Out[27]: 
  5. <xarray.DataArray (ny: 2, nx: 2)>
  6. array([[0, 1],
  7.        [2, 3]])
  8. Coordinates:
  9.     lon      (ny, nx) int64 30 40 40 50
  10.     lat      (ny, nx) int64 10 10 20 20
  11.   * ny       (ny) int64 0 1
  12.   * nx       (nx) int64 0 1