modify README for plot-tools

Author: jiyuang

tools/plot-tools/README.md

@@ -20,14 +20,14 @@ 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 BandPlot`
+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`
 2. Command-line tools are also supported in this tool. In this way, you need prepare an input file and execute some commands as follows
 
 ### Band Structure
 First, prepare a json file e.g. band-input.json:
 ```json
 {
-	"filename" : "BANDS_1.dat",
+	"bandfile" : "BANDS_1.dat",
 	"efermi" : 10.39537843955395,
 	"energy_range"	: [-1.5, 6],
 	"kptfile" : "KPT"
@@ -39,11 +39,14 @@ First, prepare a json file e.g. band-input.json:
 |    *efermi*    | `float` or `List[float]` |                       Fermi level in eV                        |
 | *energy_range* |          `list`          |                     Range of energy in eV                      |
 |    *shift*     |          `bool`          | If set `'true'`, it will evaluate band gap. Default: `'false'` |
+|   *index*    | `List[int]` or `Dict[int, List[int]]` or `Dict[int, Dict[int, List[int]]]]` | Extract PDOS of each atom e.g. [0, 1], {0:[0, 1], 1:[0]}, [0:{0:[0]}, 1:{1:[0, 1]}] |
+|   *atom_index*    | `List[int]` or `Dict[int, List[int]]` or `Dict[int, Dict[int, List[int]]]]` | Extract PDOS of each atom with same atom_index e.g. [0, 1], {0:[0, 1], 1:[0]}, [0:{0:[0]}, 1:{1:[0, 1]}] |
+|   *species*    | `List[str]` or `Dict[str, List[int]]` or `Dict[str, Dict[int, List[int]]]]` | Extract PDOS of each atom with same species e.g. ["C", "Si"], {"C":[0, 1], "Si":[0]}, ["C":{0:[0]}, "Si":{1:[0, 1]}] |
 |   *kptfile*    |          `str`           |                 K-points file in ABACUS format                 |
 |    *label*     |   `str` or `List[str]`   |                    Label of band structure                     |
 |    *color*     |   `str` or `List[str]`   |                    Color of band structure                     |
 
-If you want to plot band structure of `nspin=2` or compare two band structure on same k-path, you can set *filename*, *efermi*, *label*, *color* in `list` type. 
+If you want to plot band structure of `nspin=2` or compare two band structure on same k-path, you can set *bandfile*, *efermi*, *label*, *color* in `list` type. 
 
 The *kptfile* should be as follows, and notes after `#` will be set as k-points label automatically.
 ```shell
@@ -60,7 +63,7 @@ Line
 ```
 Then, use the following command to plot band structure:
 ```shell
-abacus-plot -b band-input.json
+abacus-plot -b -f band-input.json
 ```
 
 ### DOS

tools/plot-tools/abacus_plot/band.py

@@ -22,7 +22,7 @@
 class Band:
     """Parse Bands data"""
 
-    def __init__(self, bandfile: Union[PathLike, Sequence[PathLike]] = None, kptfile: str = '') -> None:
+    def __init__(self, bandfile: Union[PathLike, Sequence[PathLike]] = None, kptfile: PathLike = '') -> None:
         self.bandfile = bandfile
         if isinstance(bandfile, list) or isinstance(bandfile, tuple):
             self.energy = []
@@ -683,17 +683,19 @@ def plot(self,
              outdir: PathLike = './',
              cmapname='jet',
              **kwargs):
-        """Plot parsed partial dos data
+        """Plot parsed projected band data
 
         Args:
             fig (Figure): object of matplotlib.figure.Figure
             ax (Union[axes.Axes, Sequence[axes.Axes]]): object of matplotlib.axes.Axes or a list of this objects
-            index (Union[Sequence[int], Dict[int, List[int]], Dict[int, Dict[int, List[int]]]], optional): extract PDOS of each atom. Defaults to [].
-            atom_index (Union[Sequence[int], Dict[int, List[int]], Dict[int, Dict[int, List[int]]]], optional): extract PDOS of each atom with same atom_index. Defaults to [].
-            species (Union[Sequence[str], Dict[str, List[int]], Dict[str, Dict[int, List[int]]]], optional): extract PDOS of each atom with same species. Defaults to [].
+            index (Union[Sequence[int], Dict[int, List[int]], Dict[int, Dict[int, List[int]]]], optional): extract PBAND of each atom. Defaults to [].
+            atom_index (Union[Sequence[int], Dict[int, List[int]], Dict[int, Dict[int, List[int]]]], optional): extract PBAND of each atom with same atom_index. Defaults to [].
+            species (Union[Sequence[str], Dict[str, List[int]], Dict[str, Dict[int, List[int]]]], optional): extract PBAND of each atom with same species. Defaults to [].
             efermi (float, optional): fermi level in unit eV. Defaults to 0.
             energy_range (Sequence[float], optional): energy range in unit eV for plotting. Defaults to [].
             shift (bool, optional): if shift energy by fermi level and set the VBM to zero, or not. Defaults to False.
+            outdir (PathLike): Default: './'
+            cmapname (str): Default: 'jet'
 
         Returns:
             BandPlot object: for manually plotting picture with bandplot.ax 

tools/plot-tools/abacus_plot/main.py

