Get MSD (mean square displacement) from MD trajectory

The first script convert XDATCAR to XYZ file The second script extract MSD info from XYZ file

import numpy as np

franum=0
count=0
latt=[]
axyz=[]
ele=[]
etnum=[]

#fixflag=input(" Seletive Dynamics is T or F: ")
fixflag=False

tlnum=0
pf=open("XDATCAR","r")

fileend=0

tmpstr=pf.readline()
count=count+1
tmpstr=pf.readline()
count=count+1

#### Begin Read  Lattice Vector  #####
latt.append([])
for j in range(0,3):
      tl=pf.readline()
      count=count+1
      tlst=tl.split()
      tlf=[float(x) for x in tlst[0:3]]
      latt[franum].append(tlf)
print latt

#### Begin Read Element Num  #####
tmpstr=pf.readline()
count=count+1
[ele.append(x) for x in tmpstr.split()]
print ele

tnumstr=pf.readline()
print "read ele num"
print tnumstr
count=count+1
etnum=[int(x) for x in tnumstr.split()]
print etnum
totnum=sum(etnum)

if fixflag:
      tmpstr=pf.readline()
      count=count+1

#### Begin Read Atom Fraction Coordination  #####
while not fileend:
      tmpstr=pf.readline()
      if tmpstr == '':
            print "file read done"
            fileend=1
            break
      count=count+1
      axyz.append([])
      for j in range(0,totnum):
            txyz=pf.readline()
            if txyz == '':
                  print "Error:Missing Atom In Frame %d"%franum
                  fileend=1
                  break
            count=count+1
            taxyz=txyz.split()
            xyz=[float(x) for x in taxyz[0:3]]
            axyz[franum].append(xyz)
      franum=franum+1


LATTICE=np.array(latt)
ATXYZ=np.array(axyz)

print ATXYZ[0]

CXYZ=[]

outfile="LiXY"
of=open(outfile,"w")
for i in range(0,franum,10):
# ATXYZ=np.array(axyz[i])
# LATT=LATTICE[i]
# tCXYZ=np.dot(ATXYZ,LATT)
# CXYZ.append(tCXYZ)
# print >> of," %d  "%(enum[i][2])
# print >> of," %s  %d  %s  "%("### ",i+1," ####")
        for j in range(0,etnum[0]):
            print >> of,"     %f   %f  "%(ATXYZ[i][j][0],ATXYZ[i][j][1])
of.close()
# The MIT License (MIT)
#
# Copyright (c) 2014 Muratahan Aykol
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE

import numpy as np
from copy import deepcopy

# This function reads an XYZ file and a list of lattice vectors L = [x,y,z] and gives MSD + unwrapped coordinates

