Source code for pyprocar.scriptBandsDosplot

"""
Created on May 17 2020
@author: Pedram Tavadze
"""
import re

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np

from .doscarplot import DosPlot
from .vaspxml import VaspXML
from .lobsterparser import LobsterDOSParser, LobsterParser
from .qeparser import QEDOSParser, QEParser

from .abinitparser import AbinitParser
from .doscarplot import DosPlot
from .elkparser import ElkParser
from .procarparser import ProcarParser
from .procarplot import ProcarPlot
from .procarselect import ProcarSelect
from .splash import welcome
from .utilsprocar import UtilsProcar

plt.rcParams["mathtext.default"] = "regular"
# Roman ['rm', 'cal', 'it', 'tt', 'sf',
#        'bf', 'default', 'bb', 'frak',
#        'circled', 'scr', 'regular']
plt.rcParams["font.family"] = "Arial"
plt.rc("font", size=18)  # controls default text sizes
plt.rc("axes", titlesize=22)  # fontsize of the axes title
plt.rc("axes", labelsize=22)  # fontsize of the x and y labels
plt.rc("xtick", labelsize=22)  # fontsize of the tick labels
plt.rc("ytick", labelsize=22)  # fontsize of the tick labels
plt.rc("legend", fontsize=18)  # legend fontsize
# plt.rc('figure', titlesize=22)  # fontsize of the figure title


