4
Python并行处理数据
source link: https://cndaqiang.github.io/2020/12/23/py-mul/
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.
Numpy
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')
Recommend
About Joyk
Aggregate valuable and interesting links.
Joyk means Joy of geeK