Python 如何使用Dask在NumPy阵列的切片上执行并行操作？

我有一个大小为n_slice x 2048 x 3的坐标的numpy数组，其中n_slice是成千上万的。我想分别对每个2048 x 3切片应用以下操作

import numpy as np
from scipy.spatial.distance import pdist

# load coor from a binary xyz file, dcd format

n_slice, n_coor, _ = coor.shape
r = np.arange(n_coor)
dist = np.zeros([n_slice, n_coor, n_coor])

# this loop is what I want to parallelize, each slice is completely independent
for i in xrange(n_slice): 
    dist[i, r[:, None] < r] = pdist(coor[i])

但是简单地用

dcor

替换

coor

不会暴露并行性。我可以看到为每个切片设置并行线程，但是如何利用Dask来处理并行性呢

下面是使用

concurrent.futures

import concurrent.futures
import multiprocessing

n_cpu = multiprocessing.cpu_count()

def get_dist(coor, dist, r):
    dist[r[:, None] < r] = pdist(coor)

# load coor from a binary xyz file, dcd format

n_slice, n_coor, _ = coor.shape
r = np.arange(n_coor)
dist = np.zeros([n_slice, n_coor, n_coor])

with concurrent.futures.ThreadPoolExecutor(max_workers=n_cpu) as executor:
    for i in xrange(n_slice):
        executor.submit(get_dist, cool[i], dist[i], r)

import concurrent.futures
导入多处理
n\u cpu=多处理。cpu\u计数（）
def get_区（coor，dist，r）：
dist[r[：，None]

这个问题可能不适合Dask，因为没有块间计算。

map\u块

该方法可能有帮助：

dcoor.map_blocks(pdist)

非均匀阵列 看起来您正在进行一些奇特的切片，以便将特定值插入到输出数组的特定位置。这对于dask.array来说可能会很棘手。相反，我建议创建一个生成numpy数组的函数

def myfunc(chunk):
    values = pdist(chunk[0, :, :])
    output = np.zeroes((2048, 2048))
    r = np.arange(2048)
    output[r[:, None] < r] = values
    return output

dcoor.map_blocks(myfunc)

---

pdist

的输出是切片中每对坐标之间的距离数组。在我的示例中，每个切片的结果都有维度

（2048*2047/2，）

。这些数组实际上已经足够了。我认为这个答案提供了一个良好的开端，尽管我认为

map_blocks

正在尝试执行

pdist

，就好像我通过了整个3D阵列一样。我只想在2D切片上进行计算，我将其制成块；不应该有跨块计算。它对每个

（12048,3）

大小的块调用pdist。我会把你的pdist函数包装成一个东西，去掉第一个维度，然后返回完整的平方矩阵。在上面编辑了我的答案。

def myfunc(chunk):
    values = pdist(chunk[0, :, :])
    output = np.zeroes((2048, 2048))
    r = np.arange(2048)
    output[r[:, None] < r] = values
    return output

dcoor.map_blocks(myfunc)

from dask import delayed, compute
coor2 = delayed(coor)
slices = [coor2[i] for i in range(coor.shape[0])]
slices2 = [delayed(pdist)(slice) for slice in slices]
results = compute(*slices2)