from pymatgen.io.vasp import Vasprun
from pymatgen.electronic_structure.core import Spin
import matplotlib.pyplot as plt
from pymatgen.electronic_structure.plotter import DosPlotter
from matplotlib.ticker import AutoMinorLocator
vasprun = Vasprun("vasprun.xml")
dos = vasprun.complete_dos
energies = dos.energies - dos.efermi
element_dos = dos.get_element_dos()
fig, ax = plt.subplots(figsize=(7, 6))
total_up = gaussian_filter1d(dos.densities[Spin.up], sigma=50)
total_down = -gaussian_filter1d(dos.densities[Spin.down], sigma=50)
ax.plot(energies, total_up, color="black", linewidth=1.5, zorder=1, label="Total")
ax.plot(energies, total_down, color="black", linewidth=1.5, zorder=1)
ax.fill_between(energies, total_up, 0, alpha=0.1, color="gray", zorder=1)
ax.fill_between(energies, total_down, 0, alpha=0.1, color="gray", zorder=1)
for element, dos_obj in element_dos.items():
color = element_colors.get(el, "gray")
y_up = gaussian_filter1d(dos_obj.densities[Spin.up], sigma=50)
y_down = -gaussian_filter1d(dos_obj.densities[Spin.down], sigma=50)
ax.plot(energies, y_up, color=color, linewidth=1.5, label=el, zorder=3)
ax.plot(energies, y_down, color=color, linewidth=1.5, zorder=3)
ax.fill_between(energies, y_up, 0, alpha=0.15, color=color, zorder=3)
ax.fill_between(energies, y_down, 0, alpha=0.15, color=color, zorder=3)
ax.axvline(x=0, color="black", linewidth=1.5, zorder=2)
ax.axhline(y=0, color="black", linewidth=1.0, zorder=1)
ax.set_xlabel(r"$E - E_{\mathrm{f}}$ (eV)", fontsize=24)
ax.set_ylabel(r"DOS (arb. units)", fontsize=24)
lim = max(abs(np.min([total_up, total_down])), np.max([total_up, total_down]))
ax.text(0.02, 0.97, r"spin $\uparrowquot;, transform=ax.transAxes, va="top", fontsize=16) ax.text(0.02, 0.03, r"spin $\downarrowquot;, transform=ax.transAxes, va="bottom", fontsize=16) ax.tick_params(axis="both", which="major", labelsize=16, length=8, width=1.5, direction="inout")
ax.tick_params(axis="both", which="minor", length=4, width=1, direction="inout")
ax.xaxis.set_minor_locator(AutoMinorLocator())
ax.yaxis.set_minor_locator(AutoMinorLocator())
for spine in ax.spines.values():
ax.legend(fontsize=16, framealpha=0.8, edgecolor="gray", bbox_to_anchor=(0.735, .999), loc="upper left")