lammps均方根位移MSD计算---Python实现

import numpy as np
import matplotlib.pyplot as plt

dump_step=input("请输入的dump步数")      
time_step=input("请输入模拟运算步数")
LIMIT=input("请输入体系原子总数")
dump_step = float(dump_step)
time_step = float(time_step)
LIMIT=float(LIMIT)

def constrct_array(input,ii,c): #input - input array , ii - ii_th timestep , c - number of atoms
    temp=input[int(ii)*int(c):(int(ii)+1)*int(c)] #input 输出的值为字符串格式
    return temp   

def distant(l1,l2): #l1,l2 point position
    a=np.sqrt((float(l1[2])-float(l2[2]))**2+(float(l1[3])-float(l2[3]))**2+(float(l1[4])-float(l2[4]))**2)
    return a  

def MSD(list,num):
    msd3=[]
    msd2=0
    for i in range(len(list)): #number of time-point
        a=0
        for j in range(len(list[i])-1):
            #print(list[i][j])
            msd1=(abs(distant(list[i+1][j],base1_point)-distant(list[i][j],base1_point)))**2
            #print(msd1)
            a= a + msd1
        msd2=a/int(len(list[i]))
        print(msd2)
        msd3.append(msd2)
    return msd3

            
init_data=np.loadtxt("dump2.txt")
file1=[]
#print(init_data)
#with open('dump.txt') as file_project:
 #   lines =file_project.readlines()
  #  row = []
   # for line in lines:
    #    row = line.split(' ')
     #   file1.append(row)
#print(file1)
len_time = int(time_step/dump_step)
#print(len_time)
base1_point=(0,0,0,0,0,0,0)
O=[]
Al=[]
H=[]
for i in range(len_time):
    o=[]
    al=[]
    h=[]
    #print(i)
    tt=constrct_array(init_data,int(i),LIMIT) #将文件按照时间步截取
    print(tt)
    b=tt.tolist()   # tolist函数：将目标转化为列表
    #del b[0:9]
    #print(b)
    for each in b:
        #print(each)
        if distant(each,base1_point)<=31 and int(each[1])==1:
            o.append(each)
        elif distant(each,base1_point)<=31 and int(each[1])==2:
            al.append(each)
        elif distant(each,base1_point)<=31 and int(each[1])==3:
            h.append(each)
    O.append(o)
    Al.append(al)
    H.append(h)
print(O)
#print(Al)
#print(H)
y_O=MSD(O,LIMIT)
y_Al=MSD(Al,LIMIT)
y_H=MSD(H,LIMIT)

#print(y_O)
#print(y_Al)
#print(y_H)
'''
a=int(time_step/dump_step)
x=np.linspace(0,a-1,1)
plt.plot(x,y_O,c='red',label='O')
plt.plot(x,y_Al,c='black',label='Al')
plt.plot(x,y_H,c='green',label='H')

plt.show()    
    
'''