Question

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

我有一个数据数组n，其维度为('t', 'x', 'z')，我沿z维度将其分为5个块：

<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一次只对一个块执行此操作，因为我的笔记本电脑上随时只有几个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制作的数据块中？

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

Answer 1

由于它是3D阵列，因此我假设出现以下情况。

水分子在x-z平面上的速度（960μmx 512μm）随时间（801 fs）变化。求出x-z平面每个元素在整个时间的速度的均方根值。

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)>

使用XArray + Dask的内存错误-使用groupby或apply_ufunc吗？

1 个答案: