Pdos.py

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