add an example for plot-tools

Author: jiyuang

tools/plot-tools/README.md

@@ -21,7 +21,7 @@ python setup.py install
 
 ## Usage
 There are two ways to use this tool:
-1. Specify parameters in `band.py` or `dos.py` directly, and then `python band.py` or `python dos.py`. And you can also import module in your own script e.g. `from abacus_plot.band import Band`
+1. Specify parameters in `band.py` or `dos.py` directly, and then `python band.py` or `python dos.py`. And you can also import module in your own script e.g. `from abacus_plot.band import Band`. (Recommend)
 2. Command-line tools are also supported in this tool. In this way, you need prepare an input file and execute some commands (see below). You can use `abacus-plot -h` to check command-line information
 
 ### Band Structure

tools/plot-tools/abacus_plot/band.py

@@ -37,6 +37,10 @@ def __init__(self, bandfile: Union[PathLike, Sequence[PathLike]] = None, kptfile
         if self.kptfile:
             self.kpt = read_kpt(kptfile)
         self.k_index = list(map(int, self.k_index))
+        if self.kpt:
+            self._kzip = self.kpt.label_special_k
+        else:
+            self._kzip = self.k_index
 
     @classmethod
     def read(cls, filename: PathLike):
@@ -168,14 +172,14 @@ def plot_data(cls,
                   ax: axes.Axes,
                   x: Sequence,
                   y: Sequence,
-                  index: Sequence,
+                  kzip: Sequence,
                   efermi: float = 0,
                   energy_range: Sequence[float] = [],
                   **kwargs):
         """Plot band structure
 
         :params x, y: x-axis and y-axis coordinates
-        :params index: special k-points label and its index in data file
+        :params kzip: special k-points label and its k-index in data file
         :params efermi: Fermi level in unit eV
         :params energy_range: range of energy to plot, its length equals to two
         """
@@ -189,7 +193,7 @@ def plot_data(cls,
 
         bandplot.ax.plot(kpoints, energy, lw=bandplot._lw, color=bandplot._color,
                          label=bandplot._label, linestyle=bandplot._linestyle)
-        bandplot._set_figure(index, energy_range)
+        bandplot._set_figure(kzip, energy_range)
 
     def plot(self,
              fig: Figure,
@@ -231,11 +235,7 @@ def plot(self,
             bandplot.ax.plot(self.k_index, band,
                              lw=bandplot._lw, color=bandplot._color, label=bandplot._label, linestyle=bandplot._linestyle)
 
-        if self.kpt:
-            index = self.kpt.label_special_k
-        else:
-            index = self.k_index
-        bandplot._set_figure(index, energy_range)
+        bandplot._set_figure(self._kzip, energy_range)
 
         return bandplot
 
@@ -253,16 +253,16 @@ def __init__(self, fig: Figure, ax: axes.Axes, **kwargs) -> None:
         self._linestyle = kwargs.pop('linestyle', 'solid')
         self.plot_params = kwargs
 
-    def _set_figure(self, index, range: Sequence):
+    def _set_figure(self, kzip, range: Sequence):
         """set figure and axes for plotting
 
-        :params index: dict of label of points of x-axis and its index in data file. Range of x-axis based on index.value()
+        :params kzip: dict of label of points of x-axis and its index in data file. Range of x-axis based on kzip.value()
         :params range: range of y-axis
         """
 
         keys = []
         values = []
-        for t in index:
+        for t in kzip:
             if isinstance(t, tuple):
                 keys.append(t[0])
                 values.append(t[1])
@@ -355,6 +355,10 @@ def __init__(self, bandfile: Union[PathLike, Sequence[PathLike]] = None, kptfile
         if self.kptfile:
             self.kpt = read_kpt(kptfile)
         self.k_index = list(map(int, self.k_index))
+        if self.kpt:
+            self._kzip = self.kpt.label_special_k
+        else:
+            self._kzip = self.k_index
 
     def _check_energy(self, energy):
         assert energy.shape[0] == self.nkpoints, "The dimension of band structure dismatches with the number of k-points."
@@ -421,7 +425,7 @@ def read(cls, filename: PathLike):
 
         return nspin, norbitals, eunit, nbands, nkpoints, k_index, energy, orbitals
 
-    def _write(self, species: Union[Sequence[Any], Dict[Any, List[int]], Dict[Any, Dict[int, List[int]]]], keyname='', file_dir:PathLike=''):
+    def _write(self, species: Union[Sequence[Any], Dict[Any, List[int]], Dict[Any, Dict[int, List[int]]]], keyname='', file_dir: PathLike = ''):
         """Write parsed projected bands data to files
 
         Args:
@@ -581,10 +585,11 @@ def _plot(self,
             BandPlot object: for manually plotting picture with bandplot.ax 
         """
 
-        def _seg_plot(bandplot, lc, index, file_dir, name):
+        def _seg_plot(bandplot, lc, file_dir, name):
             cbar = bandplot.fig.colorbar(lc, ax=bandplot.ax)
-            bandplot._set_figure(index, energy_range)
-            bandplot.fig.savefig(file_dir/f'{keyname}-{bandplot._label}.pdf', dpi=400)
+            bandplot._set_figure(self._kzip, energy_range)
+            bandplot.fig.savefig(
+                file_dir/f'{keyname}-{bandplot._label}.pdf', dpi=400)
             cbar.remove()
             plt.cla()
 
@@ -595,15 +600,10 @@ def _seg_plot(bandplot, lc, index, file_dir, name):
         file_dir = Path(f"{outdir}", f"PBAND{out_index}_FIG")
         file_dir.mkdir(exist_ok=True)
 
-        if self.kpt:
-            index = self.kpt.label_special_k
-        else:
-            index = self.k_index
-
         if not species:
             bandplot = BandPlot(fig, ax, **kwargs)
             bandplot = super().plot(fig, ax, efermi, energy_range, shift, **kwargs)
-            bandplot._set_figure(index, energy_range)
+            bandplot._set_figure(self._kzip, energy_range)
 
             return bandplot
 
@@ -613,15 +613,19 @@ def _seg_plot(bandplot, lc, index, file_dir, name):
                 bandplot = BandPlot(fig, ax, **kwargs)
                 bandplot._label = elem
                 for ib in range(self.nbands):
-                    points = np.array((self.k_index, energy[0:, ib])).T.reshape(-1, 1, 2)
-                    segments = np.concatenate([points[:-1], points[1:]], axis=1)
-                    norm = Normalize(vmin=wei[elem][0:, ib].min(), vmax=wei[elem][0:, ib].max())
-                    lc = LineCollection(segments, cmap=plt.get_cmap(cmap), norm=norm)
+                    points = np.array(
+                        (self.k_index, energy[0:, ib])).T.reshape(-1, 1, 2)
+                    segments = np.concatenate(
+                        [points[:-1], points[1:]], axis=1)
+                    norm = Normalize(
+                        vmin=wei[elem][0:, ib].min(), vmax=wei[elem][0:, ib].max())
+                    lc = LineCollection(
+                        segments, cmap=plt.get_cmap(cmap), norm=norm)
                     lc.set_array(wei[elem][0:, ib])
                     lc.set_label(bandplot._label)
                     bandplot.ax.add_collection(lc)
-                
-                _seg_plot(bandplot, lc, index, file_dir, name=f'{elem}')
+
+                _seg_plot(bandplot, lc, file_dir, name=f'{elem}')
                 bandplots.append(bandplot)
             return bandplots
 
@@ -638,30 +642,40 @@ def _seg_plot(bandplot, lc, index, file_dir, name):
                             m_index = int(mag)
                             bandplot._label = f"{elem}-{get_angular_momentum_name(l_index, m_index)}"
                             for ib in range(self.nbands):
-                                points = np.array((self.k_index, energy[0:, ib])).T.reshape(-1, 1, 2)
-                                segments = np.concatenate([points[:-1], points[1:]], axis=1)
-                                norm = Normalize(vmin=wei[elem][ang][mag][0:, ib].min(), vmax=wei[elem][ang][mag][0:, ib].max())
-                                lc = LineCollection(segments, cmap=plt.get_cmap(cmap), norm=norm)
+                                points = np.array(
+                                    (self.k_index, energy[0:, ib])).T.reshape(-1, 1, 2)
+                                segments = np.concatenate(
+                                    [points[:-1], points[1:]], axis=1)
+                                norm = Normalize(vmin=wei[elem][ang][mag][0:, ib].min(
+                                ), vmax=wei[elem][ang][mag][0:, ib].max())
+                                lc = LineCollection(
+                                    segments, cmap=plt.get_cmap(cmap), norm=norm)
                                 lc.set_array(wei[elem][ang][mag][0:, ib])
                                 lc.set_label(bandplot._label)
                                 bandplot.ax.add_collection(lc)
-                
-                            _seg_plot(bandplot, lc, index, elem_file_dir, name=f'{elem}_{ang}_{mag}')
+
+                            _seg_plot(bandplot, lc, elem_file_dir,
+                                      name=f'{elem}_{ang}_{mag}')
                             bandplots.append(bandplot)
 
                     else:
                         bandplot = BandPlot(fig, ax, **kwargs)
                         bandplot._label = f"{elem}-{get_angular_momentum_label(l_index)}"
                         for ib in range(self.nbands):
-                            points = np.array((self.k_index, energy[0:, ib])).T.reshape(-1, 1, 2)
-                            segments = np.concatenate([points[:-1], points[1:]], axis=1)
-                            norm = Normalize(vmin=wei[elem][ang][0:, ib].min(), vmax=wei[elem][ang][0:, ib].max())
-                            lc = LineCollection(segments, cmap=plt.get_cmap(cmap), norm=norm)
+                            points = np.array(
+                                (self.k_index, energy[0:, ib])).T.reshape(-1, 1, 2)
+                            segments = np.concatenate(
+                                [points[:-1], points[1:]], axis=1)
+                            norm = Normalize(vmin=wei[elem][ang][0:, ib].min(
+                            ), vmax=wei[elem][ang][0:, ib].max())
+                            lc = LineCollection(
+                                segments, cmap=plt.get_cmap(cmap), norm=norm)
                             lc.set_array(wei[elem][ang][0:, ib])
                             lc.set_label(bandplot._label)
                             bandplot.ax.add_collection(lc)
-                
-                        _seg_plot(bandplot, lc, index, elem_file_dir, name=f'{elem}_{ang}')
+
+                        _seg_plot(bandplot, lc, elem_file_dir,
+                                  name=f'{elem}_{ang}')
                         bandplots.append(bandplot)
 
             return bandplots
@@ -742,19 +756,19 @@ def plot(self,
 if __name__ == "__main__":
     import matplotlib.pyplot as plt
     from pathlib import Path
-    parent = Path(r"C:\Users\YY.Ji\Desktop")
+    parent = Path(r"C:\Users\YY.Ji\Desktop\Si")
     name = "PBANDS_1"
     path = parent/name
     fig, ax = plt.subplots(figsize=(12, 6))
-    energy_range = [-5, 6]
-    efermi = 4.417301755850272
+    energy_range = [-5, 7]
+    efermi = 6.585653952007503
     shift = False
     #species = {"Ag": [2], "Cl": [1], "In": [0]}
-    atom_index = {8: {2: [1, 2]}, 4: {2: [1, 2]}, 10: [1, 2]}
+    atom_index = {1: {1: [0, 1]}}
     pband = PBand(str(path))
 
-    # if you want to specify `species` or `index`, you need to 
-    # set `species=species` or `index=index` in the following two functions 
+    # if you want to specify `species` or `index`, you need to
+    # set `species=species` or `index=index` in the following two functions
     pband.plot(fig, ax, atom_index=atom_index, efermi=efermi,
-               energy_range=energy_range, shift=shift)    
+               energy_range=energy_range, shift=shift)
     pband.write(atom_index=atom_index)

tools/plot-tools/abacus_plot/dos.py

@@ -154,7 +154,7 @@ def _set_figure(self, energy_range: Sequence = [], dos_range: Sequence = [], not
 class TDOS(DOS):
     """Parse total DOS data"""
 
-    def __init__(self, tdosfile: PathLike=None) -> None:
+    def __init__(self, tdosfile: PathLike = None) -> None:
         super().__init__()
         self.tdosfile = tdosfile
         self._read()
@@ -202,7 +202,7 @@ def plot(self, fig: Figure, ax: Union[axes.Axes, Sequence[axes.Axes]], efermi: f
 class PDOS(DOS):
     """Parse partial DOS data"""
 
-    def __init__(self, pdosfile: PathLike=None) -> None:
+    def __init__(self, pdosfile: PathLike = None) -> None:
         super().__init__()
         self.pdosfile = pdosfile
         self._read()
@@ -391,6 +391,8 @@ def _parial_plot(self,
         Returns:
             DOSPlot object: for manually plotting picture with dosplot.ax 
         """
+        if not isinstance(ax, list):
+            ax = [ax]
 
         dos, totnum = parse_projected_data(self.orbitals, species, keyname)
         energy_f, tdos = self._shift_energy(efermi, shift, prec)
@@ -403,7 +405,7 @@ def _parial_plot(self,
                                     notes=dosplot.plot_params["notes"])
             else:
                 dosplot._set_figure(energy_range, dos_range)
-        
+
             return dosplot
 
         if isinstance(species, (list, tuple)):
@@ -515,17 +517,18 @@ def plot(self,
     #                     energy_range=energy_range, dos_range=dos_range, notes={'s': '(a)'})
     # fig.savefig("tdos.png")
 
-    pdosfile = r"C:\Users\YY.Ji\Desktop\PDOS"
+    pdosfile = r"C:\Users\YY.Ji\Desktop\Si\PDOS"
     pdos = PDOS(pdosfile)
     #species = {"Ag": [2], "Cl": [1], "In": [0]}
-    atom_index = {8: {2: [1, 2]}, 4: {2: [1, 2]}, 10: [1, 2]}
-    fig, ax = plt.subplots(3, 1, sharex=True)
-    energy_range = [-1.5, 6]
+    atom_index = {1: {1: [0, 1]}}
+    fig, ax = plt.subplots(1, 1, sharex=True)
+    energy_range = [-5, 7]
+    efermi = 6.585653952007503
     dos_range = [0, 5]
 
-    # if you want to specify `species` or `index`, you need to 
-    # set `species=species` or `index=index` in the following two functions 
-    dosplots = pdos.plot(fig, ax, atom_index=atom_index, efermi=5, shift=True,
-                         energy_range=energy_range, dos_range=dos_range, notes=[{'s': '(a)'}, {'s': '(b)'}, {'s': '(c)'}])
+    # if you want to specify `species` or `index`, you need to
+    # set `species=species` or `index=index` in the following two functions
+    dosplots = pdos.plot(fig, ax, atom_index=atom_index, efermi=efermi, shift=True,
+                         energy_range=energy_range, dos_range=dos_range, notes=[{'s': '(a)'}])
     fig.savefig("pdos.png")
-    pdos.write(atom_index=atom_index)
+    pdos.write(atom_index=atom_index)

tools/plot-tools/abacus_plot/main.py

@@ -10,7 +10,7 @@
 
 from abacus_plot.band import Band, PBand
 from abacus_plot.dos import TDOS, PDOS
-from abacus_plot.utils import read_json
+from abacus_plot.utils import read_json, key2int
 
 
 class Show:
@@ -41,9 +41,9 @@ def show_cmdline(cls, args):
             shift = text.pop("shift", False)
             figsize = text.pop("figsize", (12, 10))
             fig, ax = plt.subplots(figsize=figsize)
-            index = text.pop("index", [])
-            atom_index = text.pop("atom_index", [])
-            species = text.pop("species", [])
+            index = key2int(text.pop("index", None))    # keys of json must be dict, I convert them to int for `index` and `atom_index`
+            atom_index = key2int(text.pop("atom_index", None))
+            species = text.pop("species", None)
             outdir = text.pop("outdir", './')
             cmapname = text.pop("cmapname", 'jet')
             pband = PBand(bandfile, kptfile)
@@ -53,9 +53,9 @@ def show_cmdline(cls, args):
             text = read_json(args.file)
             bandfile = text["bandfile"]
             kptfile = text["kptfile"]
-            index = text.pop("index", [])
-            atom_index = text.pop("atom_index", [])
-            species = text.pop("species", [])
+            index = key2int(text.pop("index", None))
+            atom_index = key2int(text.pop("atom_index", None))
+            species = text.pop("species", None)
             outdir = text.pop("outdir", './')
             pdos = PBand(bandfile, kptfile)
             pdos.write(index=index, atom_index=atom_index, species=species, outdir=outdir)
@@ -83,12 +83,19 @@ def show_cmdline(cls, args):
             energy_range = text.pop("energy_range", [])
             dos_range = text.pop("dos_range", [])
             shift = text.pop("shift", False)
-            index = text.pop("index", [])
-            atom_index = text.pop("atom_index", [])
-            species = text.pop("species", [])
+            index = key2int(text.pop("index", None))
+            atom_index = key2int(text.pop("atom_index", None))
+            species = text.pop("species", None)
             figsize = text.pop("figsize", (12, 10))
-            fig, ax = plt.subplots(
-                len(species), 1, sharex=True, figsize=figsize)
+            if index:
+                fig, ax = plt.subplots(
+                    len(index), 1, sharex=True, figsize=figsize)
+            if atom_index:
+                fig, ax = plt.subplots(
+                    len(atom_index), 1, sharex=True, figsize=figsize)
+            if species:
+                fig, ax = plt.subplots(
+                    len(species), 1, sharex=True, figsize=figsize)
             prec = text.pop("prec", 0.01)
             pdos = PDOS(pdosfile)
             pdos.plot(fig=fig, ax=ax, index=index, atom_index=atom_index, species=species, efermi=efermi, shift=shift,
@@ -99,9 +106,9 @@ def show_cmdline(cls, args):
         if args.dos and args.out:
             text = read_json(args.file)
             pdosfile = text.pop("pdosfile", '')
-            index = text.pop("index", [])
-            atom_index = text.pop("atom_index", [])
-            species = text.pop("species", [])
+            index = key2int(text.pop("index", None))
+            atom_index = key2int(text.pop("atom_index", None))
+            species = text.pop("species", None)
             outdir = text.pop("outdir", './')
             pdos = PDOS(pdosfile)
             pdos.write(index=index, atom_index=atom_index, species=species, outdir=outdir)