[docs]def bandsdosplot( bands_file="PROCAR", dos_file="vasprun.xml", outcar="OUTCAR", abinit_output=None, bands_mode="plain", dos_mode="plain", plot_total=True, fermi=None, mask=None, markersize=0.02, marker="o", atoms=None, orbitals=None, bands_spin=0, bands_separate=False, dos_spins=None, dos_labels=None, dos_spin_colors=[(1, 0, 0), (0, 0, 1)], dos_colors=None, dos_items=None, dos_interpolation_factor=None, dlimit=None, elimit=None, vmin=None, vmax=None, cmap="jet", grid=False, kpointsfile="KPOINTS", code="vasp", savefig=None, title=None, kdirect=True, discontinuities=None, plot_color_bar=True, repair=True, ): """This function creates plots containing both DOS and bands.""" welcome() # Repair PROCAR if code == "vasp" or code == "abinit": if repair: repairhandle = UtilsProcar() repairhandle.ProcarRepair(bands_file, bands_file) fig = plt.figure(figsize=(13, 7), constrained_layout=False) widths = [13, 5] heights = [9] gs = fig.add_gridspec(1, 2, width_ratios=widths, height_ratios=heights) ax1 = fig.add_subplot(gs[0, 0]) ax2 = fig.add_subplot(gs[0, 1]) total = plot_total if atoms is None: bands_atoms = [-1] else: bands_atoms = atoms.copy() if orbitals is None: bands_orbitals = [-1] else: bands_orbitals = ( orbitals.copy() ) # these copying is because orbitals select changes shifts orbitals by one if dos_spins is None: dos_spins = [0, 1] if dos_mode not in [ "plain", "parametric_line", "parametric", "stack_species", "stack_orbitals", "stack", ]: raise ValueError( "Mode should be choosed from ['plain', 'parametric_line','parametric','stack_species','stack_orbitals','stack']" ) # Verbose section print("Script initiated...") if repair: print("PROCAR repaired. Run with repair=False next time.") print("code : ", code) print("bands file : ", bands_file) print("bands mode : ", bands_mode) print("bands spin : ", bands_spin) print("dos file : ", dos_file) print("dos mode : ", dos_mode) print("dos spins : ", dos_spins) print("atoms list : ", atoms) print("orbs. list : ", orbitals) if fermi is None and outcar is None and abinit_output is None: print( "WARNING : Fermi Energy not set! Please set manually or provide output file and set code type." ) fermi = 0 print("fermi energy : ", fermi) print("energy range : ", elimit) if mask is not None: print("masking thres. : ", mask) print("colormap : ", cmap) print("markersize : ", markersize) print("vmax : ", vmax) print("vmin : ", vmin) print("grid enabled : ", grid) print("savefig : ", savefig) print("title : ", title) print("outcar : ", outcar) if kdirect: print("k-grid : reduced") else: print( "k-grid : cartesian (Remember to provide an output file for this case to work.)" ) if discontinuities is None: discontinuities = [] #### READING KPOINTS FILE IF PRESENT #### # If KPOINTS file is given: if code == "vasp": if kpointsfile is not None: # Getting the high symmetry point names from KPOINTS file f = open(kpointsfile) KPread = f.read() f.close() KPmatrix = re.findall("reciprocal[\s\S]*", KPread) tick_labels = np.array(re.findall("!\s(.*)", KPmatrix[0])) knames = [] knames = [tick_labels[0]] ################## Checking for discontinuities ######################## discont_indx = [] icounter = 1 while icounter < len(tick_labels) - 1: if tick_labels[icounter] == tick_labels[icounter + 1]: knames.append(tick_labels[icounter]) icounter = icounter + 2 else: discont_indx.append(icounter) knames.append( tick_labels[icounter] + "|" + tick_labels[icounter + 1] ) icounter = icounter + 2 knames.append(tick_labels[-1]) discont_indx = list(dict.fromkeys(discont_indx)) ################# End of discontinuity check ########################## # Added by Nicholas Pike to modify the output of seekpath to allow for # latex rendering. for i in range(len(knames)): if knames[i] == "GAMMA": knames[i] = "\Gamma" else: pass knames = [str("$" + latx + "$") for latx in knames] # getting the number of grid points from the KPOINTS file f2 = open(kpointsfile) KPreadlines = f2.readlines() f2.close() numgridpoints = int(KPreadlines[1].split()[0]) kticks = [0] gridpoint = 0 for kt in range(len(knames) - 1): gridpoint = gridpoint + numgridpoints kticks.append(gridpoint - 1) print("knames : ", knames) print("kticks : ", kticks) # creating an array for discontunuity k-points. These are the indexes # of the discontinuity k-points. for k in discont_indx: discontinuities.append(kticks[int(k / 2) + 1]) if discontinuities: print("discont. list : ", discontinuities) elif code == "lobster": procarFile = LobsterParser() kticks = procarFile.kticks knames = procarFile.knames print("knames : ", knames) print("kticks : ", kticks) if procarFile.discontinuities: discontinuities = procarFile.discontinuities vaspxml = LobsterDOSParser() dos_plot = DosPlot(dos=vaspxml.dos, structure=vaspxml.structure) if dos_spins is None: dos_spins = np.arange(len(vaspxml.dos.total)) elif code == "qe": procarFile = QEParser() kticks = procarFile.kticks knames = procarFile.knames print("knames : ", knames) print("kticks : ", kticks) # Retrieving knames and kticks from QE if procarFile.discontinuities: discontinuities = procarFile.discontinuities vaspxml = QEDOSParser() dos_plot = DosPlot(dos=vaspxml.dos, structure=vaspxml.structure) if dos_spins is None: dos_spins = np.arange(len(vaspxml.dos.total)) #### END OF KPOINTS FILE DEPENDENT SECTION #### # spin = {"0": 0, "1": 1, "2": 2, "3": 3, "st": "st"}[str(spin)] #### parsing the PROCAR file or equivalent to retrieve spd data #### code = code.lower() if code == "vasp": procarFile = ProcarParser() vaspxml = VaspXML( filename=dos_file, dos_interpolation_factor=dos_interpolation_factor ) dos_plot = DosPlot(dos=vaspxml.dos, structure=vaspxml.structure) if dos_spins is None: dos_spins = np.arange(len(vaspxml.dos.total)) # elif code == "lobster": # procarFile = LobsterParser() # vaspxml = LobsterDOSParser() # dos_plot = DosPlot(dos=vaspxml.dos, structure=vaspxml.structure) # If ticks and names are given by the user manually: if kticks is not None and knames is not None: ticks = list(zip(kticks, knames)) elif kticks is not None: ticks = list(zip(kticks, kticks)) else: ticks = None if kpointsfile is None and kticks and knames: print("knames : ", knames) print("kticks : ", kticks) if discontinuities: print("discont. list : ", discontinuities) # The second part of this function is parse/select/use the data in # OUTCAR (if given) and PROCAR # first parse the outcar if given, to get Efermi and Reciprocal lattice recLat = None if code == "vasp": if outcar: outcarparser = UtilsProcar() if fermi is None: fermi = outcarparser.FermiOutcar(outcar) print("Fermi energy : %s eV (from OUTCAR)" % str(fermi)) recLat = outcarparser.RecLatOutcar(outcar) elif code == "lobster": fermi = procarFile.fermi recLat = procarFile.reclat elif code == "qe": fermi = procarFile.fermi recLat = procarFile.reclat # if kdirect = False, then the k-points will be in cartesian coordinates. # The output should be read to find the reciprocal lattice vectors to transform # from direct to cartesian if code == "vasp": if kdirect: procarFile.readFile(bands_file, permissive=False) else: procarFile.readFile(bands_file, permissive=False, recLattice=recLat) # processing the data, getting an instance of the class that reduces the data data = ProcarSelect(procarFile, deepCopy=True, mode=bands_mode) numofbands = int(data.spd.shape[1] / 2) # handling the spin, `spin='st'` is not straightforward, needs # to calculate the k vector and its normal. Other `spin` values # are trivial. if bands_spin == "st": # two `ProcarSelect` instances, to store temporal values: spin_x, spin_y dataX = ProcarSelect(procarFile, deepCopy=True) dataX.selectIspin([1]) dataX.selectAtoms(bands_atoms, fortran=False) dataX.selectOrbital(bands_orbitals) dataY = ProcarSelect(procarFile, deepCopy=True) dataY.selectIspin([2]) dataY.selectAtoms(bands_atoms, fortran=False) dataY.selectOrbital(bands_orbitals) # getting the signed angle of each K-vector angle = np.arctan2(dataX.kpoints[:, 1], (dataX.kpoints[:, 0] + 0.000000001)) sin = np.sin(angle) cos = np.cos(angle) sin.shape = (sin.shape[0], 1) cos.shape = (cos.shape[0], 1) # print sin, cos # storing the spin projection into the original array data.spd = -sin * dataX.spd + cos * dataY.spd else: data.selectIspin([bands_spin], separate=bands_separate) data.selectAtoms(bands_atoms, fortran=False) data.selectOrbital(bands_orbitals) # Plotting the data if bands_separate: if bands_spin == 0: # plotting spin up bands separately data.bands = ( data.bands[:, :numofbands].transpose() - np.array(fermi) ).transpose() print("Plotting spin up bands...") elif bands_spin == 1: # plotting spin down bands separately data.bands = ( data.bands[:, numofbands:].transpose() - np.array(fermi) ).transpose() print("Plotting spin down bands...") plot = ProcarPlot(data.bands, data.spd, data.kpoints, ax=ax1) else: # Regular plotting method. For spin it plots density or magnetization. if code == "lobster" or code == "qe": data.bands = (data.bands.transpose()).transpose() else: data.bands = (data.bands.transpose() - np.array(fermi)).transpose() plot = ProcarPlot(data.bands, data.spd, data.kpoints) if vmin is None: vmin = plot.spd.min() if vmax is None: vmax = plot.spd.max() ###### start of mode dependent options ######### if bands_mode == "scatter": _, ax1 = plot.scatterPlot( cmap=cmap, vmin=vmin, vmax=vmax, ticks=ticks, discontinuities=discontinuities, ax=ax1, mask=mask, ) if fermi is not None: ax1.set_ylabel(r"$E-E_f$ [eV]") else: ax1.set_ylabel(r"Energy [eV]") if elimit is not None: ax1.set_ylim(elimit) elif bands_mode == "plain": _, ax1 = plot.plotBands(ticks=ticks, discontinuities=discontinuities, ax=ax1,) if fermi is not None: ax1.set_ylabel(r"$E-E_f$ [eV]") else: ax1.set_ylabel(r"Energy [eV]") if elimit: ax1.set_ylim(elimit) elif bands_mode == "parametric": if dos_mode == "parametric": plot_color_bar = False else: plot_color_bar = True _, ax1 = plot.parametricPlot( cmap=cmap, vmin=vmin, vmax=vmax, ticks=ticks, discontinuities=discontinuities, ax=ax1, plot_bar=plot_color_bar, mask=mask, ) if fermi is not None: ax1.set_ylabel(r"$E-E_f$ [eV]") else: ax1.set_ylabel(r"Energy [eV]") if elimit is not None: ax1.set_ylim(elimit) if dos_mode == "plain": _, ax2 = dos_plot.plot_total( spins=dos_spins, ax=ax2, orientation="vertical", labels=dos_labels ) elif dos_mode == "parametric_line": if not total: _, ax2 = dos_plot.plot_parametric_line( atoms=atoms, spins=dos_spins, orbitals=orbitals, spin_colors=dos_spin_colors, ax=ax2, orientation="vertical", labels=dos_labels, ) else: _, ax2 = dos_plot.plot_total( spins=dos_spins, spin_colors=[(0, 0, 0), (0, 0, 0)], ax=ax2, orientation="vertical", ) _, ax2 = dos_plot.plot_parametric_line( atoms=atoms, spins=dos_spins, orbitals=orbitals, spin_colors=dos_spin_colors, ax=ax2, orientation="vertical", labels=dos_labels, ) elif dos_mode == "parametric": if not total: _, ax2 = dos_plot.plot_parametric( atoms=atoms, spins=dos_spins, orbitals=orbitals, spin_colors=dos_spin_colors, cmap=cmap, elimit=elimit, ax=ax2, vmin=vmin, vmax=vmax, orientation="vertical", labels=dos_labels, ) else: _, ax1 = dos_plot.plot_parametric( atoms=atoms, spins=dos_spins, orbitals=orbitals, spin_colors=dos_spin_colors, cmap=cmap, elimit=elimit, ax=ax2, vmin=vmin, vmax=vmax, orientation="vertical", labels=dos_labels, ) _, ax2 = dos_plot.plot_total( spins=dos_spins, spin_colors=[(0, 0, 0), (0, 0, 0)], ax=ax2, orientation="vertical", ) ax2.yaxis.set_visible(False) elif dos_mode == "stack_species": if not total: _, ax2 = dos_plot.plot_stack_species( spins=dos_spins, spin_colors=dos_spin_colors, colors=dos_colors, elimit=elimit, figsize=(12, 6), ax=ax2, orientation="vertical", ) # dos = dos_plot.dos_total else: _, ax2 = dos_plot.plot_stack_species( spins=dos_spins, spin_colors=dos_spin_colors, colors=dos_colors, elimit=elimit, figsize=(12, 6), ax=ax2, orientation="vertical", ) # dos = dos_plot.dos_total _, ax2 = dos_plot.plot_total( spins=dos_spins, spin_colors=[(0, 0, 0), (0, 0, 0)], ax=ax2, orientation="vertical", ) elif dos_mode == "stack_orbitals": if not total: _, ax2 = dos_plot.plot_stack_orbitals( spins=dos_spins, spin_colors=dos_spin_colors, colors=dos_colors, elimit=elimit, figsize=(12, 6), ax=ax2, orientation="vertical", ) else: _, ax2 = dos_plot.plot_stack_orbitals( spins=dos_spins, spin_colors=dos_spin_colors, colors=dos_colors, elimit=elimit, figsize=(12, 6), ax=ax2, orientation="vertical", ) _, ax2 = dos_plot.plot_total( spins=dos_spins, spin_colors=[(0, 0, 0), (0, 0, 0)], ax=ax2, orientation="vertical", ) elif dos_mode == "stack": if not total: _, ax2 = dos_plot.plot_stack( items=dos_items, spins=dos_spins, spin_colors=dos_spin_colors, colors=dos_colors, elimit=elimit, figsize=(12, 6), ax=ax2, orientation="vertical", ) else: _, ax2 = dos_plot.plot_stack( items=dos_items, spins=dos_spins, spin_colors=dos_spin_colors, colors=dos_colors, elimit=elimit, figsize=(12, 6), ax=ax2, orientation="vertical", ) _, ax2 = dos_plot.plot_total( spins=dos_spins, spin_colors=[(0, 0, 0), (0, 0, 0)], ax=ax2, orientation="vertical", ) ax2.set_ylim(elimit) ax2.axhline(color="black", linestyle="--") ax2.axvline(color="black", linestyle="--") if grid: ax1.grid() ax2.grid() # ax2.yaxis.set_ticklabels([]) if dos_labels or "stack" in dos_mode: ax2.legend() ax2.yaxis.set_visible(False) cond1 = vaspxml.dos.energies >= elimit[0] cond2 = vaspxml.dos.energies <= elimit[1] cond = np.all([cond1, cond2], axis=0) if len(dos_spins) > 1: ylim = [ vaspxml.dos.total[1, cond].max() * -1.1, vaspxml.dos.total[0, cond].max() * 1.1, ] else: ylim = [0, vaspxml.dos.total[dos_spins[0], cond].max() * 1.1] if dlimit is not None: ax2.set_xlim(dlimit[0], dlimit[1]) elif ( dos_mode == "stack_species" or dos_mode == "stack_orbitals" or dos_mode == "stack" ): ax2.set_xlim(ax2.get_xlim()[0], ax2.get_xlim()[1] * 1.1) elif dos_mode == "plain": ax2.set_xlim(ylim[0], ylim[1]) if title: ax1.set_title(title) fig.tight_layout() if savefig: fig.savefig(savefig, bbox_inches="tight") plt.close() return None, None else: plt.show() return fig, ax1, ax2