def MSD(xyz_file,L):
    a = []; l = [];
    a.append(L[0]); a.append(L[1]); a.append(L[2]); #basis vectors in cartesian coords
    l.append(np.sqrt(np.dot(a[0],a[0]))); l.append(np.sqrt(np.dot(a[1],a[1]))); l.append(np.sqrt(np.dot(a[2],a[2]))); #basis vector lengths

    file = open(xyz_file, 'r')
    recorder = open("msd.out", 'w')
    coord_rec = open("unwrapped.xyz", 'w')

    origin_list = [] # Stores the origin as [element,[coords]]
    prev_list = [] # Stores the wrapped previous step
    unwrapped_list = [] # Stores the instantenous unwrapped

    msd = [] #Stores atom-wise MSD  Stores msd as [msd]
    msd_dict ={} #Stores element-wise MSD
    msd_lattice = []
    msd_dict_lattice ={}

    element_list = [] # element list
    element_dict = {} # number of elements stored

    content = file.readline()
    N = int(content)

    for i in range(N):
        msd.append(np.float64('0.0'))
        msd_lattice.append([0.0, 0.0, 0.0 ])

    file.readline()
    step = 0

    while True:
        step += 1
        # Get and store the origin coordinates in origin_dict at first step
        if step == 1:
            for i in xrange(N):
                t = file.readline().rstrip('\n').split()
                element = t[0]
                if element not in element_list:
                    element_list.append(element)
                if element not in element_dict:
                    element_dict[element] = 1.0
                else:
                    element_dict[element] += 1.0
                coords = np.array( [ float(s) for s in t[1:] ] )
                origin_list.append([element,coords])
            # Copy the first set of coordinates as prev_dict and unwrapped
            unwrapped_list = deepcopy(origin_list)
            prev_list = deepcopy(origin_list)
            recorder.write("step ")
            for element in element_list:
                recorder.write(element+" ")
            recorder.write("\n")

        # Read wrapped coordinates into wrapped_dict
        content = file.readline()
        if len(content) == 0:
            print "\n---End of file---\n"
            break
        N = int(content)
        file.readline()
        wrapped_list = [] # Erease the previous set of coordinates
        for i in xrange(N):
            t = file.readline().rstrip('\n').split()
            element = t[0]
            coords = np.array( [ float(s) for s in t[1:] ] )
            wrapped_list.append([element,coords])

        coord_rec.write(str(N)+ "\ncomment\n")

        # Unwrap coodinates and get MSD

        for atom in range(N):

            msd[atom] = 0.0

            coord_rec.write(wrapped_list[atom][0])

            # decompose wrapped atom coordinates to onto lattice vectors:
            w1 = wrapped_list[atom][1][0]
            w2 = wrapped_list[atom][1][1]
            w3 = wrapped_list[atom][1][2]

            # decompose prev atom coordinates to onto lattice vectors:
            p1 = prev_list[atom][1][0]
            p2 = prev_list[atom][1][1]
            p3 = prev_list[atom][1][2]

            #get distance between periodic images and use the smallest one
            if np.fabs(w1 - p1) > 0.5:
                 u1 = w1 - p1 - np.sign(w1 - p1)
            else:
                 u1 = w1 - p1

            if np.fabs(w2 - p2) > 0.5:
                 u2 = w2 - p2 - np.sign(w2 - p2)
            else:
                 u2 = w2 - p2

            if np.fabs(w3 - p3) > 0.5:
                 u3 = w3 - p3 - np.sign(w3 - p3)
            else:
                 u3 = w3 - p3

            #add unwrapped displacements to unwrapped coords

            unwrapped_list[atom][1][0] += u1
            unwrapped_list[atom][1][1] += u2
            unwrapped_list[atom][1][2] += u3

            uw = unwrapped_list[atom][1][0]*a[0] + unwrapped_list[atom][1][1]*a[1] +unwrapped_list[atom][1][2]*a[2]
            ol = origin_list[atom][1][0]*a[0] + origin_list[atom][1][1]*a[1] + origin_list[atom][1][2]*a[2]

            msd[atom] = np.linalg.norm(uw-ol)**2
            msd_lattice[atom] = [np.linalg.norm(uw[0]-ol[0])**2,np.linalg.norm(uw[1]-ol[1])**2,np.linalg.norm(uw[2]-ol[2])**2]

            coord_rec.write(" " + np.array_str(uw).replace("[","").replace("]",""))
            coord_rec.write("\n")

        prev_list = [] # Store current wrapped coordinates for the next step
        prev_list = deepcopy(wrapped_list)
        # record msd
        recorder.write(str(step) + " ")

        for el in element_list:
             msd_dict[el] = 0.0
             msd_dict_lattice[el]=[0.,0.,0.]

        for atom in range(len(msd)):
            msd_dict[wrapped_list[atom][0]] += msd[atom]/element_dict[wrapped_list[atom][0]]
            for i in range(3):
                msd_dict_lattice[wrapped_list[atom][0]][i] += msd_lattice[atom][i]/element_dict[wrapped_list[atom][0]]

        for el in element_list:
            recorder.write(str(msd_dict[el])+ " " + str(msd_dict_lattice[el][0])+ " " + str(msd_dict_lattice[el][1])+ " " + str(msd_dict_lattice[el][2])+ " ")

        recorder.write("\n")
        if step % 10 == 0:
            print step
    recorder.close()
    file.close()
    coord_rec.close()

def read_lat_vec():
    lat_file = open('lattice.vectors','r')
    line = []
    for i in xrange(3):
        line.append([float(x) for x in lat_file.readline().rstrip('\n').split()])
        print line[i]
    lattice = np.array([line[0],line[1],line[2]])
    return lattice


#You can read the lattice vectors from lattice.vector file
lattice = read_lat_vec()

#Or define the lattice vector manually as in
#lattice =np.array([[-12.181156,-4.306689,7.459404],[0.000000,-12.920067,7.459404],[0.000000,0.000000,14.918808]])

#Run the MSD calculator with XDATCAR_fract.xyz and the lattice vector defined above
MSD("XDATCAR_fract.xyz",lattice)