按单个维度分组对于xarray DataArrays很好：
d = xr.DataArray([1, 2, 3], coords={'a': ['x', 'x', 'y']}, dims=['a'])
d.groupby('a').mean()) # -> DataArray (a: 2) array([1.5, 3. ])`
但是，仅单个维度支持此功能，因此无法按多个维度分组：
d = DataAssembly([[1, 2, 3], [4, 5, 6]],
coords={'a': ('multi_dim', ['a', 'b']), 'c': ('multi_dim', ['c', 'c']), 'b': ['x', 'y', 'z']},
dims=['multi_dim', 'b'])
d.groupby(['a', 'b']) # TypeError: `group` must be an xarray.DataArray or the name of an xarray variable or dimension
我只有一个效率不高的解决方案，该解决方案手动执行for循环：
a, b = np.unique(d['a'].values), np.unique(d['b'].values)
result = xr.DataArray(np.zeros([len(a), len(b)]), coords={'a': a, 'b': b}, dims=['a', 'b'])
for a, b in itertools.product(a, b):
cells = d.sel(a=a, b=b)
merge = cells.mean()
result.loc[{'a': a, 'b': b}] = merge
# result = DataArray (a: 2, b: 2)> array([[2., 3.], [5., 6.]])
# Coordinates:
# * a (a) <U1 'x' 'y'
# * b (b) int64 0 1
但是，对于较大的阵列，这太慢了。 有更有效/直接的解决方法吗？
答案 0 :(得分：3)
我建立了一个手动解决方案。为了提高效率，我舍弃了所有xarray并手动重建索引和值。使用更多xarray的任何更改（例如使用sel，将单元重新包装到DataArray中；另请参见https://github.com/pydata/xarray/issues/2452）都会导致速度严重下降。
import itertools
from collections import defaultdict
import numpy as np
import xarray as xr
from xarray import DataArray
class DataAssembly(DataArray):
def multi_dim_groupby(self, groups, apply):
# align
groups = sorted(groups, key=lambda group: self.dims.index(self[group].dims[0]))
# build indices
groups = {group: np.unique(self[group]) for group in groups}
group_dims = {self[group].dims: group for group in groups}
indices = defaultdict(lambda: defaultdict(list))
result_indices = defaultdict(dict)
for group in groups:
for index, value in enumerate(self[group].values):
indices[group][value].append(index)
if value not in result_indices[group]: # if captured once, it will be "grouped away"
index = max(result_indices[group].values()) + 1 if len(result_indices[group]) > 0 else 0
result_indices[group][value] = index
coords = {coord: (dims, value) for coord, dims, value in walk_coords(self)}
def simplify(value):
return value.item() if value.size == 1 else value
def indexify(dict_indices):
return [(i,) if isinstance(i, int) else tuple(i) for i in dict_indices.values()]
# group and apply
# making this a DataArray right away and then inserting through .loc would slow things down
result = np.zeros([len(indices) for indices in result_indices.values()])
result_coords = {coord: (dims, [None] * len(result_indices[group_dims[dims]]))
for coord, (dims, value) in coords.items()}
for values in itertools.product(*groups.values()):
group_values = dict(zip(groups.keys(), values))
self_indices = {group: indices[group][value] for group, value in group_values.items()}
values_indices = indexify(self_indices)
cells = self.values[values_indices] # using DataArray would slow things down. thus we pass coords as kwargs
cells = simplify(cells)
cell_coords = {coord: (dims, value[self_indices[group_dims[dims]]])
for coord, (dims, value) in coords.items()}
cell_coords = {coord: (dims, simplify(np.unique(value))) for coord, (dims, value) in cell_coords.items()}
# ignore dims when passing to function
passed_coords = {coord: value for coord, (dims, value) in cell_coords.items()}
merge = apply(cells, **passed_coords)
result_idx = {group: result_indices[group][value] for group, value in group_values.items()}
result[indexify(result_idx)] = merge
for coord, (dims, value) in cell_coords.items():
if isinstance(value, np.ndarray): # multiple values for coord -> ignore
if coord in result_coords: # delete from result coords if not yet deleted
del result_coords[coord]
continue
assert dims == result_coords[coord][0]
coord_index = result_idx[group_dims[dims]]
result_coords[coord][1][coord_index] = value
# re-package
result = type(self)(result, coords=result_coords, dims=list(itertools.chain(*group_dims.keys())))
return result
方法multi_dim_groupby执行分组并一步应用。
传递的apply方法可以通过以坐标命名的参数接受组坐标（或者通过将**_放在函数头中来忽略坐标）。
它不是特别漂亮，不能涵盖所有可能的情况，但至少涵盖以下测试用例：
import DataAssembly
class TestMultiDimGroupby:
def test_unique_values(self):
d = DataAssembly([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]],
coords={'a': ['a', 'b', 'c', 'd'],
'b': ['x', 'y', 'z']},
dims=['a', 'b'])
g = d.multi_dim_groupby(['a', 'b'], lambda x, **_: x)
assert g.equals(d)
def test_nonunique_singledim(self):
d = DataAssembly([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]],
coords={'a': ['a', 'a', 'b', 'b'],
'b': ['x', 'y', 'z']},
dims=['a', 'b'])
g = d.multi_dim_groupby(['a', 'b'], lambda x, **_: x.mean())
assert g.equals(DataAssembly([[2.5, 3.5, 4.5], [8.5, 9.5, 10.5]],
coords={'a': ['a', 'b'], 'b': ['x', 'y', 'z']},
dims=['a', 'b']))
def test_nonunique_adjacentcoord(self):
d = DataAssembly([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]],
coords={'a': ('adim', ['a', 'a', 'b', 'b']),
'aa': ('adim', ['a', 'b', 'a', 'b']),
'b': ['x', 'y', 'z']},
dims=['adim', 'b'])
g = d.multi_dim_groupby(['a', 'b'], lambda x, **_: x.mean())
assert g.equals(DataAssembly([[2.5, 3.5, 4.5], [8.5, 9.5, 10.5]],
coords={'adim': ['a', 'b'], 'b': ['x', 'y', 'z']},
dims=['adim', 'b'])), \
"adjacent coord aa should be discarded due to non-mappability"
def test_unique_values_swappeddims(self):
d = DataAssembly([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]],
coords={'a': ['a', 'b', 'c', 'd'],
'b': ['x', 'y', 'z']},
dims=['a', 'b'])
g = d.multi_dim_groupby(['b', 'a'], lambda x, **_: x)
assert g.equals(d)
答案 1 :(得分：0)
我不知道它如何与速度进行比较，也没有足够的时间为这个特定的问题实例提供完整的解决方案，但是当我寻找解决方法时，我发现这个问题和答案非常有帮助遍历xarray中的多个维度，并希望分享我最终采用的方法。最终，我根据@RyanAbernathy的example code使用了维度堆叠：
import xarray as xr
import numpy as np
# create an example dataset
da = xr.DataArray(np.random.rand(10,30,40), dims=['dtime', 'x', 'y'])
# define a function to compute a linear trend of a timeseries
def linear_trend(x):
pf = np.polyfit(x.time, x, 1)
# we need to return a dataarray or else xarray's groupby won't be happy
return xr.DataArray(pf[0])
# stack lat and lon into a single dimension called allpoints
stacked = da.stack(allpoints=['x','y'])
# apply the function over allpoints to calculate the trend at each point
trend = stacked.groupby('allpoints').apply(linear_trend)
# unstack back to lat lon coordinates
trend_unstacked = trend.unstack('allpoints')
结合一些groupby包装器来计算多个groupby：
def _calc_allpoints(ds, function):
"""
Helper function to do a pixel-wise calculation that requires using x and y dimension values
as inputs. This version does the computation over all available timesteps as well.
"""
# note: the below code will need to be generalized for other dimensions
def _time_wrapper(gb):
gb = gb.groupby('dtime', squeeze=False).apply(function)
return gb
# stack x and y into a single dimension called allpoints
stacked = ds.stack(allpoints=['x','y'])
# groupby time and apply the function over allpoints to calculate the trend at each point
newelev = stacked.groupby('allpoints', squeeze=False).apply(_time_wrapper)
# unstack back to x y coordinates
ds = newelev.unstack('allpoints')
return ds
其中function是您正在使用的任何函数（例如linear_trend）