在 dask 中广播和调整数组大小

Question

我想广播一维 dask 数组和一个 2D dask 数组。

具体使用 numpy 会是这样的：

a = np.random.rand(20000, 3)
b = np.random.rand(16)

我想要一个大小为 (20000,16,3) 的 3D 数组，我们称之为c。因此，对于b 的每个值，我们将有 (20000,3) 个值乘以a*b[index]，其中index=0,1,...,15。在 numpy 中，使用函数 resize 非常简单。但是，调整大小在 dask 中不存在。任何人都知道我如何在 dask 上做到这一点？该数组也可以是内部带有 dask 数组的 xarray，因此，如果有人知道如何使用 xarray 进行操作，也将不胜感激。

干杯

Answer 1

如果我正确理解你的问题，你想达到这个结果

import numpy as np

a = np.random.rand(20000, 3)
b = np.random.rand(16)

result = a[:, np.newaxis, :] * b[np.newaxis, :, np.newaxis]
result.shape  # (20000, 16, 3)

该代码也直接适用于dask.array

import dask.array as da

dsk_a = da.from_array(a)
dsk_b = da.from_array(b)

result_dask = dsk_a[:, np.newaxis, :] * dsk_b[np.newaxis, :, np.newaxis]

result_dask.shape  # (20000, 16, 3)
(result_dask.compute() == result).all()  # True

如果我误解了您的问题，请告诉我。如果我有，那么如果您提供可提供所需结果的有效numpy 代码将会很有帮助。