如何将参数传递给多处理循环

时间:2019-06-18 21:37:13

标签: python multithreading parallel-processing

我想进行Hartree-Fock类的计算。简而言之,这是一个自洽矩阵收敛问题。我从矩阵H(0)开始,从中构造Fock(0)项。现在,我将H(1)= H(0)+ Fock(0)定义为新矩阵。然后我从H(1)构造Fock(1)并定义H(2)= H(0)+ Fock(1)。我重复此过程,直到H(n + 1)足够接近H(n)。

计算非常耗时,因此我正在使用多处理。但是我发现我的代码并没有像我希望的那样改变Fock术语。它始终使用Fock(0),因此结果没有任何意义。

以下是多处理部分上的一段简化代码

from numpy import *
import multiprocessing as mp
import itertools

kt = 0.3120
q = array([[0,-1],[sqrt(3)/2,1./2],[-sqrt(3)/2,1./2]])*kt

def hbz(x):
    "A set of k points"
    tL = []
    for i in range(-x,x+1):
        for j in range(-x,x+1):
            k = i*q[1]/x+j*q[2]/x
            if linalg.norm(k) <= kt+10**-5 and -q[1][0]-1e-10<=k[0]<=q[1][0]+1e-10:
                tL.append(k)
    return tL

step = 2
d = {}
def H(k):
    return 4*k
def fock(k):
    return -Hp(k)/3
Hp = lambda k: H(k)

def fill_d(k):
    return str(k), 5*fock(k)
if __name__ == '__main__':
    pool = mp.Pool(min(8,mp.cpu_count()-2))
    for ite in range(3):
        d = dict(pool.map_async(fill_d, [kpt for kpt in hbz(step)]).get())
        Hp = lambda k: H(k)+d[str(k)]
        print (d)
        print ("============================")

输出为

