Get MSD (mean square displacement) from MD trajectory (Python)

The first script convert XDATCAR to XYZ file The second script extract MSD info from XYZ file Final result is in the file msd.out. The format is: msd_of_element1,msd1_x,msd1_y,msd1_z,msd_of_element2,msd2_x,msd2_y,msd2_z,…

# 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

xdatcar = open('XDATCAR', 'r')
xyz = open('XDATCAR.xyz', 'w')
xyz_fract = open('XDATCAR_fract.xyz', 'w')


system = xdatcar.readline()
scale = float(xdatcar.readline().rstrip('\n'))
print(scale)

#get lattice vectors
a1 = np.array([ float(s)*scale for s in xdatcar.readline().rstrip('\n').split() ])
a2 = np.array([ float(s)*scale for s in xdatcar.readline().rstrip('\n').split() ])
a3 = np.array([ float(s)*scale for s in xdatcar.readline().rstrip('\n').split() ])

print(a1)
print(a2)
print(a3)

#Save scaled lattice vectors
lat_rec = open('lattice.vectors', 'w')
lat_rec.write(str(a1[0])+' '+str(a1[1])+' '+str(a1[2])+'\n')
lat_rec.write(str(a2[0])+' '+str(a2[1])+' '+str(a2[2])+'\n')
lat_rec.write(str(a3[0])+' '+str(a3[1])+' '+str(a3[2]))
lat_rec.close()


#Read xdatcar
element_names = xdatcar.readline().rstrip('\n').split()

element_dict = {}
element_numbers = xdatcar.readline().rstrip('\n').split()

i = 0
N = 0
for t in range(len(element_names)):
    element_dict[element_names[t]] = int(element_numbers[i])
    N += int(element_numbers[i])
    i += 1

print(element_dict)

while True:
    line = xdatcar.readline()
    if len(line) == 0:
        break
    xyz.write(str(N) + "\ncomment\n")
    xyz_fract.write(str(N)+"\ncomment\n")
    for el in element_names:
        for i in range(element_dict[el]):
            p = xdatcar.readline().rstrip('\n').split()
            coords = np.array([ float(s) for s in p ])
#            print coords
            cartesian_coords = coords[0]*a1+coords[1]*a2+coords[2]*a3
            xyz.write(el+ " " + str(cartesian_coords[0])+ " " + str(cartesian_coords[1]) + " " + str(cartesian_coords[2]) +"\n")
            xyz_fract.write(el+ " " + str(coords[0])+ " " + str(coords[1]) + " " + str(coords[2]) +"\n")
xdatcar.close()
xyz.close()
xyz_fract.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 range(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 range(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 range(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)