Editing Pdos.py (section)

==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