{'[ 0.    -0.312]': array([-0.  ,  2.08]), '[-0.13509996 -0.234     ]': array([0.90066642, 1.56      ]), '[-0.27019993 -0.156     ]': array([1.80133284, 1.04      ]), '[ 0.13509996 -0.234     ]': array([-0.90066642,  1.56      ]), '[ 0.    -0.156]': array([-0.  ,  1.04]), '[-0.13509996 -0.078     ]': array([0.90066642, 0.52      ]), '[-0.27019993  0.        ]': array([ 1.80133284, -0.        ]), '[ 0.27019993 -0.156     ]': array([-1.80133284,  1.04      ]), '[ 0.13509996 -0.078     ]': array([-0.90066642,  0.52      ]), '[0. 0.]': array([-0., -0.]), '[-0.13509996  0.078     ]': array([ 0.90066642, -0.52      ]), '[-0.27019993  0.156     ]': array([ 1.80133284, -1.04      ]), '[0.27019993 0.        ]': array([-1.80133284, -0.        ]), '[0.13509996 0.078     ]': array([-0.90066642, -0.52      ]), '[0.    0.156]': array([-0.  , -1.04]), '[-0.13509996  0.234     ]': array([ 0.90066642, -1.56      ]), '[0.27019993 0.156     ]': array([-1.80133284, -1.04      ]), '[0.13509996 0.234     ]': array([-0.90066642, -1.56      ]), '[0.    0.312]': array([-0.  , -2.08])}
============================
{'[ 0.    -0.312]': array([-0.  ,  2.08]), '[-0.13509996 -0.234     ]': array([0.90066642, 1.56      ]), '[-0.27019993 -0.156     ]': array([1.80133284, 1.04      ]), '[ 0.13509996 -0.234     ]': array([-0.90066642,  1.56      ]), '[ 0.    -0.156]': array([-0.  ,  1.04]), '[-0.13509996 -0.078     ]': array([0.90066642, 0.52      ]), '[-0.27019993  0.        ]': array([ 1.80133284, -0.        ]), '[ 0.27019993 -0.156     ]': array([-1.80133284,  1.04      ]), '[ 0.13509996 -0.078     ]': array([-0.90066642,  0.52      ]), '[0. 0.]': array([-0., -0.]), '[-0.13509996  0.078     ]': array([ 0.90066642, -0.52      ]), '[-0.27019993  0.156     ]': array([ 1.80133284, -1.04      ]), '[0.27019993 0.        ]': array([-1.80133284, -0.        ]), '[0.13509996 0.078     ]': array([-0.90066642, -0.52      ]), '[0.    0.156]': array([-0.  , -1.04]), '[-0.13509996  0.234     ]': array([ 0.90066642, -1.56      ]), '[0.27019993 0.156     ]': array([-1.80133284, -1.04      ]), '[0.13509996 0.234     ]': array([-0.90066642, -1.56      ]), '[0.    0.312]': array([-0.  , -2.08])}
============================
{'[ 0.    -0.312]': array([-0.  ,  2.08]), '[-0.13509996 -0.234     ]': array([0.90066642, 1.56      ]), '[-0.27019993 -0.156     ]': array([1.80133284, 1.04      ]), '[ 0.13509996 -0.234     ]': array([-0.90066642,  1.56      ]), '[ 0.    -0.156]': array([-0.  ,  1.04]), '[-0.13509996 -0.078     ]': array([0.90066642, 0.52      ]), '[-0.27019993  0.        ]': array([ 1.80133284, -0.        ]), '[ 0.27019993 -0.156     ]': array([-1.80133284,  1.04      ]), '[ 0.13509996 -0.078     ]': array([-0.90066642,  0.52      ]), '[0. 0.]': array([-0., -0.]), '[-0.13509996  0.078     ]': array([ 0.90066642, -0.52      ]), '[-0.27019993  0.156     ]': array([ 1.80133284, -1.04      ]), '[0.27019993 0.        ]': array([-1.80133284, -0.        ]), '[0.13509996 0.078     ]': array([-0.90066642, -0.52      ]), '[0.    0.156]': array([-0.  , -1.04]), '[-0.13509996  0.234     ]': array([ 0.90066642, -1.56      ]), '[0.27019993 0.156     ]': array([-1.80133284, -1.04      ]), '[0.13509996 0.234     ]': array([-0.90066642, -1.56      ]), '[0.    0.312]': array([-0.  , -2.08])}

可以看到并行化函数fill_d在不同的迭代中没有变化。由于此示例非常简单,因此无需多处理即可完成此操作:

for ite in range(3):
    for kpt in hbz(step):
        d[str(kpt)] = 5*fock(kpt)
    Hp = lambda k: H(k)+d[str(k)]
    print (d)
    print ("============================")

所以我知道预期的输出是

