Numpy 如何使用xarray实现内存高效的多维groupby/stack？

我有一个大型的

np.float64

时间序列，频率为5分钟（大小约为2500000~=24年）

我使用

Xarray

在内存中表示它，时间维度名为

'time'

我想按

'time.hour'

和

'time.dayofyear'

进行分组（反之亦然），并从时间序列中删除它们的平均值

为了有效地完成这项工作，我需要将时间序列重新排序为一个新的

xr.DataArray

，维度为

['hour'、'dayofyear'、'rest']

我编写了一个函数，它可以处理

Xarray

的

GroupBy

对象，并设法做到这一点，尽管这需要大量内存

我有一台32GB内存的机器，我仍然从

numpy

获得

MemoryError

我知道代码是有效的，因为我在原始时间序列的每小时重新采样版本上使用了它。下面是代码：

def time_series_stack(time_da, time_dim='time', grp1='hour', grp2='dayofyear'):
    """Takes a time-series xr.DataArray objects and reshapes it using
    grp1 and grp2. outout is a xr.Dataset that includes the reshaped DataArray
    , its datetime-series and the grps."""
    import xarray as xr
    import numpy as np
    import pandas as pd

    # try to infer the freq and put it into attrs for later reconstruction:
    freq = pd.infer_freq(time_da[time_dim].values)
    name = time_da.name
    time_da.attrs['freq'] = freq
    attrs = time_da.attrs

    # drop all NaNs:
    time_da = time_da.dropna(time_dim)

    # group grp1 and concat:
    grp_obj1 = time_da.groupby(time_dim + '.' + grp1)
    s_list = []
    for grp_name, grp_inds in grp_obj1.groups.items():
        da = time_da.isel({time_dim: grp_inds})
        s_list.append(da)
    grps1 = [x for x in grp_obj1.groups.keys()]
    stacked_da = xr.concat(s_list, dim=grp1)
    stacked_da[grp1] = grps1

    # group over the concatenated da and concat again:
    grp_obj2 = stacked_da.groupby(time_dim + '.' + grp2)
    s_list = []
    for grp_name, grp_inds in grp_obj2.groups.items():
        da = stacked_da.isel({time_dim: grp_inds})
        s_list.append(da)
    grps2 = [x for x in grp_obj2.groups.keys()]
    stacked_da = xr.concat(s_list, dim=grp2)
    stacked_da[grp2] = grps2

    # numpy part:
    # first, loop over both dims and drop NaNs, append values and datetimes:
    vals = []
    dts = []
    for i, grp1_val in enumerate(stacked_da[grp1]):
        da = stacked_da.sel({grp1: grp1_val})
        for j, grp2_val in enumerate(da[grp2]):
            val = da.sel({grp2: grp2_val}).dropna(time_dim)
            vals.append(val.values)
            dts.append(val[time_dim].values)

    # second, we get the max of the vals after the second groupby:
    max_size = max([len(x) for x in vals])

    # we fill NaNs and NaT for the remainder of them:
    concat_sizes = [max_size - len(x) for x in vals]
    concat_arrys = [np.empty((x)) * np.nan for x in concat_sizes]
    concat_vals = [np.concatenate(x) for x in list(zip(vals, concat_arrys))]
    # 1970-01-01 is the NaT for this time-series:
    concat_arrys = [np.zeros((x), dtype='datetime64[ns]')
                    for x in concat_sizes]
    concat_dts = [np.concatenate(x) for x in list(zip(dts, concat_arrys))]
    concat_vals = np.array(concat_vals)
    concat_dts = np.array(concat_dts)

    # finally , we reshape them:
    concat_vals = concat_vals.reshape((stacked_da[grp1].shape[0],
                                       stacked_da[grp2].shape[0],
                                       max_size))
    concat_dts = concat_dts.reshape((stacked_da[grp1].shape[0],
                                     stacked_da[grp2].shape[0],
                                     max_size))

    # create a Dataset and DataArrays for them:
    sda = xr.Dataset()
    sda.attrs = attrs
    sda[name] = xr.DataArray(concat_vals, dims=[grp1, grp2, 'rest'])
    sda[time_dim] = xr.DataArray(concat_dts, dims=[grp1, grp2, 'rest'])
    sda[grp1] = grps1
    sda[grp2] = grps2
    sda['rest'] = range(max_size)
    return sda

因此，对于2500000项的时间序列，

numpy

抛出了

MemoryError

，所以我猜这一定是我的记忆瓶颈。我能做些什么来解决这个问题

---

Dask

能帮我吗？如果是这样的话，我该如何实现它呢？

像你一样，我在输入一个小的时间序列（10000长）时没有问题地运行了它。然而，当输入一个100000长的时间序列时，

xr.DataArray

for循环的

grp_obj2

会跑掉并使用系统的所有内存

这就是我用来生成时间序列的内容

xr.DataArray

：

n = 10**5
times = np.datetime64('2000-01-01') + np.arange(n) * np.timedelta64(5,'m')
data = np.random.randn(n)
time_da = xr.DataArray(data, name='rand_data', dims=('time'), coords={'time': times})
# time_da.to_netcdf('rand_time_series.nc')

