class DosPlotterTest(unittest.TestCase):
def setUp(self):
with open(os.path.join(test_dir, "complete_dos.json"), "r") as f:
self.dos = CompleteDos.from_dict(json.load(f))
self.plotter = DosPlotter(sigma=0.2, stack=True)
def test_add_dos_dict(self):
d = self.plotter.get_dos_dict()
self.assertEqual(len(d), 0)
self.plotter.add_dos_dict(self.dos.get_element_dos(), key_sort_func=lambda x:x.X)
d = self.plotter.get_dos_dict()
self.assertEqual(len(d), 4)
def test_get_dos_dict(self):
self.plotter.add_dos_dict(self.dos.get_element_dos(), key_sort_func=lambda x:x.X)
d = self.plotter.get_dos_dict()
for el in ["Li", "Fe", "P", "O"]:
self.assertIn(el, d)
parser = argparse.ArgumentParser(description='''Convenient DOS Plotter for vasp runs.
Author: Shyue Ping Ong
Version: 1.0
Last updated: Nov 8 2011''')
parser.add_argument('filename', metavar='filename', type=str, nargs=1, help='vasprun.xml file to plot')
parser.add_argument('-s', '--site', dest='site', action='store_const', const=True, help='plot site projected DOS')
parser.add_argument('-e', '--element', dest='element', action='store_const', const=True, help='plot element projected DOS')
parser.add_argument('-o', '--orbital', dest="orbital", action='store_const', const=True, help='plot orbital projected DOS')
args = parser.parse_args()
v = Vasprun(args.filename[0])
dos = v.complete_dos
all_dos = OrderedDict()
all_dos['Total'] = dos
structure = v.final_structure
if args.site:
for i in xrange(len(structure)):
site = structure[i]
all_dos['Site ' + str(i) + " " + site.specie.symbol] = dos.get_site_dos(site)
if args.element:
all_dos.update(dos.get_element_dos())
if args.orbital:
all_dos.update(dos.get_spd_dos())
plotter = DosPlotter()
plotter.add_dos_dict(all_dos)
plotter.show()
def setUp(self):
with open(os.path.join(test_dir, "complete_dos.json"), "r") as f:
self.dos = CompleteDos.from_dict(json.load(f))
self.plotter = DosPlotter(sigma=0.2, stack=True)
