重叠数组在数据数组上的并行计算

时间:2019-12-13 08:36:04

标签: python dask dask-delayed

我想使用dask在大型数组上并行化数据帧的计算。该计算是局部的,可以在map_overlap的上下文中完成,每次计算都返回一个数据帧。最终结果将是groupby + sum的数据帧组合,大概是通过dask数据帧完成的。

我看不出这个想法如何适合map_block / map_overlapgufunc,但是我可能会遗漏一些东西。

以下尝试使用对to_delayed的调用来解决总和问题。这种情况是正确的,需要块操作。它不容易扩展到重叠类型的操作。另外,我认为from_delayed调用可能会触发部分计算,这并不理想。

是否还有其他类似问题的处理方式? 有没有适用于重叠计算的方法?

import dask
import dask.array as da
import dask.dataframe as dd
import numpy as np
import pandas as pd
from scipy import ndimage

def sum_chunk(image_chunk, labels_chunk, index):
    total = ndimage.sum(image_chunk, labels_chunk, index)
    return pd.DataFrame({'label': index, 'sum': total})

def ndimage_sum(image, labels, index):
    args_list = list(zip(
        image.to_delayed().ravel(), 
        labels.to_delayed().ravel(),
    ))
    results = [dask.delayed(sum_chunk)(*args, index) for args in args_list]
    df = dd.from_delayed(results)  # Wait -- does this trigger a computation?
    return df.groupby('label').sum()

image = da.arange(16, chunks=4)

labels = np.zeros(16, dtype=np.int)
labels[2:4] = 1
labels[5:9] = 2
labels[15] = 3
labels = da.from_array(labels, chunks=4)

index = da.from_array([1,2])

ndimage_sum(image, labels, index).compute()

0 个答案:

没有答案