Python并行处理数据

Zcore={ 'Ag':   19.00, 'Al':    3.00, 'Ar':    8.00, 'As':   15.00, 'Au':   19.00, 'Ba':   10.00, 'Be':    4.00, 
    'Bi':   15.00, 'B':    3.00, 'Br':    7.00, 'Ca':   10.00, 'Cd':   20.00, 'Cl':    7.00, 'C':    4.00, 
    'Co':   17.00, 'Cr':   14.00, 'Cs':    9.00, 'Cu':   19.00, 'Fe':   16.00, 'F':    7.00, 'Ga':   13.00, 
    'Ge':   14.00, 'He':    2.00, 'Hf':   12.00, 'Hg':   20.00, 'H':    1.00, 'In':   13.00, 'I':    7.00, 
    'Ir':   17.00, 'K':    9.00, 'Kr':    8.00, 'Li':    3.00, 'Mg':   10.00, 'Mn':   15.00, 'Mo':   14.00, 
    'Na':    9.00, 'Nb':   13.00, 'Ne':    8.00, 'Ni':   18.00, 'N':    5.00, 'O':    6.00, 'Os':   16.00, 
    'Pb':   14.00, 'Pd':   18.00, 'P':    5.00, 'Po':   16.00, 'Pt':   16.00, 'Rb':    9.00, 'Re':   15.00, 
    'Rh':   17.00, 'Rn':   18.00, 'Ru':   16.00, 'Sb':   15.00, 'Sc':   11.00, 'Se':   16.00, 'Si':    4.00, 
    'Sn':   14.00, 'S':    6.00, 'Sr':   10.00, 'Ta':   13.00, 'Tc':   15.00, 'Te':   16.00, 'Ti':   12.00, 
    'Tl':   13.00, 'V':   13.00, 'W':   14.00, 'Xe':    8.00, 'Y':   11.00, 'Zn':   20.00, 'Zr':   12.00
}
Zcore["WF"]=-2.0    #电子,一个wannier上占两个电子

def cal_dipole(ntyp,nat,xyz):
    Dipole=np.zeros([xyz.shape[0],3])
    for i in np.arange(nat.size):
        Dipole=Dipole+xyz[:,i,0:nat[i],:].sum(axis=1)*Zcore[ntyp[i]]
    return Dipole

#----------------
if 0: # 串行计算
    print("Calculate Dipole by one cpu")
    time_start=time.time()
    Dipole=cal_dipole(ntyp,nat,xyz)
    time_end=time.time()
    print('Time cost',time_end-time_start,'s')

else: #并行计算
    #multi rdf
    import  multiprocessing
    time_start=time.time()
    m_process=int(max(multiprocessing.cpu_count()/2.0,8))#选择合适的cpu数
    if m_process > nstep:
        m_process = nstep
    print("m_Calculate Dipole by "+str(m_process)+ " cpus")
    m_length=int(np.floor(nstep/m_process)) 
    m_left=nstep-m_process*m_length
    m_snstep=np.zeros(m_process).astype(np.int)
    m_enstep=np.zeros(m_process).astype(np.int)
    m_snstep[0]=0
    m_enstep[0]=m_snstep[0]+m_length
    for i in np.arange(1,m_process):
        m_snstep[i]=min(m_enstep[i-1],nstep)
        if i <= m_left: #向下取整，把剩下的以此分配到从1开始编号的CPU
            m_enstep[i]=min(m_snstep[i]+m_length+1,nstep)
        else:
            m_enstep[i]=min(m_snstep[i]+m_length,nstep)
    print(m_snstep)
    def multi_dipole(i):
        start=m_snstep[i]
        end=m_enstep[i]
        Dipole=cal_dipole(ntyp,nat,xyz[start:end])
        return Dipole

    print(m_enstep)
    m_cpu = [i for i in range(0, m_process)]
    if __name__ == '__main__':
        p = multiprocessing.Pool(m_process)
        out = p.map_async(multi_dipole,m_cpu).get()
        p.close()
        p.join()
    Dipole=np.zeros([nstep,3])
    for i in np.arange(m_process):
        start=m_snstep[i]
        end=m_enstep[i]
        Dipole[start:end]=out[i]
    time_end=time.time()
    print('Time cost',time_end-time_start,'s')