vasptdos

#!/usr/bin/env python3

import os
from sys import argv
import xml.dom.minidom
import numpy as np

set_efermi = False

if len(argv) == 1:
	fxml = "vasprun.xml"
	print("Input file is not specfied, try to read \"vasprun.xml\".")
elif len(argv) <= 3:
	fxml = argv[1]
	if len(argv) == 3:
		efermi = float(argv(2))
		set_efermi = True
		print("The fermi level is set to % .8f." % efermi)
if not os.access(fxml, os.R_OK):
	print("File \"%s\" is not readable." % fxml)
	exit()

vasp = xml.dom.minidom.parse(fxml)
dos = vasp.getElementsByTagName("dos")[0]
if not set_efermi:
	for idom in dos.getElementsByTagName("i"):
		if idom.hasAttribute("name") and idom.getAttribute("name") == "efermi":
			efermi = float(idom.firstChild.data)
			print("Read fermi level from input file: % .8f." % efermi)
			break

tdos = dos.getElementsByTagName("total")[0].getElementsByTagName("array")[0].getElementsByTagName("set")[0]
e_dos = []
for spin in tdos.getElementsByTagName("set"):
	lines = []
	for line in spin.getElementsByTagName("r"):
		lines.append(line.firstChild.data.split()[:2])
	e_dos.append(np.array(lines, dtype=np.double))
	e_dos[-1][:, 0] -= efermi
if len(e_dos) == 1:
	np.savetxt('dos.dat', e_dos[0], fmt=" %15.8f")
elif len(e_dos) == 2:
	np.savetxt('dos.dat', np.column_stack((e_dos[0], e_dos[1][:, 1] * -1)), fmt=" %15.8f")
else:
	print("The number of spin > 2.")