如何将dask数组的所有条目归零，小于top k

Question

我想将dask.array的所有元素归零，除了前几个元素。我该怎么做？

实施例

假设我有一个小的dask数组，如下所示：

import numpy as np
import dask.array as da
x = np.array([0, 4, 2, 3, 1])
x = da.from_array(x, chunks=(2,))

如何将除两个最大元素之外的所有元素归零？我想要以下内容：

>>> result.compute()
array([0, 4, 0, 3, 0])

Answer 1

您可以结合使用topk函数和inplace setitem

来完成此操作

top = x.topk(2)
x[x < top[-1]] = 0

>>> x.compute()
array([0, 4, 0, 3, 0])

请注意，这不会通过内存特别好地传输。如果您正在使用单机调度程序，那么您可能希望通过提前明确计算top来两次执行此操作：

top = x.topk(2)
top = top.compute()  # pass through data once to get top elements
x[x < top[-1]] = 0   # then pass through again applying filter

>>> x.compute()
array([0, 4, 0, 3, 0])

这只有在您尝试在单个计算机上流式传输大型数据集时才有意义，如果您使用的是分布式系统，则不应对此产生太大影响。