Skip to content

Commit

Permalink
Added dos overlay modes and added method to catch projected normaliza…
Browse files Browse the repository at this point in the history
…tion divide by 0 errors
  • Loading branch information
lllangWV committed Mar 19, 2024
1 parent f29fce2 commit d723ed4
Show file tree
Hide file tree
Showing 2 changed files with 218 additions and 110 deletions.
290 changes: 183 additions & 107 deletions pyprocar/plotter/dos_plot.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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],
Expand Down Expand Up @@ -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],
Expand All @@ -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.
Expand Down Expand Up @@ -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')
Expand All @@ -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,
Expand All @@ -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.
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down
Loading

0 comments on commit d723ed4

Please sign in to comment.