Pdos.py

From Wiki
Jump to navigation Jump to search

Source Code pdos.py

#!/usr/bin/env python
import sys, os, argparse, numpy as np

### Arguments ###
parser = argparse.ArgumentParser(description="Extracts DOS and PDOS from a DOSCAR file generated by VASP. Output: to file > extracted DOS, to terminal > number of points in grid + fermi-level. Can be spin-resolved or summed together.")
parser.add_argument("atin", type=int, help="Atom index number: Counting from 1 (0 gives the total DOS)")
parser.add_argument("magn", type=str, help="Orbital (+ magnetization): tot, s, p (px, py, pz), d (dxy, dxz, dyz, dz2, dx2), f (f-3, f-2, f-1, f0, f1, f2, f3)")
parser.add_argument("-i", "--inp", nargs='?', type=str, default='DOSCAR', help="Input file {Default: ./DOSCAR}")
parser.add_argument("-o", "--out", nargs='?', type=str, help="Output file {Default /DOS#atin_#magn.dat}")
parser.add_argument("-s", "--spin", action="store_true", help="Toggle spin-resolved (up/down)")
args = parser.parse_args()

if not args.out: args.out = "DOS"+str(args.atin)+"_"+str(args.magn)+".dat"

def read_dosfile():
    f = open(args.inp, 'r')
    lines = f.readlines()
    f.close()
    index = 5
    nedos = int(lines[index].strip().split()[2])
    efermi = float(lines[index].strip().split()[3])
    print nedos, efermi

    return lines, nedos, efermi

### WRITE DOS0 CONTAINING TOTAL DOS ###
def proj(orbit):
    	if orbit == "s": return [1]
    	elif orbit == "p": return xrange(2,4)
	elif orbit == "py": return [2]
	elif orbit == "pz": return [3]
	elif orbit == "px": return [4]
	elif orbit == "d": return xrange(5,9)
       elif orbit == "dxy": return [5]
       elif orbit == "dyz": return [6]
       elif orbit == "dz2": return [7]
       elif orbit == "dxz": return [8]
       elif orbit == "dx2": return [9]
   	elif orbit == "f": return xrange(10,16)
       elif orbit == "f-3": return [10]
       elif orbit == "f-2": return [11]
       elif orbit == "f-1": return [12]
       elif orbit == "f0": return [13]
       elif orbit == "f1": return [14]
       elif orbit == "f2": return [15]
       elif orbit == "f3": return [16]
   	elif orbit == "tot": return xrange(1,16)
   	else: return "Warning"

def proj_up(orbit):
       if orbit == "s": return [1]
       elif orbit == "p": return xrange(3,8,2)
       elif orbit == "py": return [3]
       elif orbit == "pz": return [5]
       elif orbit == "px": return [7]
       elif orbit == "d": return xrange(9,18,2)
       elif orbit == "dxy": return [9]
       elif orbit == "dyz": return [11]
       elif orbit == "dz2": return [13]
       elif orbit == "dxz": return [15]
       elif orbit == "dx2": return [17]
       elif orbit == "f": return xrange(19,32,2)
       elif orbit == "f-3": return [19]
       elif orbit == "f-2": return [21]
       elif orbit == "f-1": return [23]
       elif orbit == "f0": return [25]
       elif orbit == "f1": return [27]
       elif orbit == "f2": return [29]
       elif orbit == "f3": return [31]
       elif orbit == "tot": return xrange(1,32,2)
       else: return "Warning"

### LOOP OVER SETS OF DOS, NATOMS ###
def write_dos(lines, index, orbit, nedos, efermi, pol, spin):

f = open(args.out, 'w')

for n in xrange(nedos-1): step = int(index+n)

       	e = float(lines[step].strip().split()[0])

e_f = e-efermi f.write('%15.8f ' % (e_f)) dos = 0 dosup = 0 dosdown = 0

if pol: ncols = proj_up(orbit) for col in ncols:

                        	dosup += float(lines[step].strip().split()[col])

dosdown -= float(lines[step].strip().split()[col+1]) if spin: f.write('%15.8f %15.8f' % (dosup, dosdown)) else: dos = dosup - dosdown f.write('%15.8f' % (dos)) else: ncols = proj(orbit)

                       for col in ncols:

dos += float(lines[step].strip().split()[col]) f.write('%15.8f' % (dos))

f.write('\n')

f.close()

###  MAIN ###
if __name__ == '__main__':
    	lines, nedos, efermi = read_dosfile()

toss = lines[7].strip().split()

if len(toss)==5: pol = True elif len(toss)==3: pol = False if args.spin: print "Cannot comply: this is not a spin-polarized calculation" sys.exit(2) else: print 'Grave error' sys.exit(2)

#execute

if args.atin == 0 and args.magn == "tot": write_dos(lines, 7, "s", nedos, efermi, pol, args.spin) elif args.atin == 0: print "Use 'pdos.py 0 tot --options' to obtain the total DOS" sys.exit(2) else: index = 7 + (nedos+1)*args.atin write_dos(lines, index, args.magn, nedos, efermi, pol, args.spin)

#Written by Nathan Daelman the 31st of July, 2017