@@ -6,10 +6,9 @@
 '''
 
 import argparse
-from os import PathLike
 import matplotlib.pyplot as plt
 
-from abacus_plot.band import BandPlot
+from abacus_plot.band import Band, PBand
 from abacus_plot.dos import TDOS, PDOS
 from abacus_plot.utils import read_json
 
@@ -19,25 +18,50 @@ class Show:
 
     @classmethod
     def show_cmdline(cls, args):
-        if args.band:
-            text = read_json(args.band)
-            datafile = text["bandfile"]
+        if args.band and not args.projected and not args.out:
+            text = read_json(args.file)
+            bandfile = text["bandfile"]
             kptfile = text["kptfile"]
             efermi = text.pop("efermi", 0.0)
             energy_range = text.pop("energy_range", [])
             shift = text.pop("shift", False)
             figsize = text.pop("figsize", (12, 10))
             fig, ax = plt.subplots(figsize=figsize)
-            band = BandPlot(fig, ax)
-            if isinstance(datafile, (str, PathLike)):
-                band.singleplot(datafile, kptfile, efermi, energy_range, shift)
-            elif isinstance(datafile, (list, tuple)):
-                band.multiplot(datafile, kptfile, efermi, energy_range, shift)
+            band = Band(bandfile, kptfile)
+            band.plot(fig=fig, ax=ax, efermi=efermi, energy_range=energy_range, shift=shift, **text)
             dpi = text.pop("dpi", 400)
             fig.savefig(text.pop("outfile", "band.png"), dpi=dpi)
 
-        if args.tdos:
-            text = read_json(args.tdos)
+        if args.band and args.projected:
+            text = read_json(args.file)
+            bandfile = text["bandfile"]
+            kptfile = text["kptfile"]
+            efermi = text.pop("efermi", 0.0)
+            energy_range = text.pop("energy_range", [])
+            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", [])
+            outdir = text.pop("outdir", './')
+            cmapname = text.pop("cmapname", 'jet')
+            pband = PBand(bandfile, kptfile)
+            pband.plot(fig=fig, ax=ax, index=index, atom_index=atom_index, species=species, efermi=efermi, energy_range=energy_range, shift=shift, outdir=outdir, cmapname=cmapname, **text)
+
+        if args.band and args.out:
+            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", [])
+            outdir = text.pop("outdir", './')
+            pdos = PBand(bandfile, kptfile)
+            pdos.write(index=index, atom_index=atom_index, species=species, outdir=outdir)
+
+        if args.dos and not args.projected and not args.out:
+            text = read_json(args.file)
             tdosfile = text.pop("tdosfile", '')
             efermi = text.pop("efermi", 0.0)
             energy_range = text.pop("energy_range", [])
@@ -52,45 +76,51 @@ def show_cmdline(cls, args):
             dpi = text.pop("dpi", 400)
             fig.savefig(text.pop("tdosfig", "tdos.png"), dpi=dpi)
 
-        if args.pdos:
-            text = read_json(args.pdos)
+        if args.dos and args.projected:
+            text = read_json(args.file)
             pdosfile = text.pop("pdosfile", '')
             efermi = text.pop("efermi", 0.0)
             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", [])
             figsize = text.pop("figsize", (12, 10))
             fig, ax = plt.subplots(
                 len(species), 1, sharex=True, figsize=figsize)
             prec = text.pop("prec", 0.01)
             pdos = PDOS(pdosfile)
-            pdos.plot(fig, ax, species, efermi=efermi, shift=shift,
+            pdos.plot(fig=fig, ax=ax, index=index, atom_index=atom_index, species=species, efermi=efermi, shift=shift,
                       energy_range=energy_range, dos_range=dos_range, prec=prec, **text)
             dpi = text.pop("dpi", 400)
             fig.savefig(text.pop("pdosfig", "pdos.png"), dpi=dpi)
 
-        if args.out_pdos:
-            text = read_json(args.out_pdos)
+        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", [])
             outdir = text.pop("outdir", './')
             pdos = PDOS(pdosfile)
-            pdos.write(species, outdir)
+            pdos.write(index=index, atom_index=atom_index, species=species, outdir=outdir)
 
 def main():
     parser = argparse.ArgumentParser(
         prog='abacus-plot', description='Plotting tools for ABACUS')
 
     # Show
-    parser.add_argument('-b', '--band', dest='band', type=str,
+    parser.add_argument('-f', '--file', dest='file', type=str, nargs='?', const='config.json',
+                        default='config.json', help='profile with format json')
+    parser.add_argument('-b', '--band', dest='band', nargs='?', const=1, type=int,
                         default=None, help='plot band structure and show band information.')
-    parser.add_argument('-t', '--tdos', dest='tdos', type=str,
-                        default=None, help='plot total density of state(TDOS).')
-    parser.add_argument('-p', '--pdos', dest='pdos', type=str,
-                        default=None, help='plot partial density of state(PDOS).')
-    parser.add_argument('-o', '--out_pdos', dest='out_pdos', type=str,
-                        default=None, help='output partial density of state(PDOS).')
+    parser.add_argument('-d', '--dos', dest='dos', nargs='?', const=1, type=int,
+                        default=None, help='plot density of state(DOS).')
+    parser.add_argument('-p', '--projected', dest='projected', nargs='?', const=1, type=int,
+                        default=None, help='plot projected band structure or partial density of state(PDOS), should be used with `-b` or `-d`.')
+    parser.add_argument('-o', '--out_parsed_data', dest='out', nargs='?', const=1, type=int,
+                        default=None, help='output projected band structure or partial density of state(PDOS) to files, should be used with `-b` or `-d`.')
     parser.set_defaults(func=Show().show_cmdline)
 
     args = parser.parse_args()