Fix: dos plot for nspin=2

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):