正如你所指出的，Dask将是解决问题的一种方法，但我目前看不到一条明确的道路。。。 通常，Dask的问题类型为：

将输入设置为来自文件的数据集（如NetCDF）。这将不会在内存中加载文件，但允许Dask从磁盘一次提取一个数据块

使用

dask.delayed

或

dask.futures

方法为整个代码体定义所有计算，直到写入输出。这就是允许Dask将一小段数据分块读取然后写入的原因

计算一块工作并立即将输出写入新的数据集文件。实际上，一次只能将一个输入数据块转换为一个输出数据块（但也可以线程化/并行化）

我尝试导入Dask并将输入的

time\u da

xr.DataArray

分解成块供Dask处理，但没有任何帮助。据我所知，行

stacked_da=xr.concat（s_list，dim=grp1）

强制Dask在内存中创建

stacked_da

的完整副本，以及更多内容。。。 解决此问题的一种方法是将

stacked_da

写入磁盘，然后立即再次读取：

##For group1
xr.concat(s_list, dim=grp1).to_netcdf('stacked_da1.nc')
stacked_da = xr.load_dataset('stacked_da1.nc')
stacked_da[grp1] = grps1

##For group2
xr.concat(s_list, dim=grp2).to_netcdf('stacked_da2.nc')
stacked_da = xr.load_dataset('stacked_da2.nc')
stacked_da[grp2] = grps2

但是，

stacked_da1.nc

的文件大小为19MB，

stacked_da2.nc

的文件大小为6.5GB。这是用于包含100000个元素的

time\u da

。。。所以很明显有点不对劲

---

最初，听起来好像你想从时间序列数据中减去各组的平均值。看起来Xarray docs有一个这样的例子

关键是分组一次，在组上循环，然后在每个组上再次分组并将其附加到列表中

接下来，我对这些组使用

pd.MultiIndex.from_product

无需内存问题和

Dask

，只需几秒钟即可运行

以下是代码，请享受：

def time_series_stack(time_da, time_dim='time', grp1='hour', grp2='month',
                      plot=True):
    """Takes a time-series xr.DataArray objects and reshapes it using
    grp1 and grp2. output is a xr.Dataset that includes the reshaped DataArray
    , its datetime-series and the grps. plots the mean also"""
    import xarray as xr
    import pandas as pd
    # try to infer the freq and put it into attrs for later reconstruction:
    freq = pd.infer_freq(time_da[time_dim].values)
    name = time_da.name
    time_da.attrs['freq'] = freq
    attrs = time_da.attrs
    # drop all NaNs:
    time_da = time_da.dropna(time_dim)
    # first grouping:
    grp_obj1 = time_da.groupby(time_dim + '.' + grp1)
    da_list = []
    t_list = []
    for grp1_name, grp1_inds in grp_obj1.groups.items():
        da = time_da.isel({time_dim: grp1_inds})
        # second grouping:
        grp_obj2 = da.groupby(time_dim + '.' + grp2)
        for grp2_name, grp2_inds in grp_obj2.groups.items():
            da2 = da.isel({time_dim: grp2_inds})
            # extract datetimes and rewrite time coord to 'rest':
            times = da2[time_dim]
            times = times.rename({time_dim: 'rest'})
            times.coords['rest'] = range(len(times))
            t_list.append(times)
            da2 = da2.rename({time_dim: 'rest'})
            da2.coords['rest'] = range(len(da2))
            da_list.append(da2)
    # get group keys:
    grps1 = [x for x in grp_obj1.groups.keys()]
    grps2 = [x for x in grp_obj2.groups.keys()]
    # concat and convert to dataset:
    stacked_ds = xr.concat(da_list, dim='all').to_dataset(name=name)
    stacked_ds[time_dim] = xr.concat(t_list, 'all')
    # create a multiindex for the groups:
    mindex = pd.MultiIndex.from_product([grps1, grps2], names=[grp1, grp2])
    stacked_ds.coords['all'] = mindex
    # unstack:
    ds = stacked_ds.unstack('all')
    ds.attrs = attrs
    return ds

---

我没有提到它，但我希望时间序列重组的原因是1）我可以看到季节性与我选择的组。例如，如果我选择时间序列的2D pcolormesh图（例如，x:hour of the day，y:dayofyear），它是bc，我怀疑信号具有每日和每年的季节性。2） Xarray可以很容易地减去一个组的平均值，但到目前为止还不存在两个或多个坐标的分组（在Xarray生态系统中，我很高兴在这一点上是错误的）。现在我想起来了，这个解决方案是琐碎的、艰难的、无聊的。我可以将原始的时间序列划分成更小的块（我可以根据经验测试它的块大小是否正确，这样它就不会阻塞我的记忆），然后运行函数，然后对输出进行压缩。你可能不需要重新整形就可以实现这一点。可以这样做：1）从原始数据中按小时分组，减去每小时平均数，然后绘制图2）然后从原始数据中按天分组，减去每日平均数，然后绘制图I很难准确地描绘出你的最终结果应该是什么样的，所以我可能知道步骤的顺序，或者逻辑不正确。你也可以投票给我的答案吗？我花了大量时间研究这个……仅供参考，它根本没有效率（在组上循环），但它对我来说已经足够好了：-）