Python：如何用字典并行化循环

Question

已编辑：我的代码如下所示：

__author__ = 'feynman'

cimport cython

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.nonecheck(False)
def MC_Surface(volume, mc_vol):

    Perm_area = {
        "00000000": 0.000000,
        "11111111": 0.000000,
         ...
         ...
        "11100010": 1.515500,
        "00011101": 1.515500
    }

    cdef int j, i, k
    for k in range(volume.shape[2] - 1):
        for j in range(volume.shape[1] - 1):
            for i in range(volume.shape[0] - 1):

                pattern = '%i%i%i%i%i%i%i%i' % (
                    volume[i, j, k],
                    volume[i, j + 1, k],
                    volume[i + 1, j, k],
                    volume[i + 1, j + 1, k],
                    volume[i, j, k + 1],
                    volume[i, j + 1, k + 1],
                    volume[i + 1, j, k + 1],
                    volume[i + 1, j + 1, k + 1])

                mc_vol[i, j, k] = Perm_area[pattern]

    return mc_vol

为了提速，修改为：

    {
      ...
      ...
      "11100010": 1.515500,
      "00011101": 1.515500
    }
    keys = np.array(Perm_area.keys())
    values = np.array(Perm_area.values())

    starttime = time.time()
    tmp_vol = GetPattern(volume)
    print 'time to populate the key array: ', time.time() - starttime

    cdef int i
    starttime=time.time()
    for i, this_key in enumerate(keys):
        mc_vol[tmp_vol == this_key] = values[i]

    print 'time for the loop: ', time.time() -starttime
    return mc_vol

def GetPattern(volume):
    a = (volume.astype(np.int)).astype(np.str)
    output = a.copy()  # Central voxel
    output[:, :-1, :] = np.char.add(output[:, :-1, :], a[:, 1:, :])  # East
    output[:-1, :, :] = np.char.add(output[:-1, :, :], a[1:, :, :])  # South
    output[:-1, :-1, :] = np.char.add(output[:-1, :-1, :], a[1:, 1:, :])  # SouthEast
    output[:, :, :-1] = np.char.add(output[:, :, :-1], a[:, :, 1:])  # Down
    output[:, :-1, :-1] = np.char.add(output[:, :-1, :-1], a[:, 1:, 1:])  # DownEast
    output[:-1, :, :-1] = np.char.add(output[:-1, :, :-1], a[1:, :, 1:])  # DownSouth
    output[:-1, :-1, :-1] = np.char.add(output[:-1, :-1, :-1], a[1:, 1:, 1:])  # DownSouthEast
    output = output[:-1, :-1, :-1]
    del a
    return output

对于大小为 500^3 的 3D 数组，这需要更长的时间。这里是 tmp_vol 3D 字符串数组。例如：如果说 tmp_vol[0,0,0] = "00000000" 那么 mc_vol[0,0,0] = 0.00000。或者，我可以去掉 mc_vol 并写 if tmp_vol[0,0,0] = "00000000" then tmp_vol[0,0,0] = 0.00000。

这里，for循环需要很多时间，我看到只使用了一个CPU。我尝试使用 map 和 lambda 并行映射它们，但遇到了错误。我对 python 很陌生，所以任何提示都会很棒。

Answer 1

由于我不太了解您的代码，并且您说“任何提示都会很棒”，因此我会给您一些一般性建议。基本上你想加快一个 for 循环

for i, this_key in enumerate(keys):

您可以将keys 数组拆分为几个部分，如下所示：

length = len(keys)
part1 = keys[:length/3]
part2 = keys[length/3: 2*length/3]
part3 = keys[2*length/3:]

然后在一个子流程中处理各个部分：

from concurrent.futures import ProcessPoolExecutor

def do_work(keys):
    for i, this_key in enumerate(keys): 
        mc_vol[tmp_vol == this_key] = values[i]

with ProcessPoolExecutor(max_workers=3) as e:
    e.submit(do_work, part1)
    e.submit(do_work, part2)
    e.submit(do_work, part3)

return mc_vol

就是这样。

Answer 2

首先，虽然数组相等性检查是 O(N)，但字典查找需要非常恒定的时间。因此，您应该遍历您的数组而不是您的字典。

其次，nested list comprehension 可以节省很多时间（将 python 循环更改为 C 循环）。

mc_vol = [[Perm_area[key] for key in row] for row in tmp_vol]

这为您提供了一个列表列表，因此您可以在这种情况下完全避免使用 numpy。虽然如果你需要一个 numpy 数组，只需转换：

mc_vol = np.array(mc_vol)