Plot DOS and PDOS of Fe read from QE's output #QE #tool #plotting #Python

plot_dos_pdos.py

#!/usr/bin/env python -u -i

import os
import re

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.ticker import ScalarFormatter


def read_dos(file):
    e = []
    dos = []

    with open(file) as f:
        line = f.readline()
        fermi = re.search(r"EFermi =\s*(-?\d+\.\d*)\s*eV", line).group(1)
        for line in f:
            if not line.strip():
                continue
            energy, edos, *_ = line.split()
            e.append(energy)
            dos.append(edos)

    return np.array(e, dtype=np.float), np.array(dos, dtype=np.float), float(fermi)


def read_pdos(file):
    df = pd.read_csv(file, sep='\s+', skiprows=[0], header=None)
    e, pdos = df.iloc[:, 0], df.iloc[:, 2:].sum(axis=1)
    return e, pdos


def list_pdos_files(path):
    for f in os.listdir(path):
        if f.startswith('filepdos.out.pdos_atm'):
            match = re.search(
                r"pdos_atm#(\d+)\((\w+)\)\_wfc#(\d+)\((\w+)\)", f)
            if not match:
                raise FileNotFoundError
            yield f, match.groups()


if __name__ == '__main__':
    fig, ax = plt.subplots()
    e, dos, fermi = read_dos('hcpFe.dos.out')
    ax.plot(e, dos, label="total DOS")
    ax.plot([fermi] * 2, [min(dos), max(dos)], '-.', label="Fermi energy")

    d = {"s": dict(), "p": dict(), "d": dict()}
    for file, info in list_pdos_files('.'):
        atom_number, element, orbital_number, orbital_type = info
        energy, pdos = read_pdos(file)
        d[orbital_type].update({atom_number: pdos})

    total_pdos = {"s": None, "p": None, "d": None}
    for orbital_type in total_pdos.keys():
        total_pdos[orbital_type] = d[orbital_type]["1"] + d[orbital_type]["2"]
    total_pdos = pd.DataFrame(total_pdos)
    total_pdos.index = energy

    for orbital_type in total_pdos.keys():
        ax.plot(total_pdos.index,
                total_pdos[orbital_type], label="{} DOS".format(orbital_type))

    ax.legend(loc="upper right")
    ticks = np.linspace(min(energy), max(energy), 12)
    ticks = np.sort(np.append(ticks, fermi))
    ax.set_xticks(ticks)
    ax.set_xticklabels(ticks, rotation=45)
    ax.xaxis.set_major_formatter(ScalarFormatter())
    ax.set_xlabel("E (eV)")
    ax.set_ylabel("DOS")
    ax.set_title("DOS and partial DOS on different orbitals of Fe")
    plt.savefig('dos.pdf')
    plt.show()