如何将 mpi4py 应用于以下循环以加快速度？

Question

过去我不需要使用并行计算，因为我的 python 脚本通常不是计算密集型的。最近，我写了以下for循环，发现运行大约需要10-15分钟：

cpx_var = np.linspace(0.20,0.80,301)
horn_var = np.linspace(0.20,0.80,301)
plag_var = np.linspace(0.05,0.1,26)
mag_var = np.linspace(0.02,0.06,21)
ap_var = np.linspace(0.002,0.006,3)

poss_comb = []
count=0

for i in cpx_var:
    for j in horn_var:
        for k in plag_var:
            for l in mag_var:
                for m in ap_var:
                    count = count+1
                    if abs((i+j+k+l+m)-1.0)<0.002:
                        poss_comb.append([i,j,k,l,m])
print(count)

我一直在研究 mpi4py 作为加快我的 for 循环的一种方法，但我不确定如何将它应用到我的示例中。有人有什么建议吗？

谢谢！

扎克·埃里克森

Answer 1

虽然多处理是加快代码速度的一种方法，但您首先应该尝试的是让您的代码更高效。您正在使用 5 个嵌套循环，这在 Python 中太疯狂了！在 Python 中，您应该尽可能地使用数组而不是 for 循环来实现代码。

这是一个解决方案，我在最后只使用一个 for 循环（它实际上是一个列表推导，它是 for 循环的稍快版本）。我已经使用了 np.mesh 函数，然后进行了一些切片来摆脱您拥有的所有循环，而只使用列表理解将所有结果放入您想要的列表格式中。

所以试试这个版本吧。在我的笔记本电脑上运行只需大约 20 秒

import numpy as np

cpx_var = np.linspace(0.20,0.80,301)
horn_var = np.linspace(0.20,0.80,301)
plag_var = np.linspace(0.05,0.1,26)
mag_var = np.linspace(0.02,0.06,21)
ap_var = np.linspace(0.002,0.006,3)

cpx_vars, horn_vars, plag_vars, mag_vars, ap_vars = np.meshgrid(cpx_var, horn_var, plag_var, mag_var, ap_var)

selection = abs((cpx_vars + horn_vars + plag_vars + mag_vars + ap_vars)-1.0)<0.002

poss_comb = [list(comb) for comb in zip(cpx_vars[selection], horn_vars[selection], plag_vars[selection], mag_vars[selection], ap_vars[selection])]

Answer 2

我找到了另一种方法可以将您的代码再加快 23 倍。自从我回答了这个问题以来，我一直在了解有关名为 numba 的 Python 包的更多信息。这个包非常聪明，它本质上将你的 Python 函数转换为机器代码，并且可以为你提供巨大的加速（取决于你在做什么）。

我正在玩这个包，我在这里想到了你的问题，我想看看与使用 numpy 数组相比，numba 能做到多快。在您的代码、我在上面的答案中粘贴的 numpy 版本以及您的代码的 numba 优化版本之间进行比较。我发现 numpy（无循环）版本比你的循环版本快 23 倍，但 numba 版本快 455 倍！

不管怎样，我想你可能会喜欢看这个。 numba 版本比我上面的答案快大约 23 倍。以下是我使用的代码的 3 个版本。我把你的代码变成了一个名为original_loop_function的函数，我在上面发布的numpy版本到函数numpy_function中，更快​​的numba版本是函数numba_loop_function。希望这对您仍然有用

import numpy as np
from numba import jit

cpx_var = np.linspace(0.20,0.80,301)
horn_var = np.linspace(0.20,0.80,301)
plag_var = np.linspace(0.05,0.1,26)
mag_var = np.linspace(0.02,0.06,21)
ap_var = np.linspace(0.002,0.006,3)


def original_loop_function(cpx_var, horn_var, plag_var, mag_var, ap_var):
    '''Your original version of this code'''
    poss_comb = []
    count=0

    for i in cpx_var:
        for j in horn_var:
            for k in plag_var:
                for l in mag_var:
                    for m in ap_var:
                        count = count+1
                        if abs((i+j+k+l+m)-1.0)<0.002:
                            poss_comb.append([i,j,k,l,m])
    return poss_comb

def numpy_function(cpx_var, horn_var, plag_var, mag_var, ap_var):
    '''Numpy version of your code with no loops'''
    cpx_vars, horn_vars, plag_vars, mag_vars, ap_vars = np.meshgrid(cpx_var, horn_var, plag_var, mag_var, ap_var)

    selection = abs((cpx_vars + horn_vars + plag_vars + mag_vars + ap_vars)-1.0)<0.002

    poss_comb = [list(comb) for comb in zip(cpx_vars[selection], horn_vars[selection], plag_vars[selection], mag_vars[selection], ap_vars[selection])]

    return poss_comb


@jit(nopython=True)
def numba_loop_function(cpx_var, horn_var, plag_var, mag_var, ap_var):
    '''Numba optimised version of your code'''
    poss_comb = []
    count=0

    for i in cpx_var:
        for j in horn_var:
            for k in plag_var:
                for l in mag_var:
                    for m in ap_var:
                        count = count+1
                        if abs((i+j+k+l+m)-1.0)<0.002:
                            poss_comb.append([i,j,k,l,m])
    return poss_comb