在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对我的数据执行了一些非平凡过滤，并在条件为True且为False时返回数组。 function可以是纯python或cython编译。对行的过滤操作很复杂，可能依赖于行中的先前值，这意味着我无法以逐个元素的方式对整个数组进行操作

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

Answer 1

Dask解决方案

你可以通过逐行分组数组，调用map_blocks，然后计算结果来使用dask.array

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