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

如何解决在 dask 中广播和调整数组大小

我想广播一个一维 dask 数组和一个二维 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 中，使用函数调整大小非常简单。但是，在 dask 中不存在调整大小。任何人都知道我如何在 dask 上做到这一点？该数组也可以是带有 dask 数组的 xarray，因此如果有人知道如何使用 xarray 进行操作，也将不胜感激。

干杯

解决方法

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

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[:,:] * dsk_b[np.newaxis,np.newaxis]

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

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