python - 在 numpy 行上并行化循环

Question

我需要将相同的函数应用于 numpy 数组中的每一行，并将结果再次存储在 numpy 数组中。

# states will contain results of function applied to a row in array
states = np.empty_like(array)

for i, ar in enumerate(array):
    states[i] = function(ar, *args)

# do some other stuff on states

function对我的数据进行一些重要的过滤，并在条件为真和条件为假时返回一个数组。function可以是纯 python 或 cython 编译。对行的过滤操作很复杂，并且可能取决于行中的先前值，这意味着我不能以逐个元素的方式对整个数组进行操作

例如，有没有办法在 dask 中做这样的事情？

Answer 1

解决方案

您可以使用 dask.array 将数组按行分块，调用map_blocks，然后计算结果

ar = ...
x = da.from_array(ar, chunks=(1, arr.shape[1]))
x.map_blocks(function, *args)
states = x.compute()

默认情况下这将使用线程，您可以通过以下方式使用进程

from dask.multiprocessing import get
states = x.compute(get=get)

池解决方案

然而，对于像这样令人尴尬的并行计算来说，dask 可能有点过头了，你可以使用线程池

from multiprocessing.pool import ThreadPool
pool = ThreadPool()

ar = ...
states = np.empty_like(array)

def f(i):
    states[i] = function(ar[i], *args)

pool.map(f, range(len(ar)))

您可以通过以下更改切换到流程

from multiprocessing import Pool
pool = Pool()

Answer 2

把你的函数变成一个通用函数：http ://docs.scipy.org/doc/numpy/reference/ufuncs.html 。

然后：states = function(array, *args)。