plot_qe_bands_1d.py 2.78 KB
Newer Older
spiasko's avatar
spiasko committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
#!/users/keimre/soft/miniconda3/bin/python

import numpy as np

import xml.etree.ElementTree as et

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.colors as colors

import argparse
import os


parser = argparse.ArgumentParser(
    description='Plots quantum espresso bands')
parser.add_argument(
    'main_xml',
    metavar='XML',
    help='.xml file of the QE bands calculation')
parser.add_argument(
    '--add_xml',
    type=str,
    default=None,
    metavar='DIR',
    help='Additional .xml file of the QE bands calculation')
parser.add_argument(
    '--fold',
    type=int,
    default=1,
    metavar='N',
    help='Number of times to unfold the additional bands')
parser.add_argument(
    '--emin',
    type=float,
    default=-2.0,
    metavar='EMIN',
    help='Emin of the plot')
parser.add_argument(
    '--emax',
    type=float,
    default=+3.5,
    metavar='EMAX',
    help='Emax of the plot')

args = parser.parse_args()


def read_band_data_new_xml(xml_file):
    """
    Reads data from QE bands calculations (new XML)
    Returns:
      - kpts[i_kpt] = [kx, ky, kz] in [2*pi/a] 
      - eigvals[i_kpt, i_band] in [eV]
      - fermi_en in [eV]
    """
    
    data_file_xml = et.parse(xml_file)
    data_file_root = data_file_xml.getroot()

    output_node = data_file_root.find('output')

    # Find fermi
    band_node = output_node.find('band_structure')
    fermi_en = -3# float(band_node.find('fermi_energy').text)*27.21138602
    lsda = band_node.find('spinorbit').text

    kpts = []
    eigvals = []

    for kpt in band_node.findall("ks_energies"):
        k_coords = np.array(kpt.find('k_point').text.split(), dtype=float)
        kpts.append(k_coords)

        eig_vals = np.array(kpt.find('eigenvalues').text.split(), dtype=float)
        eigvals.append(eig_vals*27.21138602)
    kpts = np.array(kpts)
    eigvals = np.array(eigvals)
    
    return kpts, eigvals, fermi_en

plt.figure(figsize=(6, 10))

if args.add_xml != None:
        
    kpts2, bands2, fermi_en2 = read_band_data_new_xml(args.add_xml)
    bands2 -= fermi_en2

    for i_band2 in range(bands2.shape[1]):
        plt.plot(kpts2[:, 0]/args.fold, bands2[:, i_band2], '-', color='b')

kpts, bands, fermi_en = read_band_data_new_xml(args.main_xml)
bands -= fermi_en

for i_band in range(bands.shape[1]):
    plt.plot(kpts[:, 0], bands[:, i_band], '--', color='r')

plt.ylim([args.emin, args.emax])
plt.xlim([np.min(kpts[:, 0]), np.max(kpts[:, 0])])

plt.ylabel("Energy [eV]")
plt.xlabel(r"k [$2\pi /a$]")

# Make a suitable name
name = os.path.basename(args.main_xml)
name = name.split('.')[0]
if args.add_xml != None:
    add_name = os.path.basename(args.add_xml)
    add_name = add_name.split('.')[0]
    name += "-" + add_name

plt.savefig("./%s.png"%name, dpi=300)
plt.close()
print("Wrote %s.png!"%name)