{'[ 0.    -0.312]': array([-0.  ,  2.08]), '[-0.13509996 -0.234     ]': array([0.90066642, 1.56      ]), '[-0.27019993 -0.156     ]': array([1.80133284, 1.04      ]), '[ 0.13509996 -0.234     ]': array([-0.90066642,  1.56      ]), '[ 0.    -0.156]': array([-0.  ,  1.04]), '[-0.13509996 -0.078     ]': array([0.90066642, 0.52      ]), '[-0.27019993  0.        ]': array([ 1.80133284, -0.        ]), '[ 0.27019993 -0.156     ]': array([-1.80133284,  1.04      ]), '[ 0.13509996 -0.078     ]': array([-0.90066642,  0.52      ]), '[0. 0.]': array([-0., -0.]), '[-0.13509996  0.078     ]': array([ 0.90066642, -0.52      ]), '[-0.27019993  0.156     ]': array([ 1.80133284, -1.04      ]), '[0.27019993 0.        ]': array([-1.80133284, -0.        ]), '[0.13509996 0.078     ]': array([-0.90066642, -0.52      ]), '[0.    0.156]': array([-0.  , -1.04]), '[-0.13509996  0.234     ]': array([ 0.90066642, -1.56      ]), '[0.27019993 0.156     ]': array([-1.80133284, -1.04      ]), '[0.13509996 0.234     ]': array([-0.90066642, -1.56      ]), '[0.    0.312]': array([-0.  , -2.08])}
============================
{'[ 0.    -0.312]': array([-0.        , -1.38666667]), '[-0.13509996 -0.234     ]': array([-0.60044428, -1.04      ]), '[-0.27019993 -0.156     ]': array([-1.20088856, -0.69333333]), '[ 0.13509996 -0.234     ]': array([ 0.60044428, -1.04      ]), '[ 0.    -0.156]': array([-0.        , -0.69333333]), '[-0.13509996 -0.078     ]': array([-0.60044428, -0.34666667]), '[-0.27019993  0.        ]': array([-1.20088856, -0.        ]), '[ 0.27019993 -0.156     ]': array([ 1.20088856, -0.69333333]), '[ 0.13509996 -0.078     ]': array([ 0.60044428, -0.34666667]), '[0. 0.]': array([-0., -0.]), '[-0.13509996  0.078     ]': array([-0.60044428,  0.34666667]), '[-0.27019993  0.156     ]': array([-1.20088856,  0.69333333]), '[0.27019993 0.        ]': array([ 1.20088856, -0.        ]), '[0.13509996 0.078     ]': array([0.60044428, 0.34666667]), '[0.    0.156]': array([-0.        ,  0.69333333]), '[-0.13509996  0.234     ]': array([-0.60044428,  1.04      ]), '[0.27019993 0.156     ]': array([1.20088856, 0.69333333]), '[0.13509996 0.234     ]': array([0.60044428, 1.04      ]), '[0.    0.312]': array([-0.        ,  1.38666667])}
============================
{'[ 0.    -0.312]': array([-0.        ,  4.39111111]), '[-0.13509996 -0.234     ]': array([1.90140689, 3.29333333]), '[-0.27019993 -0.156     ]': array([3.80281377, 2.19555556]), '[ 0.13509996 -0.234     ]': array([-1.90140689,  3.29333333]), '[ 0.    -0.156]': array([-0.        ,  2.19555556]), '[-0.13509996 -0.078     ]': array([1.90140689, 1.09777778]), '[-0.27019993  0.        ]': array([ 3.80281377, -0.        ]), '[ 0.27019993 -0.156     ]': array([-3.80281377,  2.19555556]), '[ 0.13509996 -0.078     ]': array([-1.90140689,  1.09777778]), '[0. 0.]': array([-0., -0.]), '[-0.13509996  0.078     ]': array([ 1.90140689, -1.09777778]), '[-0.27019993  0.156     ]': array([ 3.80281377, -2.19555556]), '[0.27019993 0.        ]': array([-3.80281377, -0.        ]), '[0.13509996 0.078     ]': array([-1.90140689, -1.09777778]), '[0.    0.156]': array([-0.        , -2.19555556]), '[-0.13509996  0.234     ]': array([ 1.90140689, -3.29333333]), '[0.27019993 0.156     ]': array([-3.80281377, -2.19555556]), '[0.13509996 0.234     ]': array([-1.90140689, -3.29333333]), '[0.    0.312]': array([-0.        , -4.39111111])}

如何在迭代时根据H矩阵修改多处理代码以计算Fock部分?

1 个答案:

答案 0 :(得分:0)

所以我自己找到了一条路。对于那些可能遇到相同问题的人,我将解决方案发布在这里。基本上, Pool 中有一个名为 starmap 的类似函数,可以容纳更多参数。而且我稍微调整了顺序。

def fock(k, n):
    return -Hp(k,n)/3
def Hp(k,n):
    if n==0: return H(k)
    else: return H(k)+d[n-1][str(k)] 
d={}
def fill_d(k, n):
    return str(k), 5*fock(k,n)
for ite in range(3):
    if __name__ == '__main__':
        pool = mp.Pool(min(8,mp.cpu_count()-2))    
        d[ite] = dict(pool.starmap_async(fill_d, [(kpt,ite) for kpt in hbz(step)]).get())
        print (d[ite])
        print ("============================")
        pool.close()