diff --git a/pyprocar/plotter/dos_plot.py b/pyprocar/plotter/dos_plot.py index 5431fef8..b82aeaaa 100644 --- a/pyprocar/plotter/dos_plot.py +++ b/pyprocar/plotter/dos_plot.py @@ -400,6 +400,7 @@ def plot_stack_species( principal_q_numbers:List[int]=[-1], orbitals:List[int]=None, spins:List[int]=None, + overlay_mode:bool=False, orientation:str="horizontal", ): """A method to plot the dos with the species contribution stacked on eachother @@ -498,27 +499,37 @@ def plot_stack_species( x = self.dos.energies y = (dos_projected[ispin] / dos_projected_total[ispin] ) * dos_total[ispin] - - if ispin > 0 and len(spins) > 1: - y *= -1 - handle = self.ax.fill_between(x, - bottom + y, - bottom, - color=self.config['colors']['value'][specie], - ) + # replace very large values above a threshold with 0. + # These are artifacts from the division of samll total values. + y = np.nan_to_num(y, 0) + threshold=50 + y[np.abs(y) > threshold] = 0 + if overlay_mode: + if ispin > 0 and len(spins) > 1: + y *= -1 + handle, = self.ax.plot(x,y,color=self.config['colors']['value'][specie]) + else: + handle, = self.ax.plot(x,y,color=self.config['colors']['value'][specie]) else: - handle = self.ax.fill_between( - x, - bottom + y, - bottom, - color=self.config['colors']['value'][specie], - ) + if ispin > 0 and len(spins) > 1: + y *= -1 + handle = self.ax.fill_between(x, + bottom + y, + bottom, + color=self.config['colors']['value'][specie], + ) + else: + handle = self.ax.fill_between( + x, + bottom + y, + bottom, + color=self.config['colors']['value'][specie], + ) self.handles.append(handle) - label=self.structure.species[specie] + label + self.config['spin_labels']['value'][ispin] - self.labels.append(label) + self.labels.append(self.structure.species[specie] + label + self.config['spin_labels']['value'][ispin]) bottom += y - if spins_index == 0: + if ispin == 0: self.ax.plot( self.dos.energies, self.dos.total[ispin, :], color= 'black', alpha=self.config['opacity']['value'][ispin], @@ -562,26 +573,37 @@ def plot_stack_species( x = self.dos.energies y = (dos[ispin] * dos_total[ispin]) / dos_projected_total[ispin] - if ispin > 0 and len(spins) > 1: - y *= -1 - handle = self.ax.fill_betweenx(x, - bottom + y, - bottom, - color=self.config['colors']['value'][specie], - ) + # replace very large values above a threshold with 0. + # These are artifacts from the division of samll total values. + y = np.nan_to_num(y, 0) + threshold=50 + y[np.abs(y) > threshold] = 0 + if overlay_mode: + if ispin > 0 and len(spins) > 1: + y *= -1 + handle, = self.ax.plot(y,x,color=self.config['colors']['value'][specie]) + else: + handle, = self.ax.plot(y,x,color=self.config['colors']['value'][specie]) else: - handle = self.ax.fill_betweenx(x, - bottom + y, - bottom, - color=self.config['colors']['value'][specie], - ) + if ispin > 0 and len(spins) > 1: + y *= -1 + handle = self.ax.fill_betweenx(x, + bottom + y, + bottom, + color=self.config['colors']['value'][specie], + ) + else: + handle = self.ax.fill_betweenx(x, + bottom + y, + bottom, + color=self.config['colors']['value'][specie], + ) self.handles.append(handle) - label=self.structure.species[specie] + label + self.config['spin_labels']['value'][ispin] - self.labels.append(label) + self.labels.append(self.structure.species[specie] + label + self.config['spin_labels']['value'][ispin]) bottom += y if self.config['plot_total']['value'] == True: - if spins_index == 0: + if ispin == 0: self.ax.plot( self.dos.total[ispin, :], self.dos.energies, color= 'black', alpha=self.config['opacity']['value'][ispin], @@ -602,6 +624,7 @@ def plot_stack_orbitals(self, atoms:List[int]=None, spins:List[int]=None, principal_q_numbers:List[int]=[-1], + overlay_mode:bool=False, orientation:str="horizontal", ): """A method to plot dos orbitals contribution stacked. @@ -678,45 +701,57 @@ def plot_stack_orbitals(self, x = self.dos.energies y = (dos[ispin] * dos_total[ispin]) / dos_projected_total[ispin] + # replace very large values above a threshold with 0. + # These are artifacts from the division of samll total values. y = np.nan_to_num(y, 0) + threshold=50 + y[np.abs(y) > threshold] = 0 + + if overlay_mode: + if ispin > 0 and len(spins) > 1: + y *= -1 + handle, = self.ax.plot(x,y,color=self.config['colors']['value'][iorb]) + else: + handle, = self.ax.plot(x,y,color=self.config['colors']['value'][iorb],) - if ispin > 0 and len(spins) > 1: - y *= -1 - handle = self.ax.fill_between(x, - bottom + y, - bottom, - color=self.config['colors']['value'][iorb]) - else: - handle = self.ax.fill_between( - x, - bottom + y, - bottom, - color=self.config['colors']['value'][iorb], - ) + if ispin > 0 and len(spins) > 1: + y *= -1 + handle = self.ax.fill_between(x, + bottom + y, + bottom, + color=self.config['colors']['value'][iorb]) + + else: + handle = self.ax.fill_between( + x, + bottom + y, + bottom, + color=self.config['colors']['value'][iorb], + ) self.labels.append(atom_names + orb_names[iorb] + self.config['spin_labels']['value'][ispin]) self.handles.append(handle) bottom += y - if self.config['plot_total']['value'] == True: - if ispin == 0: - self.ax.plot( - self.dos.energies, self.dos.total[ispin, :], color= 'black', - alpha=self.config['opacity']['value'][ispin], - linestyle=self.config['linestyle']['value'][ispin], - label=self.config['spin_labels']['value'][ispin], - linewidth=self.config['linewidth']['value'][ispin], - ) - else: - self.ax.plot( - self.dos.energies, -self.dos.total[ispin, :], color= 'black', - alpha=self.config['opacity']['value'][ispin], - linestyle=self.config['linestyle']['value'][ispin], - label=self.config['spin_labels']['value'][ispin], - linewidth=self.config['linewidth']['value'][ispin], - ) + if self.config['plot_total']['value'] == True: + if ispin == 0: + self.ax.plot( + self.dos.energies, self.dos.total[ispin, :], color= 'black', + alpha=self.config['opacity']['value'][ispin], + linestyle=self.config['linestyle']['value'][ispin], + label=self.config['spin_labels']['value'][ispin], + linewidth=self.config['linewidth']['value'][ispin], + ) + else: + self.ax.plot( + self.dos.energies, -self.dos.total[ispin, :], color= 'black', + alpha=self.config['opacity']['value'][ispin], + linestyle=self.config['linestyle']['value'][ispin], + label=self.config['spin_labels']['value'][ispin], + linewidth=self.config['linewidth']['value'][ispin], + ) elif orientation == 'vertical': self.set_xlabel('DOS Cumlative') @@ -739,8 +774,18 @@ def plot_stack_orbitals(self, x = self.dos.energies y = (dos[ispin] * dos_total[ispin]) / dos_projected_total[ispin] + # replace very large values above a threshold with 0. + # These are artifacts from the division of samll total values. y = np.nan_to_num(y, 0) - + threshold=50 + y[np.abs(y) > threshold] = 0 + + if overlay_mode: + if ispin > 0 and len(spins) > 1: + y *= -1 + handle, = self.ax.plot(y,x,color=self.config['colors']['value'][iorb]) + else: + handle, = self.ax.plot(y,x,color=self.config['colors']['value'][iorb]) if ispin > 0 and len(spins) > 1: y *= -1 handle = self.ax.fill_betweenx(x, @@ -760,28 +805,29 @@ def plot_stack_orbitals(self, self.labels.append(atom_names + orb_names[iorb] + self.config['spin_labels']['value'][ispin]) self.handles.append(handle) bottom += y - if self.config['plot_total']['value'] == True: - if ispin == 0: - self.ax.plot( - self.dos.total[ispin, :], self.dos.energies, color= 'black', - alpha=self.config['opacity']['value'][ispin], - linestyle=self.config['linestyle']['value'][ispin], - label=self.config['spin_labels']['value'][ispin], - linewidth=self.config['linewidth']['value'][ispin], - ) - else: - self.ax.plot( - -self.dos.total[ispin, :], self.dos.energies, color= 'black', - alpha=self.config['opacity']['value'][ispin], - linestyle=self.config['linestyle']['value'][ispin], - label=self.config['spin_labels']['value'][ispin], - linewidth=self.config['linewidth']['value'][ispin], - ) + if self.config['plot_total']['value'] == True: + if ispin == 0: + self.ax.plot( + self.dos.total[ispin, :], self.dos.energies, color= 'black', + alpha=self.config['opacity']['value'][ispin], + linestyle=self.config['linestyle']['value'][ispin], + label=self.config['spin_labels']['value'][ispin], + linewidth=self.config['linewidth']['value'][ispin], + ) + else: + self.ax.plot( + -self.dos.total[ispin, :], self.dos.energies, color= 'black', + alpha=self.config['opacity']['value'][ispin], + linestyle=self.config['linestyle']['value'][ispin], + label=self.config['spin_labels']['value'][ispin], + linewidth=self.config['linewidth']['value'][ispin], + ) def plot_stack(self, items:dict=None, spins:List[int]=None, plot_total:bool= True, + overlay_mode=False, orientation:str='horizontal', ): """A method to plot dos contributions stacked. @@ -888,23 +934,40 @@ def plot_stack(self, label = "" x = self.dos.energies + + y = (dos[ispin] * dos_total[ispin]) / dos_projected_total[ispin] - if ispin > 0 and len(spins) > 1: - y *= -1 - handle = self.ax.fill_between(x, - bottom + y, - bottom, - color=colors_dict[specie]) - else: - handle = self.ax.fill_between( - x, - bottom + y, - bottom, - color=colors_dict[specie], - ) - self.handles.append(handle) + # replace very large values above a threshold with 0. + # These are artifacts from the division of samll total values. + y = np.nan_to_num(y, 0) + threshold=50 + y[np.abs(y) > threshold] = 0 + + if overlay_mode: + if ispin > 0 and len(spins) > 1: + y *= -1 + handle, = self.ax.plot(x,y,color=colors_dict[specie]) + else: + handle, = self.ax.plot(x,y,color=colors_dict[specie]) + else: + if ispin > 0 and len(spins) > 1: + y *= -1 + handle = self.ax.fill_between(x, + bottom + y, + bottom, + color=colors_dict[specie]) + else: + handle = self.ax.fill_between( + x, + bottom + y, + bottom, + color=colors_dict[specie], + ) + self.labels.append(specie + label + self.config['spin_labels']['value'][ispin]) + self.handles.append(handle) + bottom += y if plot_total == True: if ispin == 0: @@ -978,20 +1041,31 @@ def plot_stack(self, x = self.dos.energies y = (dos[ispin] * dos_total[ispin]) / dos_projected_total[ispin] - - if ispin > 0 and len(spins) > 1: - y *= -1 - handle = self.ax.fill_betweenx(x, - bottom + y, - bottom, - color=colors_dict[specie]) + # replace very large values above a threshold with 0. + # These are artifacts from the division of samll total values. + y = np.nan_to_num(y, 0) + threshold=50 + y[np.abs(y) > threshold] = 0 + if overlay_mode: + if ispin > 0 and len(spins) > 1: + y *= -1 + handle, = self.ax.plot(y,x,color=colors_dict[specie]) + else: + handle, = self.ax.plot(y,x,color=colors_dict[specie]) else: - handle = self.ax.fill_betweenx( - x, - bottom + y, - bottom, - color=self.config['colors']['value'][specie], - ) + if ispin > 0 and len(spins) > 1: + y *= -1 + handle = self.ax.fill_betweenx(x, + bottom + y, + bottom, + color=colors_dict[specie]) + else: + handle = self.ax.fill_betweenx( + x, + bottom + y, + bottom, + color=self.config['colors']['value'][specie], + ) self.handles.append(handle) self.labels.append(specie + label + self.config['spin_labels']['value'][ispin]) bottom += y @@ -1128,6 +1202,8 @@ def legend(self, if labels == None: labels = self.labels if self.config['legend']['value'] and len(labels) != 0: + if len(self.handles) != len(labels): + raise ValueError(f"The number of labels and handles should be the same, currently there are {len(self.handles)} handles and {len(labels)} labels") self.ax.legend(self.handles, labels) return None diff --git a/pyprocar/scripts/scriptDosplot.py b/pyprocar/scripts/scriptDosplot.py index 1530528e..4f0ed62d 100644 --- a/pyprocar/scripts/scriptDosplot.py +++ b/pyprocar/scripts/scriptDosplot.py @@ -33,6 +33,7 @@ def dosplot( elimit:List[float]=None, dos_limit:List[float]=None, savefig:str=None, + labels:List[str]=None, projection_mask=None, ax:plt.Axes=None, show:bool=True, @@ -158,6 +159,12 @@ def dosplot( e.g. ``dos_limit=[0, 30]`` + labels : list str, optional + ``labels`` is a list of strings that will be used as the + legend of the plot. The length of the list should be equal to + the number of curves being plotted. If not provided the + default labels will be used. + savefig : str , optional (default None) ``savefig`` defines the file that the plot is going to be saved in. ``savefig`` accepts all the formats accepted by @@ -329,20 +336,39 @@ def dosplot( orbitals=orbitals, orientation=orientation, ) - elif mode == "stack_orbitals": edos_plot.plot_stack_orbitals( spins=spins, atoms=atoms, orientation=orientation, ) - elif mode == "stack": edos_plot.plot_stack( spins=spins, items=items, orientation=orientation, ) + elif mode == "overlay_species": + edos_plot.plot_stack_species( + spins=spins, + orbitals=orbitals, + orientation=orientation, + overlay_mode=True + ) + elif mode == "overlay_orbitals": + edos_plot.plot_stack_orbitals( + spins=spins, + atoms=atoms, + orientation=orientation, + overlay_mode=True + ) + elif mode == "overlay": + edos_plot.plot_stack( + spins=spins, + items=items, + orientation=orientation, + overlay_mode=True + ) else: raise ValueError("The mode needs to be in the List [plain,parametric,parametric_line,stack_species,stack_orbitals,stack]") @@ -360,7 +386,13 @@ def dosplot( edos_plot.set_xlim(dos_limit) edos_plot.grid() - edos_plot.legend(edos_plot.labels) + + if labels: + labels=labels + else: + labels=edos_plot.labels + edos_plot.legend(labels) + if savefig is not None: edos_plot.save(savefig) if show: