diff --git a/tools/plot-tools/abacus_plot/dos.py b/tools/plot-tools/abacus_plot/dos.py index 394cf23e6fd..816afb6ccb2 100644 --- a/tools/plot-tools/abacus_plot/dos.py +++ b/tools/plot-tools/abacus_plot/dos.py @@ -20,8 +20,8 @@ class DOS: """Parse DOS data""" - def __init__(self) -> None: - self.nspin = 1 + def __init__(self, nspin) -> None: + self.nspin = nspin if self.nspin in [1, 4]: self._nsplit = 1 elif self.nspin == 2: @@ -89,10 +89,11 @@ def bandgap(cls, vb: namedtuple, cb: namedtuple): return gap -class DOSPlot: +class : """Plot density of state(DOS)""" - def __init__(self, fig: Figure = None, ax: axes.Axes = None, **kwargs) -> None: + def __init__(self, fig: Figure = None, ax: axes.Axes = None, nspin: int = 1, **kwargs) -> None: + self.nspin = nspin self.fig = fig self.ax = ax self._lw = kwargs.pop('lw', 2) @@ -145,19 +146,26 @@ def _set_figure(self, energy_range: Sequence = [], dos_range: Sequence = [], not else: self.ax.axvline(0, linestyle="--", c='b', lw=1.0) + if self.nspin == 2: + self.ax.axhline(0, linestyle="--", c='gray', lw=1.0) + + handles, labels = self.ax.get_legend_handles_labels() + by_label = OrderedDict(zip(labels, handles)) if "legend_prop" in self.plot_params.keys(): - self.ax.legend(prop=self.plot_params["legend_prop"]) + self.ax.legend(by_label.values(), by_label.keys(), + prop=self.plot_params["legend_prop"]) else: - self.ax.legend(prop={'size': 15}) + self.ax.legend(by_label.values(), + by_label.keys(), prop={'size': 15}) class TDOS(DOS): """Parse total DOS data""" def __init__(self, tdosfile: PathLike = None) -> None: - super().__init__() self.tdosfile = tdosfile self._read() + super().__init__(self.nspin) def _read(self) -> tuple: """Read total DOS data file @@ -171,7 +179,7 @@ def _read(self) -> tuple: self.energy, self.dos = np.split(data, self.nspin+1, axis=1) elif self.nspin == 2: self.energy, dos_up, dos_dw = np.split(data, self.nspin+1, axis=1) - self.dos = np.hstack(dos_up, dos_dw) + self.dos = np.hstack((dos_up, dos_dw)) def _shift_energy(self, efermi: float = 0, shift: bool = False, prec: float = 0.01): if shift: @@ -188,7 +196,7 @@ def plot(self, fig: Figure, ax: Union[axes.Axes, Sequence[axes.Axes]], efermi: f energy_f = self._shift_energy(efermi, shift, prec) - dosplot = DOSPlot(fig, ax, **kwargs) + dosplot = DOSPlot(fig, ax, self.nspin, **kwargs) dosplot.ax = self._plot(dosplot, energy_f, self.dos, "TDOS") if "notes" in dosplot.plot_params.keys(): dosplot._set_figure(energy_range, dos_range, @@ -203,9 +211,9 @@ class PDOS(DOS): """Parse partial DOS data""" def __init__(self, pdosfile: PathLike = None) -> None: - super().__init__() self.pdosfile = pdosfile self._read() + super().__init__(self.nspin) def _read(self): """Read partial DOS data file @@ -398,7 +406,7 @@ def _parial_plot(self, energy_f, tdos = self._shift_energy(efermi, shift, prec) if not species: - dosplot = DOSPlot(fig, ax, **kwargs) + dosplot = DOSPlot(fig, ax, self.nspin, **kwargs) dosplot.ax = self._plot(dosplot, energy_f, tdos, "TDOS") if "notes" in dosplot.plot_params.keys(): dosplot._set_figure(energy_range, dos_range, @@ -409,7 +417,7 @@ def _parial_plot(self, return dosplot if isinstance(species, (list, tuple)): - dosplot = DOSPlot(fig, ax, **kwargs) + dosplot = DOSPlot(fig, ax, self.nspin, **kwargs) if "xlabel_params" in dosplot.plot_params.keys(): dosplot.ax.set_xlabel("Energy(eV)", ** dosplot.plot_params["xlabel_params"]) @@ -431,7 +439,7 @@ def _parial_plot(self, assert len(ax) >= len( dos.keys()), "There must be enough `axes` to plot." for i, elem in enumerate(dos.keys()): - dosplot = DOSPlot(fig, ax[i], **kwargs) + dosplot = DOSPlot(fig, ax[i], self.nspin, **kwargs) for ang in dos[elem].keys(): l_index = int(ang) if isinstance(dos[elem][ang], dict):