Source code for pyprocar.doscarplot.dosplot

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

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

np.seterr(divide="ignore", invalid="ignore")
figsize = (12, 6)


[docs]class DosPlot: def __init__(self, dos=None, structure=None): """ Parameters ---------- vaspxml : TYPE, optional DESCRIPTION. The default is "vasprun.xml". Returns ------- None. """ self.dos = dos self.structure = structure return
[docs] def plot_total(self, spins=None, spin_colors=None, figsize=figsize, ax=None, orientation="horizontal", labels=None, linewidth=1, ): """ Parameters ---------- spins : TYPE, optional DESCRIPTION. The default is None. spin_colors : TYPE, optional DESCRIPTION. The default is None. figsize : TYPE, optional DESCRIPTION. The default is figsize. ax : TYPE, optional DESCRIPTION. The default is None. orientation : TYPE, optional DESCRIPTION. The default is "horizontal". labels : TYPE, optional DESCRIPTION. The default is None. : TYPE DESCRIPTION. Returns ------- fig : TYPE DESCRIPTION. ax : TYPE DESCRIPTION. """ if spin_colors is None: spin_colors = [(1, 0, 0), (0, 0, 1)] energies = self.dos.energies dos = np.array(self.dos.total) if spins is None: spins = np.arange(len(self.dos.total)) fig, ax = plotter(energies, dos, spins, spin_colors, figsize, ax, orientation, linewidth, labels) return fig, ax
[docs] def plot_parametric_line(self, atoms=None, principal_q_numbers=[-1], orbitals=None, spins=None, spin_colors=None, figsize=(12, 6), ax=None, orientation="horizontal", labels=None, linewidth=1, ): if ax is None: if orientation == "horizontal": fig = plt.figure(figsize=figsize) ax = fig.add_subplot(111) elif orientation == "vertical": fig = plt.figure(figsize=(figsize[1], figsize[0])) ax = fig.add_subplot(111) else: fig = ax.get_figure() if spin_colors is None: spin_colors = [(1, 0, 0), (0, 0, 1)] # dos = self.parsedData.dos_parametric(atoms=atoms, # spin=spins, # orbitals=orbitals) if spins is None: spins = np.arange(len(self.dos.total)) dos = self.dos.dos_sum(atoms, principal_q_numbers, orbitals, spins) fig, ax = plotter(self.dos.energies, dos, spins, spin_colors, figsize, ax, orientation, linewidth, labels) return fig, ax
[docs] def plot_parametric(self, atoms=None, principal_q_numbers=[-1], orbitals=None, spins=None, spin_colors=None, cmap="jet", vmin=0, vmax=1, elimit=None, figsize=(12, 6), ax=None, orientation="horizontal", labels=None, plot_bar=True): if ax is None: if orientation == "horizontal": fig = plt.figure(figsize=figsize) ax = fig.add_subplot(111) elif orientation == "vertical": fig = plt.figure(figsize=(figsize[1], figsize[0])) ax = fig.add_subplot(111) else: fig = ax.get_figure() if spin_colors is None: spin_colors = [(1, 0, 0), (0, 0, 1)] cmap = mpl.cm.get_cmap(cmap) dos_total = np.array(self.dos.total) dos_total_projected = self.dos.dos_sum() if spins is None: spins = np.arange(len(self.dos.total)) dos = self.dos.dos_sum(atoms=atoms, principal_q_numbers=principal_q_numbers, orbitals=orbitals, spins=spins) # dos_total = dos # dos_total = self.parsedData.dos_total # dos_total_projected = self.parsedData.dos_parametric() # dos = self.parsedData.dos_parametric(atoms=atoms, # spin=spins, # orbitals=orbitals) if vmin is None or vmax is None: vmin = 0 vmax = (dos.max() / dos_total.max()) if plot_bar: norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax) fig.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap), ax=ax) if not elimit: elimit = [self.dos.energies.min(), self.dos.energies.max()] deltaE = elimit[1]-elimit[0] cond1 = self.dos.energies >= elimit[0] - deltaE*0.05 cond2 = self.dos.energies <= elimit[1] + deltaE*0.05 cond = np.all([cond1, cond2], axis=0) # dE = dos.energies[1] - dos.energies[0] if spins is None: spins = np.arange(len(self.dos.total)) for ispin in spins: x = [] y_total = [] bar_color = [] for idos in range(len(self.dos.energies[cond])): x.append(self.dos.energies[cond][idos]) y = dos[ispin, cond][idos] y_total.append(dos_total[ispin, cond][idos]) y_total_projected = dos_total_projected[ispin, cond][idos] if ispin > 0 and len(spins) > 1: y *= -1 y_total[-1] *= -1 y_total_projected *= -1 bar_color.append(cmap(y / (y_total_projected * (vmax - vmin)))) if orientation == "horizontal": # ax.bar(x,y_total,dE,color=bar_color) for idos in range(len(x) - 1): ax.fill_between([x[idos], x[idos + 1]], [y_total[idos], y_total[idos + 1]], color=bar_color[idos]) elif orientation == "vertical": for idos in range(len(x) - 1): ax.fill_betweenx([x[idos], x[idos + 1]], [y_total[idos], y_total[idos + 1]], color=bar_color[idos]) return fig, ax
[docs] def plot_stack_species( self, principal_q_numbers=[-1], orbitals=None, spins=None, spin_colors=None, colors=None, elimit=None, figsize=(12, 6), ax=None, orientation="horizontal", ): if spin_colors is None: spin_colors = [(0, 0, 1), (1, 0, 0)] if colors is None: colors = [ 'red', 'green', 'blue', 'cyan', 'magenta', 'yellow', 'orange', 'purple', 'brown', 'navy', 'maroon', 'olive' ] if ax is None: if orientation == "horizontal": fig = plt.figure(figsize=figsize) ax = fig.add_subplot(111) elif orientation == "vertical": fig = plt.figure(figsize=(figsize[1], figsize[0])) ax = fig.add_subplot(111) else: fig = ax.get_figure() if spins is None: spins = np.arange(len(self.dos.total)) if not elimit: elimit = [self.dos.energies.min(), self.dos.energies.max()] # dos_projected_total = self.parsedData.dos_parametric(spin=spins, # orbitals=orbitals) if len(self.dos.projected[0][0]) == 1 + 3 + 5: all_orbitals = "spd" elif len(self.dos.projected[0][0]) == 1 + 3 + 5 + 7: all_orbitals = "spdf" else: all_orbitals = "" label = "" if orbitals: print("The plot only considers orbitals", orbitals) label = "-" if sum([x in orbitals for x in [0]]) == 1: label += "s" if sum([x in orbitals for x in [1, 2, 3]]) == 3: label += "p" if sum([x in orbitals for x in [4, 5, 6, 7, 8]]) == 5: label += "d" if sum([x in orbitals for x in [9, 10, 11, 12, 13, 14, 15]]) == 7: label += "f" if label == "-" + all_orbitals: label = "" # dos_total = self.parsedData.dos_total # cond1 = self.parsedData.dos_total.energies >= elimit[0] # cond2 = self.parsedData.dos_total.energies <= elimit[1] # cond = np.all([cond1, cond2], axis=0) dos_total = self.dos.total dos_projected_total = self.dos.dos_sum() deltaE = elimit[1]-elimit[0] cond1 = self.dos.energies >= elimit[0] - deltaE*0.05 cond2 = self.dos.energies <= elimit[1] + deltaE*0.05 cond = np.all([cond1, cond2], axis=0) for ispin in spins: # bottom = np.zeros_like(self.parsedData.dos_total.energies[cond]) bottom = np.zeros_like(self.dos.energies[cond]) for ispc in range(len(self.structure.species)): idx = (np.array( self.structure.atoms) == self.structure.species[ispc]) atoms = list(np.where(idx)[0]) # dos = self.parsedData.dos_parametric(atoms=atoms, # spin=spins, # orbitals=orbitals) dos = self.dos.dos_sum(atoms=atoms, principal_q_numbers=principal_q_numbers, orbitals=orbitals, spins=spins) x = self.dos.energies[cond] y = (dos[ispin, cond] * dos_total[ispin, cond]) / dos_projected_total[ispin, cond] if ispin > 0 and len(spins) > 1: y *= -1 if orientation == "horizontal": ax.fill_between(x, bottom + y, bottom, color=colors[ispc]) elif orientation == "vertical": ax.fill_betweenx(x, bottom + y, bottom, color=colors[ispc]) else: if orientation == "horizontal": ax.fill_between( x, bottom + y, bottom, color=colors[ispc], label=self.structure.species[ispc] + label, ) elif orientation == "vertical": ax.fill_betweenx( x, bottom + y, bottom, color=colors[ispc], label=self.structure.species[ispc] + label, ) bottom += y return fig, ax
[docs] def plot_stack_orbitals( self, atoms=None, principal_q_numbers=[-1], spins=None, spin_colors=None, colors=None, elimit=None, figsize=(12, 6), ax=None, orientation="horizontal", ): if spin_colors is None: spin_colors = [(0, 0, 0), (0, 0, 0)] if colors is None: colors = [ 'red', 'blue', 'green', 'cyan', 'magenta', 'yellow', 'orange', 'purple', 'brown', 'navy', 'maroon', 'olive' ] if spins is None: spins = np.arange(len(self.dos.total), dtype=np.int) if ax is None: if orientation == "horizontal": fig = plt.figure(figsize=figsize) ax = fig.add_subplot(111) elif orientation == "vertical": fig = plt.figure(figsize=(figsize[1], figsize[0])) ax = fig.add_subplot(111) else: fig = ax.get_figure() atom_names = "" if atoms: print( "The plot only considers atoms", np.array(self.structure.atoms)[atoms], ) atom_names = "" for ispc in np.unique(np.array(self.structure.atoms)[atoms]): atom_names += ispc + "-" all_atoms = "" for ispc in np.unique(np.array(self.structure.atoms)): all_atoms += ispc + "-" if atom_names == all_atoms: atom_names = "" dos_total = self.dos.total dos_projected_total = self.dos.dos_sum() if not elimit: elimit = [self.dos.energies.min(), self.dos.energies.max()] # dos_projected_total = self.parsedData.dos_parametric() deltaE = elimit[1]-elimit[0] cond1 = self.dos.energies >= elimit[0] - deltaE*0.05 cond2 = self.dos.energies <= elimit[1] + deltaE*0.05 cond = np.all([cond1, cond2], axis=0) orb_names = ["s", "p", "d"] orb_l = [[0], [1, 2, 3], [4, 5, 6, 7, 8]] for ispin in spins: bottom = np.zeros_like(self.dos.energies[cond]) for iorb in range(3): # dos = self.parsedData.dos_parametric(atoms=atoms, # spin=spins, # orbitals=orb_l[iorb]) dos = self.dos.dos_sum(atoms=atoms, principal_q_numbers=principal_q_numbers, orbitals=orb_l[iorb], spins=spins) x = self.dos.energies[cond] y = (dos[ispin, cond] * dos_total[ispin, cond]) / dos_projected_total[ispin, cond] y = np.nan_to_num(y, 0) if ispin > 0 and len(spins) > 1: y *= -1 if orientation == "horizontal": ax.fill_between(x, bottom + y, bottom, color=colors[iorb]) elif orientation == "vertical": ax.fill_betweenx(x, bottom + y, bottom, color=colors[iorb]) else: if orientation == "horizontal": ax.fill_between( x, bottom + y, bottom, color=colors[iorb], label=atom_names + orb_names[iorb], ) elif orientation == "vertical": ax.fill_betweenx( x, bottom + y, bottom, color=colors[iorb], label=atom_names + orb_names[iorb], ) bottom += y return fig, ax
[docs] def plot_stack( self, items={}, spins=None, spin_colors=None, colors=None, elimit=None, figsize=(12, 6), ax=None, orientation="horizontal", ): if len(items) == 0: print("""Please provide the stacking items in which you want to plot, example : {'Sr':[1,2,3],'O':[4,5,6,7,8]} will plot the stacked plots of p orbitals of Sr and d orbitals of Oxygen.""") if spin_colors is None: spin_colors = [(0, 0, 1), (1, 0, 0)] src_colors = colors if src_colors is None: src_colors = [ 'red', 'blue', 'green', 'cyan', 'magenta', 'yellow', 'orange', 'purple', 'brown', 'navy', 'maroon', 'olive' ] if spins is None: spins = np.arange(len(self.dos.total), dtype=np.int) if ax is None: if orientation == "horizontal": fig = plt.figure(figsize=figsize) ax = fig.add_subplot(111) elif orientation == "vertical": fig = plt.figure(figsize=(figsize[1], figsize[0])) ax = fig.add_subplot(111) else: fig = ax.get_figure() if not elimit: elimit = [self.dos.energies.min(), self.dos.energies.max()] # dos_total = self.parsedData.dos_total dos_total = self.dos.total # if self.dos.dos_projected[0].ncols == (1 + 3 + # 5) * dos_total.ncols: # all_orbitals = "spd" # elif self.parsedData.dos_projected[0].ncols == (1 + 3 + 5 + # 7) * dos_total.ncols: # all_orbitals = "spdf" # else: # all_orbitals = "" if len(self.dos.projected[0][0]) == (1 + 3 + 5): all_orbitals = "spd" elif len(self.dos.projected[0][0]) == (1 + 3 + 5 + 7): all_orbitals = "spdf" else: all_orbitals = "" deltaE = elimit[1]-elimit[0] cond1 = self.dos.energies >= elimit[0] - deltaE*0.05 cond2 = self.dos.energies <= elimit[1] + deltaE*0.05 cond = np.all([cond1, cond2], axis=0) counter = 0 colors = {} for ispc in items: colors[ispc] = src_colors[counter] counter += 1 # dos_projected_total = self.parsedData.dos_parametric(spin=spins) dos_projected_total = self.dos.dos_sum(spins=spins) for ispin in spins: bottom = np.zeros_like(self.dos.energies[cond]) # bottom = np.zeros_like(self.parsedData.dos_total.energies[cond]) # bottom = np.zeros_like(self.VaspXML.dos_total.energies[:]) for ispc in items: idx = np.array(self.structure.atoms) == ispc atoms = list(np.where(idx)[0]) orbitals = items[ispc] # dos = self.parsedData.dos_parametric(atoms=atoms, # spin=spins, # orbitals=orbitals) dos = self.dos.dos_sum(atoms=atoms, spins=spins, orbitals=orbitals) label = "-" if sum([x in orbitals for x in [0]]) == 1: label += "s" if sum([x in orbitals for x in [1, 2, 3]]) == 3: label += "p" if sum([x in orbitals for x in [4, 5, 6, 7, 8]]) == 5: label += "d" if sum([x in orbitals for x in [9, 10, 11, 12, 13, 14, 15]]) == 7: label += "f" if label == "-" + all_orbitals: label = "" x = self.dos.energies[cond] # x = dos.energies[:] y = (dos[ispin, cond] * dos_total[ispin, cond]) / dos_projected_total[ispin, cond] # y = ( dos.dos[:, ispin + 1] * dos_total.dos[:, ispin + 1] # ) / dos_projected_total.dos[:, ispin + 1] # y = dos.dos[cond, ispin + 1] if ispin > 0 and len(spins) > 1: y *= -1 if orientation == "horizontal": ax.fill_between(x, bottom + y, bottom, color=colors[ispc]) elif orientation == "vertical": ax.fill_betweenx(x, bottom + y, bottom, color=colors[ispc]) else: if orientation == "horizontal": ax.fill_between( x, bottom + y, bottom, color=colors[ispc], label=ispc + label, ) elif orientation == "vertical": ax.fill_betweenx( x, bottom + y, bottom, color=colors[ispc], label=ispc + label, ) bottom += y return fig, ax
[docs]def plotter(energies, dos, spins, spin_colors, figsize, ax, orientation, linewidth, labels=None, ): if spins is None: spins = np.arange(dos.shape[0]) if orientation == "horizontal": if ax is None: fig = plt.figure(figsize=figsize) ax = fig.add_subplot(111) else: fig = ax.get_figure() x = energies for iy in spins: y = dos[iy, :] if iy > 0 and len(spins) > 1: y *= -1 if labels is not None: ax.plot(x, y, "r-", color=spin_colors[iy], label=labels[iy], linewidth=linewidth) else: ax.plot(x, y, "r-", color=spin_colors[iy], linewidth=linewidth) elif orientation == "vertical": if ax is None: fig = plt.figure(figsize=(figsize[1], figsize[0])) ax = fig.add_subplot(111) else: fig = plt.gca() y = energies for ix in spins: x = dos[ix, :] if ix > 0 and len(spins) > 1: x *= -1 if labels is not None: ax.plot(x, y, "r-", color=spin_colors[ix], label=labels[ix], linewidth=linewidth) else: ax.plot(x, y, "r-", color=spin_colors[ix], linewidth=linewidth) return fig, ax