在xarray中，沿一个维度对多个坐标进行分组

Question

我有一个带有单个维度上多个坐标的xarray。在下面的例子中，coords a 和 b 在维度 dim1 上定义。如何使用在同一维度上定义的两个坐标groupby？与这个问题不同，我不是尝试沿不同的维度进行分组，而是沿单个维度进行分组。

import xarray as xr

d = xr.DataArray([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]],
    coords={
        'a': ('dim1',['A', 'A', 'B', 'B']),
        'b': ('dim1',['1', '2', '1', '2']),
        'c': ('dim2',['x', 'y', 'z'])
    },
    dims=['dim1', 'dim2'])
d.groupby(['a','b']) # this gives: TypeError: `group` must be an xarray.DataArray or the name of an xarray variable or dimension

- Patrickens

Answer 1

这是我的当前解决方案：

import numpy as np
import xarray as xr

def groupby_multicoords(da, fields):
    common_dim = da.coords[fields[0]].dims[0]
    tups_arr = np.empty(len(da[common_dim]), dtype=object)
    tups_arr[:] = list(zip(*(da[f].values for f in fields)))
    return da.assign_coords(grouping_zip=xr.DataArray(tups_arr, dims=common_dim)).groupby('grouping_zip')

然后，groupby_multicoords(da=d, fields=['a', 'b'])

然而，分组后仍然存在“grouping_zip”坐标。我希望用d.groupby(['a','b'])替换它。

- Oren Talmor

---

如果“a”和“b”的长度不一样呢？ - user29988

Answer 2

您可以使用.stack(new=[“dim1”,”dim2”)将它们堆叠成一个多级索引，然后按该维度进行分组。

- Maximilian

---

我不想按维度分组，而是按坐标分组。因此，我想通过类似于 ('A', '1') 或 ('B','2') 的分组方式来对 a 和 b 进行分组。在这种情况下，这两个坐标沿着同一维度定义。 - Patrickens

好的，感谢澄清。您可以将 a 和 b 转换为堆叠的维度坐标，并对其进行分组。不幸的是，目前无法将多个坐标传递给 groupby。 - Maximilian