python - 使用 xarray + dask 的内存错误 - 使用 groupby 或 apply_ufunc？

Question

我使用 xarray 作为分析流体湍流数据的工作流程的基础，但我无法正确利用 dask 来限制笔记本电脑的内存使用。

我有一个维度为 ('t', 'x', 'z') 的数据数组 n，我已经沿着 z 将其分成 5 个 block 维度:

<xarray.DataArray 'n' (t: 801, x: 960, z: 512)>
dask.array<shape=(801, 960, 512), dtype=float32, chunksize=(801, 960, 5)>
Coordinates:
* t              (t) int64 200 201 202 203 204 205 206 207 208 209 210 211 ...
* x              (x) int64 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 ...
* z              (z) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 ...  

我想计算 n 在 t 上的均方根波动，并返回维度为 ('x', 'z') 的简化数据数组。我想利用 dask 一次只对一个 block 执行此操作，因为我的笔记本电脑上任何时候都只有几 GB 的 RAM 可用。

我写了一个通用的 ufunc 来计算 dask 数组的均方根:

def rms_gufunc(x, axis):
    """Generalized ufunc to calculate root mean squared value of a 1d dask array."""
    squares = np.square(x)
    return np.sqrt(squares.mean(axis=-1))

但现在我不确定应用它的最佳方法是什么。据我所知，我可以使用 (1) xarray.apply_ufunc，或 (2) groupby.reduce。

<强>1。使用apply_ufunc:

我可以使用 xarray.apply_ufunc 应用这个函数:

def rms(da, dim=None):
    """
    Reduces a dataarray by calculating the root mean square along dimension dim.
    """

    if dim is None:
        raise ValueError('Must supply a dimension along which to calculate rms')

    return xr.apply_ufunc(rms_gufunc, da,
                          input_core_dims=[[dim]],
                          dask='parallelized', output_dtypes=[da.dtype])

n_rms = rms(data['n'], dim='t')
n_rms.load()  # Trigger computation

这似乎可行，但似乎比必要的更复杂？

<强>2。使用 groupby.reduce:

xarray 文档似乎表明这是一个“split-apply-combine”操作，我应该能够通过类似的东西来完成

n_rms = data['n'].groupby('z').reduce(rms_gufunc, dim='t')

然而，这会导致 MemoryError，而且我很确定这不是我想要通过 groupby 步骤实现的。我是否应该使用 groupby_bins 将数据分箱到我沿 z 制作的 block 中？

我想知道 a) 我是否正确使用 apply_ufunc，以及 b) 我如何使用 groupby 做同样的事情。

Answer 1

因为它是一个 3D 数组，所以我假设以下情况。

水分子在 x-z 平面 (960 μm x 512 μm) 上的速度随时间变化 (801 fs)。找出所有时间 x-z 平面每个元素的速度 rmsf。

numpy 代码为:

xz_plane_rmsf = (data ** 2).mean(axis=0)

其中数据是 shape=(801, 960, 512) 的 3D numpy 数组。data 的第 0、1 和 2 维表示时间、x 坐标和 z 坐标。 data 的每个元素表示时间 t 和坐标 x 和 z 处水分子的平均速度。

Dask 数组的等效代码是:

# Make lazy array
xz_plane_rmsf = (data ** 2).mean(axis=0)
# Evaluate the array
xz_plane_rmsf = xz_plane_rmsf.compute()

其中 data 是一个 3D Dask 数组。

唯一剩下的问题是将 xarray 转换为 Dask 数组。我不使用 xarray，但看起来它已经是一个 Dask 数组:

<xarray.DataArray 'n' (t: 801, x: 960, z: 512)>
dask.array<shape=(801, 960, 512), dtype=float32, chunksize=(801, 960, 5)>