From c907538c69ca0ca338684b49ddcd113f5330bcd0 Mon Sep 17 00:00:00 2001
From: jiyuang <jiyuyang@mail.ustc.com>
Date: Fri, 22 Apr 2022 18:30:15 +0800
Subject: [PATCH 01/12] extract PDOS by index(atom_index) using plot-tools

---
 tools/plot-tools/README.md           |  16 ++-
 tools/plot-tools/abacus_plot/band.py |  46 +++++--
 tools/plot-tools/abacus_plot/dos.py  | 178 +++++++++++++++++++++------
 3 files changed, 185 insertions(+), 55 deletions(-)

diff --git a/tools/plot-tools/README.md b/tools/plot-tools/README.md
index 743e1d45dc..97db50fb4d 100644
--- a/tools/plot-tools/README.md
+++ b/tools/plot-tools/README.md
@@ -19,6 +19,10 @@ 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`
+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
@@ -39,7 +43,7 @@ First, prepare a json file e.g. band-input.json:
 |    *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 to 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 *filename*, *efermi*, *label*, *color* in `list` type. 
 
 The *kptfile* should be as follows, and notes after `#` will be set as k-points label automatically.
 ```shell
@@ -75,10 +79,10 @@ First, prepare a json file e.g. dos-input.json:
 	],
 	"species": {
 			"C": {
-				"0": [
+				0: [
 					0
 				],
-				"1": [
+				1: [
 					0,
 					1
 				]
@@ -88,7 +92,7 @@ First, prepare a json file e.g. dos-input.json:
 				1
 			]
 		},
-	"pdosfig": "pdos.png",
+	"pdosfig": "pdos.png"
 }
 ```
 If you only want to plot total DOS, you can modify `pdosfile` to `tdosfile` and do not set `species` and `pdosfig`.
@@ -99,7 +103,9 @@ If you only want to plot total DOS, you can modify `pdosfile` to `tdosfile` and
 |    *efermi*    |                                   `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'`                        |
-|   *species*    | `List[str]` or `Dict[str, List[int]]` or `Dict[str, Dict[str, List[int]]]]` | Three ways to plot partial DOS e.g. ["C", "Si"], {"C":[0, 1], "Si":[0]}, ["C":{"0":[0]}, "Si":{"1":[0, 1]}] |
+|   *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]}] |
 |   *tdosfig*    |                                    `str`                                    |                                         Output picture of total DOS                                         |
 |   *pdosfig*    |                                    `str`                                    |                                        Output picture of partial DOS                                        |
 
diff --git a/tools/plot-tools/abacus_plot/band.py b/tools/plot-tools/abacus_plot/band.py
index 982ff597e3..66051baed6 100644
--- a/tools/plot-tools/abacus_plot/band.py
+++ b/tools/plot-tools/abacus_plot/band.py
@@ -239,6 +239,19 @@ def multiplot(self, datafile: Sequence[PathLike], kptfile: str = '', efermi: Seq
         index = kpt.label_special_k
         self._set_figure(index, energy_range)
 
+    @classmethod
+    def direct_bandgap(cls, vb: namedtuple, cb: namedtuple, klength: int):
+        """Calculate direct band gap"""
+
+        gap_list = []
+        i_index = []
+        for i in range(klength):
+            gap_list.append(np.min(cb.band[:, i])-np.max(vb.band[:, i]))
+            i_index.append(i)
+        dgap = np.min(gap_list)
+
+        return dgap, i_index[np.argmin(gap_list)]
+
     @classmethod
     def bandgap(cls, vb: namedtuple, cb: namedtuple):
         """Calculate band gap"""
@@ -247,6 +260,19 @@ def bandgap(cls, vb: namedtuple, cb: namedtuple):
 
         return gap
 
+    @classmethod
+    def band_type(cls, vb: namedtuple, cb: namedtuple):
+        vbm_x, cbm_x = vb.k_index, cb.k_index
+        longone, shortone = (vbm_x, cbm_x) if len(
+                vbm_x) >= len(cbm_x) else (cbm_x, vbm_x)
+        for i in shortone:
+            if i in longone:
+                btype = "Direct"
+            else:
+                btype = "Indirect"
+
+        return btype
+
     @classmethod
     def info(cls, kpath: Sequence, vb: namedtuple, cb: namedtuple):
         """Output the information of band structure
@@ -255,22 +281,18 @@ def info(cls, kpath: Sequence, vb: namedtuple, cb: namedtuple):
         :params energy: band energy after subtracting the Fermi level
         """
 
-        def band_type(vbm_x, cbm_x):
-            longone, shortone = (vbm_x, cbm_x) if len(
-                vbm_x) >= len(cbm_x) else (cbm_x, vbm_x)
-            for i in shortone:
-                if i in longone:
-                    btype = "Direct"
-                else:
-                    btype = "Indirect"
-            return btype
-
         gap = cls.bandgap(vb, cb)
+        dgap, d_i = cls.direct_bandgap(vb, cb, len(kpath))
+        btype = cls.band_type(vb, cb)
         print(
             "--------------------------Band Structure--------------------------", flush=True)
         print(
-            f"{'Band character:'.ljust(30)}{band_type(vb.k_index, cb.k_index)}", flush=True)
-        print(f"{'Band gap(eV):'.ljust(30)}{gap: .4f}", flush=True)
+            f"{'Band character:'.ljust(30)}{btype}", flush=True)
+        if btype == "Indirect":
+            print(f"{'Direct Band gap(eV):'.ljust(30)}{dgap: .4f}", flush=True)
+            print(f"{'Indirect Band gap(eV):'.ljust(30)}{gap: .4f}", flush=True)
+        elif btype == "Direct":
+            print(f"{'Band gap(eV):'.ljust(30)}{gap: .4f}", flush=True)
         print(f"{'Band index:'.ljust(30)}{'HOMO'.ljust(10)}{'LUMO'}", flush=True)
         print(
             f"{''.ljust(30)}{str(vb.band_index[-1]).ljust(10)}{str(cb.band_index[0])}", flush=True)
diff --git a/tools/plot-tools/abacus_plot/dos.py b/tools/plot-tools/abacus_plot/dos.py
index a289f78641..db7a1cef34 100644
--- a/tools/plot-tools/abacus_plot/dos.py
+++ b/tools/plot-tools/abacus_plot/dos.py
@@ -9,7 +9,7 @@
 import numpy as np
 from os import PathLike
 from pathlib import Path
-from typing import Dict, List, Sequence, Tuple, Union
+from typing import Dict, List, Sequence, Tuple, Union, Any
 from matplotlib.figure import Figure
 from matplotlib import axes
 
@@ -99,7 +99,7 @@ def __init__(self, fig: Figure, ax: axes.Axes, **kwargs) -> None:
         self._bwidth = kwargs.pop('bwdith', 3)
         self.plot_params = kwargs
 
-    def _set_figure(self, energy_range: Sequence, dos_range: Sequence, notes: Dict = {}):
+    def _set_figure(self, energy_range: Sequence = [], dos_range: Sequence = [], notes: Dict = {}):
         """set figure and axes for plotting
 
         :params energy_range: range of energy
@@ -257,11 +257,12 @@ def _all_sum(self) -> Tuple[np.ndarray, int]:
             res = res + orb['data']
         return res
 
-    def parse(self, species: Union[Sequence[str], Dict[str, List[int]], Dict[str, Dict[str, List[int]]]]):
+    def parse(self, species: Union[Sequence[Any], Dict[Any, List[int]], Dict[Any, Dict[Any, List[int]]]], keyname=''):
         """Extract partial dos from file
 
         Args:
-            species (Union[Sequence[str], Dict[str, List[int]], Dict[str, Dict[str, List[int]]]], optional): list of atomic species or dict of atomic species and its angular momentum list. Defaults to [].
+            species (Union[Sequence[Any], Dict[Any, List[int]], Dict[Any, Dict[str, List[int]]]], optional): list of atomic species(index or atom index) or dict of atomic species(index or atom index) and its angular momentum list. Defaults to [].
+            keyname (str): the keyword that extracts the PDOS
         """
 
         if isinstance(species, (list, tuple)):
@@ -271,7 +272,7 @@ def parse(self, species: Union[Sequence[str], Dict[str, List[int]], Dict[str, Di
                 count = 0
                 dos_temp = np.zeros_like(self.orbitals[0]["data"], dtype=float)
                 for orb in self.orbitals:
-                    if orb["species"] == elem:
+                    if orb[keyname] == elem:
                         dos_temp += orb["data"]
                         count += 1
                 if count:
@@ -294,7 +295,7 @@ def parse(self, species: Union[Sequence[str], Dict[str, List[int]], Dict[str, Di
                             dos_temp = np.zeros_like(
                                 self.orbitals[0]["data"], dtype=float)
                             for orb in self.orbitals:
-                                if orb["species"] == elem and orb["l"] == l_index and orb["m"] == m_index:
+                                if orb[keyname] == elem and orb["l"] == l_index and orb["m"] == m_index:
                                     dos_temp += orb["data"]
                                     m_count += 1
                                     l_count += 1
@@ -309,7 +310,7 @@ def parse(self, species: Union[Sequence[str], Dict[str, List[int]], Dict[str, Di
                         dos_temp = np.zeros_like(
                             self.orbitals[0]["data"], dtype=float)
                         for orb in self.orbitals:
-                            if orb["species"] == elem and orb["l"] == l_index:
+                            if orb[keyname] == elem and orb["l"] == l_index:
                                 dos_temp += orb["data"]
                                 count += 1
                         if count:
@@ -317,14 +318,15 @@ def parse(self, species: Union[Sequence[str], Dict[str, List[int]], Dict[str, Di
 
             return dos
 
-    def write(self, species: Union[Sequence[str], Dict[str, List[int]], Dict[str, Dict[str, List[int]]]], outdir: PathLike = './'):
+    def _write(self, species: Union[Sequence[Any], Dict[Any, List[int]], Dict[Any, Dict[int, List[int]]]], keyname='', outdir: PathLike = './'):
         """Write parsed partial dos data to files
 
         Args:
-            species (Union[Sequence[str], Dict[str, List[int]], Dict[str, Dict[str, List[int]]]], optional): list of atomic species or dict of atomic species and its angular momentum list. Defaults to [].
+            species (Union[Sequence[Any], Dict[Any, List[int]], Dict[Any, Dict[int, List[int]]]], optional): list of atomic species(index or atom index) or dict of atomic species(index or atom index) and its angular momentum list. Defaults to [].
+            keyname (str): the keyword that extracts the PDOS
         """
 
-        dos = self.parse(species)
+        dos = self.parse(species, keyname)
         fmt = ['%13.7f', '%15.8f'] if self._nsplit == 1 else [
             '%13.7f', '%15.8f', '%15.8f']
         file_dir = Path(f"{outdir}", "PDOS_FILE")
@@ -334,14 +336,14 @@ def write(self, species: Union[Sequence[str], Dict[str, List[int]], Dict[str, Di
             for elem in dos.keys():
                 header_list = ['']
                 data = np.hstack((self.energy.reshape(-1, 1), dos[elem]))
-                with open(file_dir/f"{elem}.dat", 'w') as f:
+                with open(file_dir/f"{keyname}-{elem}.dat", 'w') as f:
                     header_list.append(
                         f"\tpartial DOS for atom species: {elem}")
                     header_list.append('')
                     for orb in self.orbitals:
-                        if orb["species"] == elem:
+                        if orb[keyname] == elem:
                             header_list.append(
-                                f"\tAdd data for atom_index ={orb['atom_index']:4d},  l,m,z={orb['l']:3d}, {orb['m']:3d}, {orb['z']:3d}")
+                                f"\tAdd data for index ={orb['index']:4d}, atom_index ={orb['atom_index']:4d}, element ={orb['species']:4s},  l,m,z={orb['l']:3d}, {orb['m']:3d}, {orb['z']:3d}")
                     header_list.append('')
                     header_list.append('\tEnergy'+10*' ' +
                                        'spin 1'+8*' '+'spin 2')
@@ -351,7 +353,7 @@ def write(self, species: Union[Sequence[str], Dict[str, List[int]], Dict[str, Di
 
         elif isinstance(species, dict):
             for elem in dos.keys():
-                elem_file_dir = file_dir/f"{elem}"
+                elem_file_dir = file_dir/f"{keyname}-{elem}"
                 elem_file_dir.mkdir(exist_ok=True)
                 for ang in dos[elem].keys():
                     l_index = int(ang)
@@ -361,14 +363,14 @@ def write(self, species: Union[Sequence[str], Dict[str, List[int]], Dict[str, Di
                             data = np.hstack(
                                 (self.energy.reshape(-1, 1), dos[elem][ang][mag]))
                             m_index = int(mag)
-                            with open(elem_file_dir/f"{elem}_{ang}_{mag}.dat", 'w') as f:
+                            with open(elem_file_dir/f"{keyname}-{elem}_{ang}_{mag}.dat", 'w') as f:
                                 header_list.append(
                                     f"\tpartial DOS for atom species: {elem}")
                                 header_list.append('')
                                 for orb in self.orbitals:
-                                    if orb["species"] == elem and orb["l"] == l_index and orb["m"] == m_index:
+                                    if orb[keyname] == elem and orb["l"] == l_index and orb["m"] == m_index:
                                         header_list.append(
-                                            f"\tAdd data for atom_index ={orb['atom_index']:4d},  l,m,z={orb['l']:3d}, {orb['m']:3d}, {orb['z']:3d}")
+                                            f"\tAdd data for index ={orb['index']:4d}, atom_index ={orb['atom_index']:4d}, element ={orb['species']:4s},  l,m,z={orb['l']:3d}, {orb['m']:3d}, {orb['z']:3d}")
                                 header_list.append('')
                                 header_list.append(
                                     '\tEnergy'+10*' '+'spin 1'+8*' '+'spin 2')
@@ -380,14 +382,14 @@ def write(self, species: Union[Sequence[str], Dict[str, List[int]], Dict[str, Di
                         header_list = ['']
                         data = np.hstack(
                             (self.energy.reshape(-1, 1), dos[elem][ang]))
-                        with open(elem_file_dir/f"{elem}_{ang}.dat", 'w') as f:
+                        with open(elem_file_dir/f"{keyname}-{elem}_{ang}.dat", 'w') as f:
                             header_list.append(
                                 f"\tpartial DOS for atom species: {elem}")
                             header_list.append('')
                             for orb in self.orbitals:
-                                if orb["species"] == elem and orb["l"] == l_index:
+                                if orb[keyname] == elem and orb["l"] == l_index:
                                     header_list.append(
-                                        f"\tAdd data for atom_index ={orb['atom_index']:4d},  l,m,z={orb['l']:3d}, {orb['m']:3d}, {orb['z']:3d}")
+                                        f"\tAdd data for index ={orb['index']:4d}, atom_index ={orb['atom_index']:4d}, element ={orb['species']:4s},  l,m,z={orb['l']:3d}, {orb['m']:3d}, {orb['z']:3d}")
                             header_list.append('')
                             header_list.append(
                                 '\tEnergy'+10*' '+'spin 1'+8*' '+'spin 2')
@@ -395,6 +397,30 @@ def write(self, species: Union[Sequence[str], Dict[str, List[int]], Dict[str, Di
                             header = '\n'.join(header_list)
                             np.savetxt(f, data, fmt=fmt, header=header)
 
+    def write(self,
+              index: Union[Sequence[int], Dict[int, List[int]],
+                           Dict[int, Dict[int, List[int]]]] = [],
+              atom_index: Union[Sequence[int], Dict[int, List[int]],
+                                Dict[int, Dict[int, List[int]]]] = [],
+              species: Union[Sequence[str], Dict[str, List[int]],
+                             Dict[str, Dict[int, List[int]]]] = [],
+              outdir: PathLike = './'
+              ):
+        """Write parsed partial dos data to files
+
+        Args:
+            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 [].
+            outdir (PathLike, optional): directory of parsed PDOS files. Defaults to './'.
+        """
+        if index:
+            self._write(index, keyname='index', outdir=outdir)
+        if atom_index:
+            self._write(atom_index, keyname='atom_index', outdir=outdir)
+        if species:
+            self._write(species, keyname='species', outdir=outdir)
+
     def _shift_energy(self, efermi: float = 0, shift: bool = False, prec: float = 0.01):
         tdos = self._all_sum()
         if shift:
@@ -407,10 +433,36 @@ def _shift_energy(self, efermi: float = 0, shift: bool = False, prec: float = 0.
 
         return energy_f, tdos
 
-    def plot(self, fig: Figure, ax: Union[axes.Axes, Sequence[axes.Axes]], species: Union[Sequence[str], Dict[str, List[int]], Dict[str, Dict[str, List[int]]]] = [], efermi: float = 0, energy_range: Sequence[float] = [], dos_range: Sequence[float] = [], shift: bool = False, prec: float = 0.01, **kwargs):
-        """Plot partial DOS"""
+    def _parial_plot(self,
+                     fig: Figure,
+                     ax: Union[axes.Axes, Sequence[axes.Axes]],
+                     species: Union[Sequence[Any], Dict[Any, List[int]],
+                                    Dict[Any, Dict[int, List[int]]]] = [],
+                     efermi: float = 0,
+                     energy_range: Sequence[float] = [],
+                     dos_range: Sequence[float] = [],
+                     shift: bool = False,
+                     prec: float = 0.01,
+                     keyname: str = '',
+                     **kwargs):
+        """Plot parsed partial dos data
 
-        dos = self.parse(species)
+        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
+            species (Union[Sequence[Any], Dict[Any, List[int]], Dict[Any, Dict[int, List[int]]]], optional): list of atomic species(index or atom index) or dict of atomic species(index or atom index) and its angular momentum list. 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 [].
+            dos_range (Sequence[float], optional): dos range for plotting. Defaults to [].
+            shift (bool, optional): if shift energy by fermi level and set the VBM to zero, or not. Defaults to False.
+            prec (float, optional): precision for treating dos as zero. Defaults to 0.01.
+            keyname (str, optional): the keyword that extracts the PDOS. Defaults to ''.
+
+        Returns:
+            DOSPlot object: for manually plotting picture with dosplot.ax 
+        """
+
+        dos = self.parse(species, keyname)
         energy_f, tdos = self._shift_energy(efermi, shift, prec)
 
         if not species:
@@ -424,7 +476,7 @@ def plot(self, fig: Figure, ax: Union[axes.Axes, Sequence[axes.Axes]], species:
 
             return dosplot
 
-        elif isinstance(species, (list, tuple)):
+        if isinstance(species, (list, tuple)):
             dosplot = DOSPlot(fig, ax, **kwargs)
             if "xlabel_params" in dosplot.plot_params.keys():
                 dosplot.ax.set_xlabel("Energy(eV)", **
@@ -473,24 +525,74 @@ def plot(self, fig: Figure, ax: Union[axes.Axes, Sequence[axes.Axes]], species:
 
             return dosplots
 
+    def plot(self,
+             fig: Figure,
+             ax: Union[axes.Axes, Sequence[axes.Axes]],
+             index: Union[Sequence[int], Dict[int, List[int]],
+                          Dict[int, Dict[int, List[int]]]] = [],
+             atom_index: Union[Sequence[int], Dict[int, List[int]],
+                               Dict[int, Dict[int, List[int]]]] = [],
+             species: Union[Sequence[str], Dict[str, List[int]],
+                            Dict[str, Dict[int, List[int]]]] = [],
+             efermi: float = 0,
+             energy_range: Sequence[float] = [],
+             dos_range: Sequence[float] = [],
+             shift: bool = False,
+             prec: float = 0.01,
+             **kwargs):
+        """Plot parsed partial dos 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 [].
+            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 [].
+            dos_range (Sequence[float], optional): dos range for plotting. Defaults to [].
+            shift (bool, optional): if shift energy by fermi level and set the VBM to zero, or not. Defaults to False.
+            prec (float, optional): precision for treating dos as zero. Defaults to 0.01.
+
+        Returns:
+            DOSPlot object: for manually plotting picture with dosplot.ax 
+        """
+
+        if not index and not atom_index and not species:
+            dosplot = self._parial_plot(fig=fig, ax=ax, species=[
+            ], efermi=efermi, energy_range=energy_range, dos_range=dos_range, shift=shift, prec=prec, keyname='', **kwargs)
+        if index:
+            dosplot = self._parial_plot(fig=fig, ax=ax, species=index, efermi=efermi, energy_range=energy_range,
+                                        dos_range=dos_range, shift=shift, prec=prec, keyname='index', **kwargs)
+        if atom_index:
+            dosplot = self._parial_plot(fig=fig, ax=ax, species=atom_index, efermi=efermi, energy_range=energy_range,
+                                        dos_range=dos_range, shift=shift, prec=prec, keyname='atom_index', **kwargs)
+        if species:
+            dosplot = self._parial_plot(fig=fig, ax=ax, species=species, efermi=efermi, energy_range=energy_range,
+                                        dos_range=dos_range, shift=shift, prec=prec, keyname='species', **kwargs)
+
+        return dosplot
+
 
 if __name__ == "__main__":
     import matplotlib.pyplot as plt
     pdosfile = r"C:\Users\YY.Ji\Desktop\PDOS"
     pdos = PDOS(pdosfile)
-    species = {"Cs": [0, 1], "Na": [0, 1]}
-    fig, ax = plt.subplots(2, 1, sharex=True)
-    energy_range = [-6, 10]
+    #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]
     dos_range = [0, 5]
-    dosplots = pdos.plot(fig, ax, species, efermi=5, shift=True,
-                         energy_range=energy_range, dos_range=dos_range, notes=[{'s': '(a)'}, {'s': '(b)'}])
+    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)'}])
     fig.savefig("pdos.png")
-
-    tdosfile = r"C:\Users\YY.Ji\Desktop\TDOS"
-    tdos = TDOS(tdosfile)
-    fig, ax = plt.subplots()
-    energy_range = [-6, 10]
-    dos_range = [0, 5]
-    dosplots = tdos.plot(fig, ax, efermi=5, shift=True,
-                         energy_range=energy_range, dos_range=dos_range, notes={'s': '(a)'})
-    fig.savefig("tdos.png")
+    pdos.write(atom_index=atom_index)
+
+    #tdosfile = r"C:\Users\YY.Ji\Desktop\TDOS"
+    #tdos = TDOS(tdosfile)
+    #fig, ax = plt.subplots()
+    #energy_range = [-6, 10]
+    #dos_range = [0, 5]
+    # dosplots = tdos.plot(fig, ax, efermi=5, shift=True,
+    #                     energy_range=energy_range, dos_range=dos_range, notes={'s': '(a)'})
+    # fig.savefig("tdos.png")

From b9ea27520032d174ea33cc84acff0507c65182cc Mon Sep 17 00:00:00 2001
From: jiyuang <jiyuyang@mail.ustc.com>
Date: Mon, 25 Apr 2022 15:39:11 +0800
Subject: [PATCH 02/12] add function to parse projected band in plot-tools

---
 tools/plot-tools/abacus_plot/band.py  | 464 +++++++++++++++-----------
 tools/plot-tools/abacus_plot/dos.py   |  88 +----
 tools/plot-tools/abacus_plot/utils.py |  73 +++-
 tools/plot-tools/setup.py             |   2 +-
 4 files changed, 356 insertions(+), 271 deletions(-)

diff --git a/tools/plot-tools/abacus_plot/band.py b/tools/plot-tools/abacus_plot/band.py
index 66051baed6..4c8a98a8cd 100644
--- a/tools/plot-tools/abacus_plot/band.py
+++ b/tools/plot-tools/abacus_plot/band.py
@@ -8,25 +8,111 @@
 from collections import OrderedDict, namedtuple
 import numpy as np
 from os import PathLike
-from typing import Sequence, Tuple
+from typing import Sequence, Tuple, Union, Any, List, Dict
 from matplotlib.figure import Figure
 from matplotlib import axes
 
-from abacus_plot.utils import energy_minus_efermi, list_elem2str, read_kpt
+from abacus_plot.utils import energy_minus_efermi, list_elem2str, read_kpt, remove_empty, parse_projected_data, handle_data
 
 
-class BandPlot:
-    """Plot band structure"""
+class Band:
+    """Parse Bands data"""
 
-    def __init__(self, fig: Figure, ax: axes.Axes, **kwargs) -> None:
-        self.fig = fig
-        self.ax = ax
-        self._lw = kwargs.pop('lw', 2)
-        self._bwidth = kwargs.pop('bwdith', 3)
-        self._label = kwargs.pop('label', None)
-        self._color = kwargs.pop('color', None)
-        self._linestyle = kwargs.pop('linestyle', 'solid')
-        self.plot_params = kwargs
+    def __init__(self, bandfile: Union[PathLike, Sequence[PathLike]] = None, kptfile: str = '') -> None:
+        self.bandfile = bandfile
+        self.read()
+        self.kptfile = kptfile
+        self.kpt = None
+        if self.kptfile:
+            self.kpt = read_kpt(kptfile)
+        self.nspin = None
+        self.norbitals = None
+        self.orbitals = []
+
+    def read(self):
+        """Read band data file and return k-points and energy
+
+        :params filename: string of band data file
+        """
+
+        data = np.loadtxt(self.bandfile)
+        X, self.k_index = np.split(data, (1, ), axis=1)
+        self.nkpoints = len(self.k_index)
+        self.energy = X.flatten()
+        self.nbands = self.energy.shape[-1]
+        self.eunit = 'eV'
+
+    @classmethod
+    def direct_bandgap(cls, vb: namedtuple, cb: namedtuple, klength: int):
+        """Calculate direct band gap"""
+
+        gap_list = []
+        i_index = []
+        for i in range(klength):
+            gap_list.append(np.min(cb.band[:, i])-np.max(vb.band[:, i]))
+            i_index.append(i)
+        dgap = np.min(gap_list)
+
+        return dgap, i_index[np.argmin(gap_list)]
+
+    @classmethod
+    def bandgap(cls, vb: namedtuple, cb: namedtuple):
+        """Calculate band gap"""
+
+        gap = cb.value-vb.value
+
+        return gap
+
+    @classmethod
+    def band_type(cls, vb: namedtuple, cb: namedtuple):
+        vbm_x, cbm_x = vb.k_index, cb.k_index
+        longone, shortone = (vbm_x, cbm_x) if len(
+            vbm_x) >= len(cbm_x) else (cbm_x, vbm_x)
+        for i in shortone:
+            if i in longone:
+                btype = "Direct"
+            else:
+                btype = "Indirect"
+
+        return btype
+
+    @classmethod
+    def info(cls, kpath: Sequence, vb: namedtuple, cb: namedtuple):
+        """Output the information of band structure
+
+        :params kpath: k-points path 
+        :params energy: band energy after subtracting the Fermi level
+        """
+
+        gap = cls.bandgap(vb, cb)
+        dgap, d_i = cls.direct_bandgap(vb, cb, len(kpath))
+        btype = cls.band_type(vb, cb)
+        print(
+            "--------------------------Band Structure--------------------------", flush=True)
+        print(
+            f"{'Band character:'.ljust(30)}{btype}", flush=True)
+        if btype == "Indirect":
+            print(f"{'Direct Band gap(eV):'.ljust(30)}{dgap: .4f}", flush=True)
+            print(f"{'Indirect Band gap(eV):'.ljust(30)}{gap: .4f}", flush=True)
+        elif btype == "Direct":
+            print(f"{'Band gap(eV):'.ljust(30)}{gap: .4f}", flush=True)
+        print(f"{'Band index:'.ljust(30)}{'HOMO'.ljust(10)}{'LUMO'}", flush=True)
+        print(
+            f"{''.ljust(30)}{str(vb.band_index[-1]).ljust(10)}{str(cb.band_index[0])}", flush=True)
+        print(f"{'Eigenvalue of VBM(eV):'.ljust(30)}{vb.value: .4f}", flush=True)
+        print(f"{'Eigenvalue of CBM(eV):'.ljust(30)}{cb.value: .4f}", flush=True)
+        vbm_k = np.unique(kpath[vb.k_index], axis=0)
+        cbm_k = np.unique(kpath[cb.k_index], axis=0)
+        print(
+            f"{'Location of VBM'.ljust(30)}{' '.join(list_elem2str(vbm_k[0]))}", flush=True)
+        for i, j in enumerate(vbm_k):
+            if i != 0:
+                print(f"{''.ljust(30)}{' '.join(list_elem2str(j))}", flush=True)
+        print(
+            f"{'Location of CBM'.ljust(30)}{' '.join(list_elem2str(cbm_k[0]))}", flush=True)
+        for i, j in enumerate(cbm_k):
+            if i != 0:
+                print(f"{''.ljust(30)}{' '.join(list_elem2str(j))}", flush=True)
 
     @classmethod
     def set_vcband(cls, energy: Sequence) -> Tuple[namedtuple, namedtuple]:
@@ -56,18 +142,101 @@ def set_vcband(cls, energy: Sequence) -> Tuple[namedtuple, namedtuple]:
 
         return vb, cb
 
-    @classmethod
-    def read(cls, filename: PathLike) -> Tuple[np.ndarray, np.ndarray]:
-        """Read band data file and return k-points and energy
+    def plot(self, x: Sequence, y: Sequence, index: Sequence, efermi: float = 0, energy_range: Sequence[float] = []):
+        """Plot band structure
 
-        :params filename: string of band data file
+        :params x, y: x-axis and y-axis coordinates
+        :params index: special k-points label and its index in data file
+        :params efermi: Fermi level in unit eV
+        :params energy_range: range of energy to plot, its length equals to two
         """
 
-        data = np.loadtxt(filename)
-        X, y = np.split(data, (1, ), axis=1)
-        x = X.flatten()
+        if not self._color:
+            self._color = 'black'
+
+        kpoints, energy = x, y
+        energy = energy_minus_efermi(energy, efermi)
+
+        self.ax.plot(kpoints, energy, lw=self._lw, color=self._color,
+                     label=self._label, linestyle=self._linestyle)
+        self._set_figure(index, energy_range)
 
-        return x, y
+    def singleplot(self, efermi: float = 0, energy_range: Sequence[float] = [], shift: bool = False):
+        """Plot band structure using data file
+
+        :params efermi: Fermi level in unit eV
+        :params energy_range: range of energy to plot, its length equals to two
+        :params shift: if sets True, it will calculate band gap. This parameter usually is suitable for semiconductor and insulator. Default: False
+        """
+
+        if not self._color:
+            self._color = 'black'
+
+        energy = energy_minus_efermi(self.energy, efermi)
+        if shift:
+            vb, cb = self.set_vcband(energy)
+            self.ax.plot(self.k_index, np.vstack((vb.band, cb.band)).T,
+                         lw=self._lw, color=self._color, label=self._label, linestyle=self._linestyle)
+            self.info(self.kpt.full_kpath, vb, cb)
+        else:
+            self.ax.plot(self.k_index, energy,
+                         lw=self._lw, color=self._color, label=self._label, linestyle=self._linestyle)
+        index = self.kpt.label_special_k
+        self._set_figure(index, energy_range)
+
+    def multiplot(self, efermi: Sequence[float] = [], energy_range: Sequence[float] = [], shift: bool = True):
+        """Plot more than two band structures using data file
+
+        :params efermi: list of Fermi levels in unit eV, its length equals to `filename`
+        :params energy_range: range of energy to plot, its length equals to two
+        :params shift: if sets True, it will calculate band gap. This parameter usually is suitable for semiconductor and insulator. Default: False
+        """
+
+        if not efermi:
+            efermi = [0.0 for i in range(len(self.bandfile))]
+        if not self._label:
+            self._label = ['' for i in range(len(self.bandfile))]
+        if not self._color:
+            self._color = ['black' for i in range(len(self.bandfile))]
+        if not self._linestyle:
+            self._linestyle = ['solid' for i in range(len(self.bandfile))]
+
+        emin = -np.inf
+        emax = np.inf
+        for i, file in enumerate(self.bandfile):
+            if shift:
+                vb, cb = self.set_vcband(
+                    energy_minus_efermi(self.energy, efermi[i]))
+                energy_min = np.min(vb.band)
+                energy_max = np.max(cb.band)
+                if energy_min > emin:
+                    emin = energy_min
+                if energy_max < emax:
+                    emax = energy_max
+
+                self.ax.plot(self.k_index, np.vstack((vb.band, cb.band)).T,
+                             lw=self._lw, color=self._color[i], label=self._label[i], linestyle=self._linestyle[i])
+                self.info(self.kpt.full_kpath, vb, cb)
+            else:
+                self.ax.plot(self.k_index, energy_minus_efermi(self.energy, efermi[i]),
+                             lw=self._lw, color=self._color[i], label=self._label[i], linestyle=self._linestyle[i])
+
+        index = self.kpt.label_special_k
+        self._set_figure(index, energy_range)
+
+
+class BandPlot:
+    """Plot band structure"""
+
+    def __init__(self, fig: Figure, ax: axes.Axes, **kwargs) -> None:
+        self.fig = fig
+        self.ax = ax
+        self._lw = kwargs.pop('lw', 2)
+        self._bwidth = kwargs.pop('bwdith', 3)
+        self._label = kwargs.pop('label', None)
+        self._color = kwargs.pop('color', None)
+        self._linestyle = kwargs.pop('linestyle', 'solid')
+        self.plot_params = kwargs
 
     def _set_figure(self, index: dict, range: Sequence):
         """set figure and axes for plotting
@@ -147,188 +316,103 @@ def _set_figure(self, index: dict, range: Sequence):
                 self.ax.legend(by_label.values(),
                                by_label.keys(), prop={'size': 15})
 
-    def plot(self, x: Sequence, y: Sequence, index: Sequence, efermi: float = 0, energy_range: Sequence[float] = []):
-        """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 efermi: Fermi level in unit eV
-        :params energy_range: range of energy to plot, its length equals to two
-        """
+class PBand(Band):
+    def __init__(self, bandfile: PathLike = None) -> None:
+        self.bandfile = bandfile
+        self.read()
 
-        if not self._color:
-            self._color = 'black'
+    def _check_energy(self):
+        assert self.energy.shape[0] == self.nkpoints, "The dimension of band structure dismatches with the number of k-points."
+        assert self.energy.shape[1] == self.nbands, "The dimension of band structure dismatches with the number of bands."
 
-        kpoints, energy = x, y
-        energy = energy_minus_efermi(energy, efermi)
+    def _check_weights(self, weights:np.ndarray, prec=1e-5):
+        assert weights.shape[0] == self.norbitals, "The dimension of weights dismatches with the number of orbitals."
+        assert weights.shape[1] == self.nkpoints, "The dimension of weights dismatches with the number of k-points."
+        assert weights.shape[2] == self.nbands, "The dimension of weights dismatches with the number of bands."
+        one_mat = np.ones((self.nkpoints, self.nbands))
+        assert (np.abs(weights.sum(axis=0)-one_mat) < prec).all(), f"np.abs(weights.sum(axis=0)-np.ones(({self.nkpoints}, {self.nbands}))) < {prec}"
 
-        self.ax.plot(kpoints, energy, lw=self._lw, color=self._color,
-                     label=self._label, linestyle=self._linestyle)
-        self._set_figure(index, energy_range)
+    @property
+    def weights(self):
+        data = np.empty((self.norbitals, self.nkpoints, self.nbands))
+        for i, orb in enumerate(self.orbitals):
+            data[i] = orb['data']
+        self._check_weights(data)
+        return data
 
-    def singleplot(self, datafile: PathLike, kptfile: str = '', efermi: float = 0, energy_range: Sequence[float] = [], shift: bool = False):
-        """Plot band structure using data file
+    def read(self):
+        """Read projected band data file and return k-points, energy and Mulliken weights
 
-        :params datafile: string of band date file
-        :params kptfile: k-point file
-        :params efermi: Fermi level in unit eV
-        :params energy_range: range of energy to plot, its length equals to two
-        :params shift: if sets True, it will calculate band gap. This parameter usually is suitable for semiconductor and insulator. Default: False
+        :params filename: string of projected band data file
         """
 
-        kpt = read_kpt(kptfile)
-
-        if not self._color:
-            self._color = 'black'
-
-        kpoints, energy = self.read(datafile)
-        energy = energy_minus_efermi(energy, efermi)
-        if shift:
-            vb, cb = self.set_vcband(energy)
-            self.ax.plot(kpoints, np.vstack((vb.band, cb.band)).T,
-                         lw=self._lw, color=self._color, label=self._label, linestyle=self._linestyle)
-            self.info(kpt.full_kpath, vb, cb)
-        else:
-            self.ax.plot(kpoints, energy,
-                         lw=self._lw, color=self._color, label=self._label, linestyle=self._linestyle)
-        index = kpt.label_special_k
-        self._set_figure(index, energy_range)
-
-    def multiplot(self, datafile: Sequence[PathLike], kptfile: str = '', efermi: Sequence[float] = [], energy_range: Sequence[float] = [], shift: bool = True):
-        """Plot more than two band structures using data file
-
-        :params datafile: list of path of band date file 
-        :params kptfile: k-point file
-        :params efermi: list of Fermi levels in unit eV, its length equals to `filename`
-        :params energy_range: range of energy to plot, its length equals to two
-        :params shift: if sets True, it will calculate band gap. This parameter usually is suitable for semiconductor and insulator. Default: False
-        """
-
-        kpt = read_kpt(kptfile)
-
-        if not efermi:
-            efermi = [0.0 for i in range(len(datafile))]
-        if not self._label:
-            self._label = ['' for i in range(len(datafile))]
-        if not self._color:
-            self._color = ['black' for i in range(len(datafile))]
-        if not self._linestyle:
-            self._linestyle = ['solid' for i in range(len(datafile))]
-
-        emin = -np.inf
-        emax = np.inf
-        for i, file in enumerate(datafile):
-            kpoints, energy = self.read(file)
-            if shift:
-                vb, cb = self.set_vcband(
-                    energy_minus_efermi(energy, efermi[i]))
-                energy_min = np.min(vb.band)
-                energy_max = np.max(cb.band)
-                if energy_min > emin:
-                    emin = energy_min
-                if energy_max < emax:
-                    emax = energy_max
-
-                self.ax.plot(kpoints, np.vstack((vb.band, cb.band)).T,
-                             lw=self._lw, color=self._color[i], label=self._label[i], linestyle=self._linestyle[i])
-                self.info(kpt.full_kpath, vb, cb)
-            else:
-                self.ax.plot(kpoints, energy_minus_efermi(energy, efermi[i]),
-                             lw=self._lw, color=self._color[i], label=self._label[i], linestyle=self._linestyle[i])
-
-        index = kpt.label_special_k
-        self._set_figure(index, energy_range)
-
-    @classmethod
-    def direct_bandgap(cls, vb: namedtuple, cb: namedtuple, klength: int):
-        """Calculate direct band gap"""
-
-        gap_list = []
-        i_index = []
-        for i in range(klength):
-            gap_list.append(np.min(cb.band[:, i])-np.max(vb.band[:, i]))
-            i_index.append(i)
-        dgap = np.min(gap_list)
-
-        return dgap, i_index[np.argmin(gap_list)]
-
-    @classmethod
-    def bandgap(cls, vb: namedtuple, cb: namedtuple):
-        """Calculate band gap"""
-
-        gap = cb.value-vb.value
-
-        return gap
-
-    @classmethod
-    def band_type(cls, vb: namedtuple, cb: namedtuple):
-        vbm_x, cbm_x = vb.k_index, cb.k_index
-        longone, shortone = (vbm_x, cbm_x) if len(
-                vbm_x) >= len(cbm_x) else (cbm_x, vbm_x)
-        for i in shortone:
-            if i in longone:
-                btype = "Direct"
-            else:
-                btype = "Indirect"
-
-        return btype
-
-    @classmethod
-    def info(cls, kpath: Sequence, vb: namedtuple, cb: namedtuple):
-        """Output the information of band structure
-
-        :params kpath: k-points path 
-        :params energy: band energy after subtracting the Fermi level
+        from lxml import etree
+        pbanddata = etree.parse(self.bandfile)
+        root = pbanddata.getroot()
+        self.nspin = int(root.xpath('//nspin')[0].text.replace(' ', ''))
+        self.norbitals = int(root.xpath('//norbitals')
+                             [0].text.replace(' ', ''))
+        self.eunit = root.xpath('//band_structure/@units')[0].replace(' ', '')
+        self.nbands = int(root.xpath('//band_structure/@nbands')
+                          [0].replace(' ', ''))
+        self.nkpoints = int(root.xpath('//band_structure/@nkpoints')
+                            [0].replace(' ', ''))
+        self.k_index = np.arange(self.nkpoints)
+        self.energy = root.xpath('//band_structure')[0].text.split('\n')
+        self.energy = handle_data(self.energy)
+        remove_empty(self.energy)
+        self.energy = np.asarray(self.energy, dtype=float)
+        self._check_energy()
+
+        self.orbitals = []
+        for i in range(self.norbitals):
+            orb = OrderedDict()
+            o_index_str = root.xpath(
+                '//orbital/@index')[i]
+            orb['index'] = int(o_index_str.replace(' ', ''))
+            orb['atom_index'] = int(root.xpath(
+                '//orbital/@atom_index')[i].replace(' ', ''))
+            orb['species'] = root.xpath(
+                '//orbital/@species')[i].replace(' ', '')
+            orb['l'] = int(root.xpath('//orbital/@l')[i].replace(' ', ''))
+            orb['m'] = int(root.xpath('//orbital/@m')[i].replace(' ', ''))
+            orb['z'] = int(root.xpath('//orbital/@z')[i].replace(' ', ''))
+            data = root.xpath('//data')[i].text.split('\n')
+            data = handle_data(data)
+            remove_empty(data)
+            orb['data'] = np.asarray(data, dtype=float)
+            self.orbitals.append(orb)
+
+    def _write(self, species: Union[Sequence[Any], Dict[Any, List[int]], Dict[Any, Dict[int, List[int]]]], keyname='', outdir: PathLike = './'):
+        """Write parsed partial dos data to files
+
+        Args:
+            species (Union[Sequence[Any], Dict[Any, List[int]], Dict[Any, Dict[int, List[int]]]], optional): list of atomic species(index or atom index) or dict of atomic species(index or atom index) and its angular momentum list. Defaults to [].
+            keyname (str): the keyword that extracts the PDOS. Allowed values: 'index', 'atom_index', 'species'
         """
 
-        gap = cls.bandgap(vb, cb)
-        dgap, d_i = cls.direct_bandgap(vb, cb, len(kpath))
-        btype = cls.band_type(vb, cb)
-        print(
-            "--------------------------Band Structure--------------------------", flush=True)
-        print(
-            f"{'Band character:'.ljust(30)}{btype}", flush=True)
-        if btype == "Indirect":
-            print(f"{'Direct Band gap(eV):'.ljust(30)}{dgap: .4f}", flush=True)
-            print(f"{'Indirect Band gap(eV):'.ljust(30)}{gap: .4f}", flush=True)
-        elif btype == "Direct":
-            print(f"{'Band gap(eV):'.ljust(30)}{gap: .4f}", flush=True)
-        print(f"{'Band index:'.ljust(30)}{'HOMO'.ljust(10)}{'LUMO'}", flush=True)
-        print(
-            f"{''.ljust(30)}{str(vb.band_index[-1]).ljust(10)}{str(cb.band_index[0])}", flush=True)
-        print(f"{'Eigenvalue of VBM(eV):'.ljust(30)}{vb.value: .4f}", flush=True)
-        print(f"{'Eigenvalue of CBM(eV):'.ljust(30)}{cb.value: .4f}", flush=True)
-        vbm_k = np.unique(kpath[vb.k_index], axis=0)
-        cbm_k = np.unique(kpath[cb.k_index], axis=0)
-        print(
-            f"{'Location of VBM'.ljust(30)}{' '.join(list_elem2str(vbm_k[0]))}", flush=True)
-        for i, j in enumerate(vbm_k):
-            if i != 0:
-                print(f"{''.ljust(30)}{' '.join(list_elem2str(j))}", flush=True)
-        print(
-            f"{'Location of CBM'.ljust(30)}{' '.join(list_elem2str(cbm_k[0]))}", flush=True)
-        for i, j in enumerate(cbm_k):
-            if i != 0:
-                print(f"{''.ljust(30)}{' '.join(list_elem2str(j))}", flush=True)
+        return 
 
 
 if __name__ == "__main__":
     import matplotlib.pyplot as plt
     from pathlib import Path
-    parent = Path(r"D:\ustc\TEST\HOIP\double HOIP\result\bond")
-    name = "CsAgBiBr"
+    parent = Path(r"C:\Users\YY.Ji\Desktop")
+    name = "PBANDS_1"
     path = parent/name
-    notes = {'s': '(b)'}
-    datafile = [path/"soc.dat", path/"non-soc.dat"]
-    kptfile = path/"KPT"
-    fig, ax = plt.subplots(figsize=(12, 12))
-    label = ["with SOC", "without SOC"]
-    color = ["r", "g"]
-    linestyle = ["solid", "dashed"]
-    band = BandPlot(fig, ax, notes=notes, label=label,
-                    color=color, linestyle=linestyle)
-    energy_range = [-5, 6]
-    efermi = [4.417301755850272, 4.920435541999894]
-    shift = True
-    band.multiplot(datafile, kptfile, efermi, energy_range, shift)
-    fig.savefig("band.png")
+    # notes = {'s': '(b)'}
+    # datafile = [path/"soc.dat", path/"non-soc.dat"]
+    # kptfile = path/"KPT"
+    #fig, ax = plt.subplots(figsize=(12, 12))
+    # label = ["with SOC", "without SOC"]
+    # color = ["r", "g"]
+    # linestyle = ["solid", "dashed"]
+    # band = BandPlot(fig, ax, notes=notes, label=label,
+    #                 color=color, linestyle=linestyle)
+    # energy_range = [-5, 6]
+    # efermi = [4.417301755850272, 4.920435541999894]
+    # shift = True
+    # band.multiplot(datafile, kptfile, efermi, energy_range, shift)
+    # fig.savefig("band.png")
+    pband = PBand(str(path))
diff --git a/tools/plot-tools/abacus_plot/dos.py b/tools/plot-tools/abacus_plot/dos.py
index db7a1cef34..c9ad5714f0 100644
--- a/tools/plot-tools/abacus_plot/dos.py
+++ b/tools/plot-tools/abacus_plot/dos.py
@@ -5,7 +5,7 @@
 Mail: jiyuyang@mail.ustc.edu.cn, 1041176461@qq.com
 '''
 
-from collections import OrderedDict, defaultdict, namedtuple
+from collections import OrderedDict, namedtuple
 import numpy as np
 from os import PathLike
 from pathlib import Path
@@ -14,7 +14,7 @@
 from matplotlib import axes
 
 from abacus_plot.utils import (energy_minus_efermi, get_angular_momentum_label,
-                               get_angular_momentum_name, remove_empty)
+                               get_angular_momentum_name, remove_empty, handle_data, parse_projected_data)
 
 
 class DOS:
@@ -92,7 +92,7 @@ def bandgap(cls, vb: namedtuple, cb: namedtuple):
 class DOSPlot:
     """Plot density of state(DOS)"""
 
-    def __init__(self, fig: Figure, ax: axes.Axes, **kwargs) -> None:
+    def __init__(self, fig: Figure = None, ax: axes.Axes = None, **kwargs) -> None:
         self.fig = fig
         self.ax = ax
         self._lw = kwargs.pop('lw', 2)
@@ -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:
+    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:
+    def __init__(self, pdosfile: PathLike=None) -> None:
         super().__init__()
         self.pdosfile = pdosfile
         self._read()
@@ -213,15 +213,6 @@ def _read(self):
         :params pdosfile: string of PDOS data file
         """
 
-        def handle_data(data):
-            data.remove('')
-
-            def handle_elem(elem):
-                elist = elem.split(' ')
-                remove_empty(elist)  # `list` will be modified in function
-                return elist
-            return list(map(handle_elem, data))
-
         from lxml import etree
         pdosdata = etree.parse(self.pdosfile)
         root = pdosdata.getroot()
@@ -257,76 +248,15 @@ def _all_sum(self) -> Tuple[np.ndarray, int]:
             res = res + orb['data']
         return res
 
-    def parse(self, species: Union[Sequence[Any], Dict[Any, List[int]], Dict[Any, Dict[Any, List[int]]]], keyname=''):
-        """Extract partial dos from file
-
-        Args:
-            species (Union[Sequence[Any], Dict[Any, List[int]], Dict[Any, Dict[str, List[int]]]], optional): list of atomic species(index or atom index) or dict of atomic species(index or atom index) and its angular momentum list. Defaults to [].
-            keyname (str): the keyword that extracts the PDOS
-        """
-
-        if isinstance(species, (list, tuple)):
-            dos = {}
-            elements = species
-            for elem in elements:
-                count = 0
-                dos_temp = np.zeros_like(self.orbitals[0]["data"], dtype=float)
-                for orb in self.orbitals:
-                    if orb[keyname] == elem:
-                        dos_temp += orb["data"]
-                        count += 1
-                if count:
-                    dos[elem] = dos_temp
-
-            return dos
-
-        elif isinstance(species, dict):
-            dos = defaultdict(dict)
-            elements = list(species.keys())
-            l = list(species.values())
-            for i, elem in enumerate(elements):
-                if isinstance(l[i], dict):
-                    for ang, mag in l[i].items():
-                        l_count = 0
-                        l_index = int(ang)
-                        l_dos = {}
-                        for m_index in mag:
-                            m_count = 0
-                            dos_temp = np.zeros_like(
-                                self.orbitals[0]["data"], dtype=float)
-                            for orb in self.orbitals:
-                                if orb[keyname] == elem and orb["l"] == l_index and orb["m"] == m_index:
-                                    dos_temp += orb["data"]
-                                    m_count += 1
-                                    l_count += 1
-                            if m_count:
-                                l_dos[m_index] = dos_temp
-                        if l_count:
-                            dos[elem][l_index] = l_dos
-
-                elif isinstance(l[i], list):
-                    for l_index in l[i]:
-                        count = 0
-                        dos_temp = np.zeros_like(
-                            self.orbitals[0]["data"], dtype=float)
-                        for orb in self.orbitals:
-                            if orb[keyname] == elem and orb["l"] == l_index:
-                                dos_temp += orb["data"]
-                                count += 1
-                        if count:
-                            dos[elem][l_index] = dos_temp
-
-            return dos
-
     def _write(self, species: Union[Sequence[Any], Dict[Any, List[int]], Dict[Any, Dict[int, List[int]]]], keyname='', outdir: PathLike = './'):
         """Write parsed partial dos data to files
 
         Args:
             species (Union[Sequence[Any], Dict[Any, List[int]], Dict[Any, Dict[int, List[int]]]], optional): list of atomic species(index or atom index) or dict of atomic species(index or atom index) and its angular momentum list. Defaults to [].
-            keyname (str): the keyword that extracts the PDOS
+            keyname (str): the keyword that extracts the PDOS. Allowed values: 'index', 'atom_index', 'species'
         """
 
-        dos = self.parse(species, keyname)
+        dos = parse_projected_data(self.orbitals, species, keyname)
         fmt = ['%13.7f', '%15.8f'] if self._nsplit == 1 else [
             '%13.7f', '%15.8f', '%15.8f']
         file_dir = Path(f"{outdir}", "PDOS_FILE")
@@ -462,7 +392,7 @@ def _parial_plot(self,
             DOSPlot object: for manually plotting picture with dosplot.ax 
         """
 
-        dos = self.parse(species, keyname)
+        dos = parse_projected_data(self.orbitals, species, keyname)
         energy_f, tdos = self._shift_energy(efermi, shift, prec)
 
         if not species:
@@ -579,7 +509,7 @@ def plot(self,
     pdosfile = r"C:\Users\YY.Ji\Desktop\PDOS"
     pdos = PDOS(pdosfile)
     #species = {"Ag": [2], "Cl": [1], "In": [0]}
-    atom_index = {8: {2:[1, 2]}, 4: {2:[1, 2]}, 10: [1, 2]}
+    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]
     dos_range = [0, 5]
diff --git a/tools/plot-tools/abacus_plot/utils.py b/tools/plot-tools/abacus_plot/utils.py
index 72be9e0b2a..7d5012cfa6 100644
--- a/tools/plot-tools/abacus_plot/utils.py
+++ b/tools/plot-tools/abacus_plot/utils.py
@@ -9,7 +9,8 @@
 import re
 import string
 from os import PathLike
-from typing import List, Sequence, Union
+from typing import List, Sequence, Union, Any, Dict
+from collections import defaultdict
 
 import numpy as np
 
@@ -21,6 +22,14 @@ def remove_empty(a: list) -> list:
     while [] in a:
         a.remove([])
 
+def handle_data(data):
+    data.remove('')
+
+    def handle_elem(elem):
+        elist = elem.split(' ')
+        remove_empty(elist)  # `list` will be modified in function
+        return elist
+    return list(map(handle_elem, data))
 
 def skip_notes(line: str) -> str:
     """Delete comments lines with '#' or '//'
@@ -257,3 +266,65 @@ def get_angular_momentum_name(l_index: int, m_index: int) -> str:
     """
 
     return angular_momentum_name[l_index][m_index]
+
+
+def parse_projected_data(orbitals, species: Union[Sequence[Any], Dict[Any, List[int]], Dict[Any, Dict[Any, List[int]]]], keyname=''):
+    """Extract projected data from file
+
+    Args:
+        species (Union[Sequence[Any], Dict[Any, List[int]], Dict[Any, Dict[str, List[int]]]], optional): list of atomic species(index or atom index) or dict of atomic species(index or atom index) and its angular momentum list. Defaults to [].
+        keyname (str): the keyword that extracts the projected data. Allowed values: 'index', 'atom_index', 'species'
+    """
+
+    if isinstance(species, (list, tuple)):
+        data = {}
+        elements = species
+        for elem in elements:
+            count = 0
+            data_temp = np.zeros_like(orbitals[0]["data"], dtype=float)
+            for orb in orbitals:
+                if orb[keyname] == elem:
+                    data_temp += orb["data"]
+                    count += 1
+            if count:
+                data[elem] = data_temp
+
+        return data
+
+    elif isinstance(species, dict):
+        data = defaultdict(dict)
+        elements = list(species.keys())
+        l = list(species.values())
+        for i, elem in enumerate(elements):
+            if isinstance(l[i], dict):
+                for ang, mag in l[i].items():
+                    l_count = 0
+                    l_index = int(ang)
+                    l_data = {}
+                    for m_index in mag:
+                        m_count = 0
+                        data_temp = np.zeros_like(
+                            orbitals[0]["data"], dtype=float)
+                        for orb in orbitals:
+                            if orb[keyname] == elem and orb["l"] == l_index and orb["m"] == m_index:
+                                data_temp += orb["data"]
+                                m_count += 1
+                                l_count += 1
+                        if m_count:
+                            l_data[m_index] = data_temp
+                    if l_count:
+                        data[elem][l_index] = l_data
+
+            elif isinstance(l[i], list):
+                for l_index in l[i]:
+                    count = 0
+                    data_temp = np.zeros_like(
+                        orbitals[0]["data"], dtype=float)
+                    for orb in orbitals:
+                        if orb[keyname] == elem and orb["l"] == l_index:
+                            data_temp += orb["data"]
+                            count += 1
+                    if count:
+                        data[elem][l_index] = data_temp
+
+        return data
\ No newline at end of file
diff --git a/tools/plot-tools/setup.py b/tools/plot-tools/setup.py
index b76066b0b0..265dfd7d9d 100644
--- a/tools/plot-tools/setup.py
+++ b/tools/plot-tools/setup.py
@@ -10,7 +10,7 @@
 if __name__ == "__main__":
     setup(
         name='abacus_plot',
-        version='1.1.0',
+        version='1.2.0',
         packages=find_packages(),
         description='Ploting tools for ABACUS',
         author='jiyuyang',

From 38e13fe49983b34a1fcf98ddae5a77f7dc7f4e61 Mon Sep 17 00:00:00 2001
From: jiyuang <jiyuyang@mail.ustc.com>
Date: Wed, 27 Apr 2022 16:24:04 +0800
Subject: [PATCH 03/12] add projected band plotting tools

---
 tools/plot-tools/abacus_plot/band.py  | 581 ++++++++++++++++++++------
 tools/plot-tools/abacus_plot/dos.py   |  12 +-
 tools/plot-tools/abacus_plot/utils.py |   7 +-
 3 files changed, 473 insertions(+), 127 deletions(-)

diff --git a/tools/plot-tools/abacus_plot/band.py b/tools/plot-tools/abacus_plot/band.py
index 4c8a98a8cd..41d30b52e3 100644
--- a/tools/plot-tools/abacus_plot/band.py
+++ b/tools/plot-tools/abacus_plot/band.py
@@ -1,7 +1,7 @@
 '''
 Date: 2021-12-29 10:27:01
 LastEditors: jiyuyang
-LastEditTime: 2022-01-03 17:06:14
+LastEditTime: 2022-04-26 00:00:41
 Mail: jiyuyang@mail.ustc.edu.cn, 1041176461@qq.com
 '''
 
@@ -10,9 +10,13 @@
 from os import PathLike
 from typing import Sequence, Tuple, Union, Any, List, Dict
 from matplotlib.figure import Figure
+from matplotlib.colors import Normalize
+from matplotlib.collections import LineCollection
 from matplotlib import axes
+import matplotlib.pyplot as plt
+from pathlib import Path
 
-from abacus_plot.utils import energy_minus_efermi, list_elem2str, read_kpt, remove_empty, parse_projected_data, handle_data
+from abacus_plot.utils import energy_minus_efermi, list_elem2str, read_kpt, remove_empty, parse_projected_data, handle_data, get_angular_momentum_name, get_angular_momentum_label
 
 
 class Band:
@@ -20,27 +24,32 @@ class Band:
 
     def __init__(self, bandfile: Union[PathLike, Sequence[PathLike]] = None, kptfile: str = '') -> None:
         self.bandfile = bandfile
-        self.read()
+        if isinstance(bandfile, list) or isinstance(bandfile, tuple):
+            self.energy = []
+            for file in self.bandfile:
+                self.k_index, e = self.read(file)
+                self.energy.append(e)
+        else:
+            self.k_index, self.energy = self.read(self.bandfile)
+        self.energy = np.asarray(self.energy)
         self.kptfile = kptfile
         self.kpt = None
         if self.kptfile:
             self.kpt = read_kpt(kptfile)
-        self.nspin = None
-        self.norbitals = None
-        self.orbitals = []
+        self.k_index = list(map(int, self.k_index))
 
-    def read(self):
+    @classmethod
+    def read(cls, filename: PathLike):
         """Read band data file and return k-points and energy
 
         :params filename: string of band data file
         """
 
-        data = np.loadtxt(self.bandfile)
-        X, self.k_index = np.split(data, (1, ), axis=1)
-        self.nkpoints = len(self.k_index)
-        self.energy = X.flatten()
-        self.nbands = self.energy.shape[-1]
-        self.eunit = 'eV'
+        data = np.loadtxt(filename, dtype=float)
+        X, y = np.split(data, (1, ), axis=1)
+        x = X.flatten()
+
+        return x, y
 
     @classmethod
     def direct_bandgap(cls, vb: namedtuple, cb: namedtuple, klength: int):
@@ -80,7 +89,7 @@ def band_type(cls, vb: namedtuple, cb: namedtuple):
     def info(cls, kpath: Sequence, vb: namedtuple, cb: namedtuple):
         """Output the information of band structure
 
-        :params kpath: k-points path 
+        :params kpath: k-points path
         :params energy: band energy after subtracting the Fermi level
         """
 
@@ -142,7 +151,27 @@ def set_vcband(cls, energy: Sequence) -> Tuple[namedtuple, namedtuple]:
 
         return vb, cb
 
-    def plot(self, x: Sequence, y: Sequence, index: Sequence, efermi: float = 0, energy_range: Sequence[float] = []):
+    def _shift_energy(self, energy, efermi: float = 0, shift: bool = False):
+        energy = energy_minus_efermi(energy, efermi)
+        if shift:
+            vb, cb = self.set_vcband(energy)
+            refine_E = np.vstack((vb.band, cb.band)).T
+            self.info(self.kpt.full_kpath, vb, cb)
+        else:
+            refine_E = energy
+
+        return refine_E
+
+    @classmethod
+    def plot_data(cls,
+                  fig: Figure,
+                  ax: axes.Axes,
+                  x: Sequence,
+                  y: Sequence,
+                  index: Sequence,
+                  efermi: float = 0,
+                  energy_range: Sequence[float] = [],
+                  **kwargs):
         """Plot band structure
 
         :params x, y: x-axis and y-axis coordinates
@@ -151,78 +180,64 @@ def plot(self, x: Sequence, y: Sequence, index: Sequence, efermi: float = 0, ene
         :params energy_range: range of energy to plot, its length equals to two
         """
 
-        if not self._color:
-            self._color = 'black'
+        bandplot = BandPlot(fig, ax, **kwargs)
+        if not bandplot._color:
+            bandplot._color = 'black'
 
         kpoints, energy = x, y
         energy = energy_minus_efermi(energy, efermi)
 
-        self.ax.plot(kpoints, energy, lw=self._lw, color=self._color,
-                     label=self._label, linestyle=self._linestyle)
-        self._set_figure(index, energy_range)
-
-    def singleplot(self, efermi: float = 0, energy_range: Sequence[float] = [], shift: bool = False):
-        """Plot band structure using data file
+        bandplot.ax.plot(kpoints, energy, lw=bandplot._lw, color=bandplot._color,
+                         label=bandplot._label, linestyle=bandplot._linestyle)
+        bandplot._set_figure(index, energy_range)
+
+    def plot(self,
+             fig: Figure,
+             ax: axes.Axes,
+             efermi: Union[float, Sequence[float]] = [],
+             energy_range: Sequence[float] = [],
+             shift: bool = True,
+             **kwargs):
+        """Plot more than two band structures using data file
 
-        :params efermi: Fermi level in unit eV
+        :params efermi: Fermi levels in unit eV, its length equals to `filename`
         :params energy_range: range of energy to plot, its length equals to two
         :params shift: if sets True, it will calculate band gap. This parameter usually is suitable for semiconductor and insulator. Default: False
         """
 
-        if not self._color:
-            self._color = 'black'
+        bandplot = BandPlot(fig, ax, **kwargs)
+        nums = len(self.bandfile)
 
-        energy = energy_minus_efermi(self.energy, efermi)
-        if shift:
-            vb, cb = self.set_vcband(energy)
-            self.ax.plot(self.k_index, np.vstack((vb.band, cb.band)).T,
-                         lw=self._lw, color=self._color, label=self._label, linestyle=self._linestyle)
-            self.info(self.kpt.full_kpath, vb, cb)
-        else:
-            self.ax.plot(self.k_index, energy,
-                         lw=self._lw, color=self._color, label=self._label, linestyle=self._linestyle)
-        index = self.kpt.label_special_k
-        self._set_figure(index, energy_range)
+        if isinstance(self.bandfile, list):
+            if not efermi:
+                efermi = [0.0 for i in range(nums)]
+            if not kwargs.pop('label', None):
+                bandplot._label = ['' for i in range(nums)]
+            if not kwargs.pop('color', None):
+                bandplot._color = ['black' for i in range(nums)]
+            if not kwargs.pop('linestyle', None):
+                bandplot._linestyle = ['solid' for i in range(nums)]
 
-    def multiplot(self, efermi: Sequence[float] = [], energy_range: Sequence[float] = [], shift: bool = True):
-        """Plot more than two band structures using data file
+            for i, band in enumerate(self.energy):
+                band = self._shift_energy(band, efermi[i], shift)
+                bandplot.ax.plot(self.k_index, band,
+                                 lw=bandplot._lw, color=bandplot._color[i], label=bandplot._label[i], linestyle=bandplot._linestyle[i])
 
-        :params efermi: list of Fermi levels in unit eV, its length equals to `filename`
-        :params energy_range: range of energy to plot, its length equals to two
-        :params shift: if sets True, it will calculate band gap. This parameter usually is suitable for semiconductor and insulator. Default: False
-        """
+        else:
+            if not efermi:
+                efermi = 0.0
 
-        if not efermi:
-            efermi = [0.0 for i in range(len(self.bandfile))]
-        if not self._label:
-            self._label = ['' for i in range(len(self.bandfile))]
-        if not self._color:
-            self._color = ['black' for i in range(len(self.bandfile))]
-        if not self._linestyle:
-            self._linestyle = ['solid' for i in range(len(self.bandfile))]
-
-        emin = -np.inf
-        emax = np.inf
-        for i, file in enumerate(self.bandfile):
-            if shift:
-                vb, cb = self.set_vcband(
-                    energy_minus_efermi(self.energy, efermi[i]))
-                energy_min = np.min(vb.band)
-                energy_max = np.max(cb.band)
-                if energy_min > emin:
-                    emin = energy_min
-                if energy_max < emax:
-                    emax = energy_max
-
-                self.ax.plot(self.k_index, np.vstack((vb.band, cb.band)).T,
-                             lw=self._lw, color=self._color[i], label=self._label[i], linestyle=self._linestyle[i])
-                self.info(self.kpt.full_kpath, vb, cb)
-            else:
-                self.ax.plot(self.k_index, energy_minus_efermi(self.energy, efermi[i]),
-                             lw=self._lw, color=self._color[i], label=self._label[i], linestyle=self._linestyle[i])
+            band = self._shift_energy(self.energy, efermi, shift)
+            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)
 
-        index = self.kpt.label_special_k
-        self._set_figure(index, energy_range)
+        return bandplot
 
 
 class BandPlot:
@@ -234,11 +249,11 @@ def __init__(self, fig: Figure, ax: axes.Axes, **kwargs) -> None:
         self._lw = kwargs.pop('lw', 2)
         self._bwidth = kwargs.pop('bwdith', 3)
         self._label = kwargs.pop('label', None)
-        self._color = kwargs.pop('color', None)
+        self._color = kwargs.pop('color', 'black')
         self._linestyle = kwargs.pop('linestyle', 'solid')
         self.plot_params = kwargs
 
-    def _set_figure(self, index: dict, range: Sequence):
+    def _set_figure(self, index, 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()
@@ -251,7 +266,7 @@ def _set_figure(self, index: dict, range: Sequence):
             if isinstance(t, tuple):
                 keys.append(t[0])
                 values.append(t[1])
-            elif isinstance(t, (int, float)):
+            elif isinstance(t, int):
                 keys.append('')
                 values.append(t)
 
@@ -297,10 +312,11 @@ def _set_figure(self, index: dict, range: Sequence):
         self.ax.spines['bottom'].set_linewidth(bwidth)
 
         # guides
-        if "grid_params" in self.plot_params.keys():
-            self.ax.grid(axis='x', **self.plot_params["grid_params"])
-        else:
-            self.ax.grid(axis='x', lw=1.2)
+        if '' not in keys:
+            if "grid_params" in self.plot_params.keys():
+                self.ax.grid(axis='x', **self.plot_params["grid_params"])
+            else:
+                self.ax.grid(axis='x', lw=1.2)
         if "hline_params" in self.plot_params.keys():
             self.ax.axhline(0, **self.plot_params["hline_params"])
         else:
@@ -318,20 +334,39 @@ def _set_figure(self, index: dict, range: Sequence):
 
 
 class PBand(Band):
-    def __init__(self, bandfile: PathLike = None) -> None:
+    def __init__(self, bandfile: Union[PathLike, Sequence[PathLike]] = None, kptfile: str = '') -> None:
         self.bandfile = bandfile
-        self.read()
+        if isinstance(bandfile, list) or isinstance(bandfile, tuple):
+            self.energy = []
+            self.orbitals = []
+            for file in self.bandfile:
+                self.nspin, self.norbitals, self.eunit, self.nbands, self.nkpoints, self.k_index, e, orb = self.read(
+                    file)
+                self._check_energy(e)
+                self.energy.append(e)
+                self.orbitals.append(orb)
+        else:
+            self.nspin, self.norbitals, self.eunit, self.nbands, self.nkpoints, self.k_index, self.energy, self.orbitals = self.read(
+                self.bandfile)
+            self._check_energy(self.energy)
+        self.energy = np.asarray(self.energy)
+        self.kptfile = kptfile
+        self.kpt = None
+        if self.kptfile:
+            self.kpt = read_kpt(kptfile)
+        self.k_index = list(map(int, self.k_index))
 
-    def _check_energy(self):
-        assert self.energy.shape[0] == self.nkpoints, "The dimension of band structure dismatches with the number of k-points."
-        assert self.energy.shape[1] == self.nbands, "The dimension of band structure dismatches with the number of bands."
+    def _check_energy(self, energy):
+        assert energy.shape[0] == self.nkpoints, "The dimension of band structure dismatches with the number of k-points."
+        assert energy.shape[1] == self.nbands, "The dimension of band structure dismatches with the number of bands."
 
-    def _check_weights(self, weights:np.ndarray, prec=1e-5):
+    def _check_weights(self, weights: np.ndarray, prec=1e-5):
         assert weights.shape[0] == self.norbitals, "The dimension of weights dismatches with the number of orbitals."
         assert weights.shape[1] == self.nkpoints, "The dimension of weights dismatches with the number of k-points."
         assert weights.shape[2] == self.nbands, "The dimension of weights dismatches with the number of bands."
         one_mat = np.ones((self.nkpoints, self.nbands))
-        assert (np.abs(weights.sum(axis=0)-one_mat) < prec).all(), f"np.abs(weights.sum(axis=0)-np.ones(({self.nkpoints}, {self.nbands}))) < {prec}"
+        assert (np.abs(weights.sum(axis=0)-one_mat) < prec).all(
+        ), f"np.abs(weights.sum(axis=0)-np.ones(({self.nkpoints}, {self.nbands}))) < {prec}"
 
     @property
     def weights(self):
@@ -341,32 +376,32 @@ def weights(self):
         self._check_weights(data)
         return data
 
-    def read(self):
+    @classmethod
+    def read(cls, filename: PathLike):
         """Read projected band data file and return k-points, energy and Mulliken weights
 
-        :params filename: string of projected band data file
+        :params bandfile: string of projected band data file
         """
 
         from lxml import etree
-        pbanddata = etree.parse(self.bandfile)
+        pbanddata = etree.parse(filename)
         root = pbanddata.getroot()
-        self.nspin = int(root.xpath('//nspin')[0].text.replace(' ', ''))
-        self.norbitals = int(root.xpath('//norbitals')
-                             [0].text.replace(' ', ''))
-        self.eunit = root.xpath('//band_structure/@units')[0].replace(' ', '')
-        self.nbands = int(root.xpath('//band_structure/@nbands')
-                          [0].replace(' ', ''))
-        self.nkpoints = int(root.xpath('//band_structure/@nkpoints')
-                            [0].replace(' ', ''))
-        self.k_index = np.arange(self.nkpoints)
-        self.energy = root.xpath('//band_structure')[0].text.split('\n')
-        self.energy = handle_data(self.energy)
-        remove_empty(self.energy)
-        self.energy = np.asarray(self.energy, dtype=float)
-        self._check_energy()
-
-        self.orbitals = []
-        for i in range(self.norbitals):
+        nspin = int(root.xpath('//nspin')[0].text.replace(' ', ''))
+        norbitals = int(root.xpath('//norbitals')
+                        [0].text.replace(' ', ''))
+        eunit = root.xpath('//band_structure/@units')[0].replace(' ', '')
+        nbands = int(root.xpath('//band_structure/@nbands')
+                     [0].replace(' ', ''))
+        nkpoints = int(root.xpath('//band_structure/@nkpoints')
+                       [0].replace(' ', ''))
+        k_index = np.arange(nkpoints)
+        energy = root.xpath('//band_structure')[0].text.split('\n')
+        energy = handle_data(energy)
+        remove_empty(energy)
+        energy = np.asarray(energy, dtype=float)
+
+        orbitals = []
+        for i in range(norbitals):
             orb = OrderedDict()
             o_index_str = root.xpath(
                 '//orbital/@index')[i]
@@ -382,17 +417,324 @@ def read(self):
             data = handle_data(data)
             remove_empty(data)
             orb['data'] = np.asarray(data, dtype=float)
-            self.orbitals.append(orb)
+            orbitals.append(orb)
+
+        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='', outdir: 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:
+            orbital (dict): parsed data
+            species (Union[Sequence[Any], Dict[Any, List[int]], Dict[Any, Dict[int, List[int]]]], optional): list of atomic species(index or atom index) or dict of atomic species(index or atom index) and its angular momentum list. Defaults to [].
+            keyname (str): the keyword that extracts the PBAND. Allowed values: 'index', 'atom_index', 'species'
+        """
+
+        band, totnum = parse_projected_data(self.orbitals, species, keyname)
+
+        if isinstance(species, (list, tuple)):
+            for elem in band.keys():
+                header_list = ['']
+                with open(file_dir/f"{keyname}-{elem}.dat", 'w') as f:
+                    header_list.append(
+                        f"Projected band structure for {keyname}: {elem}")
+                    header_list.append('')
+                    header_list.append(
+                        f'\tNumber of k-points: {self.nkpoints}')
+                    header_list.append(f'\tNumber of bands: {self.nbands}')
+                    header_list.append('')
+                    for orb in self.orbitals:
+                        if orb[keyname] == elem:
+                            header_list.append(
+                                f"\tAdd data for index ={orb['index']:4d}, atom_index ={orb['atom_index']:4d}, element ={orb['species']:4s},  l,m,z={orb['l']:3d}, {orb['m']:3d}, {orb['z']:3d}")
+                    header_list.append('')
+                    header_list.append(
+                        f'Data shape: ({self.nkpoints}, {self.nbands})')
+                    header_list.append('')
+                    header = '\n'.join(header_list)
+                    np.savetxt(f, band[elem], header=header)
+
+        elif isinstance(species, dict):
+            for elem in band.keys():
+                elem_file_dir = file_dir/f"{keyname}-{elem}"
+                elem_file_dir.mkdir(exist_ok=True)
+                for ang in band[elem].keys():
+                    l_index = int(ang)
+                    if isinstance(band[elem][ang], dict):
+                        for mag in band[elem][ang].keys():
+                            header_list = ['']
+                            m_index = int(mag)
+                            with open(elem_file_dir/f"{keyname}-{elem}_{ang}_{mag}.dat", 'w') as f:
+                                header_list.append(
+                                    f"Projected band structure for {keyname}: {elem}")
+                                header_list.append('')
+                                header_list.append(
+                                    f'\tNumber of k-points: {self.nkpoints}')
+                                header_list.append(
+                                    f'\tNumber of bands: {self.nbands}')
+                                header_list.append('')
+                                for orb in self.orbitals:
+                                    if orb[keyname] == elem and orb["l"] == l_index and orb["m"] == m_index:
+                                        header_list.append(
+                                            f"\tAdd data for index ={orb['index']:4d}, atom_index ={orb['atom_index']:4d}, element ={orb['species']:4s},  l,m,z={orb['l']:3d}, {orb['m']:3d}, {orb['z']:3d}")
+                                header_list.append('')
+                                header_list.append(
+                                    f'Data shape: ({self.nkpoints}, {self.nbands})')
+                                header_list.append('')
+                                header = '\n'.join(header_list)
+                                np.savetxt(f, band[elem][ang]
+                                           [mag], header=header)
+
+                    else:
+                        header_list = ['']
+                        with open(elem_file_dir/f"{keyname}-{elem}_{ang}.dat", 'w') as f:
+                            header_list.append(
+                                f"Projected band structure for {keyname}: {elem}")
+                            header_list.append('')
+                            header_list.append(
+                                f'\tNumber of k-points: {self.nkpoints}')
+                            header_list.append(
+                                f'\tNumber of bands: {self.nbands}')
+                            header_list.append('')
+                            for orb in self.orbitals:
+                                if orb[keyname] == elem and orb["l"] == l_index:
+                                    header_list.append(
+                                        f"\tAdd data for index ={orb['index']:4d}, atom_index ={orb['atom_index']:4d}, element ={orb['species']:4s},  l,m,z={orb['l']:3d}, {orb['m']:3d}, {orb['z']:3d}")
+                            header_list.append('')
+                            header_list.append(
+                                f'Data shape: ({self.nkpoints}, {self.nbands})')
+                            header_list.append('')
+                            header = '\n'.join(header_list)
+                            np.savetxt(f, band[elem][ang], header=header)
+
+    def write(self,
+              index: Union[Sequence[int], Dict[int, List[int]],
+                           Dict[int, Dict[int, List[int]]]] = [],
+              atom_index: Union[Sequence[int], Dict[int, List[int]],
+                                Dict[int, Dict[int, List[int]]]] = [],
+              species: Union[Sequence[str], Dict[str, List[int]],
+                             Dict[str, Dict[int, List[int]]]] = [],
+              outdir: PathLike = './'
+              ):
         """Write parsed partial dos data to files
 
         Args:
+            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 [].
+            outdir (PathLike, optional): directory of parsed PDOS files. Defaults to './'.
+        """
+
+        if isinstance(self.bandfile, list):
+            for i in range(len(self.bandfile)):
+                file_dir = Path(f"{outdir}", f"PBAND{i}_FILE")
+                file_dir.mkdir(exist_ok=True)
+                if index:
+                    self._write(index, keyname='index',
+                                file_dir=file_dir)
+                if atom_index:
+                    self._write(atom_index, keyname='atom_index',
+                                file_dir=file_dir)
+                if species:
+                    self._write(species, keyname='species',
+                                file_dir=file_dir)
+
+        else:
+            file_dir = Path(f"{outdir}", f"PBAND{1}_FILE")
+            file_dir.mkdir(exist_ok=True)
+            if index:
+                self._write(index, keyname='index',
+                            file_dir=file_dir)
+            if atom_index:
+                self._write(atom_index, keyname='atom_index',
+                            file_dir=file_dir)
+            if species:
+                self._write(species, keyname='species',
+                            file_dir=file_dir)
+
+    def _plot(self,
+              fig: Figure,
+              ax: axes.Axes,
+              energy: np.ndarray,
+              species: Union[Sequence[Any], Dict[Any, List[int]],
+                             Dict[Any, Dict[int, List[int]]]] = [],
+              efermi: float = 0,
+              energy_range: Sequence[float] = [],
+              shift: bool = False,
+              keyname: str = '',
+              outdir: PathLike = './',
+              out_index: int = 1,
+              cmap='jet',
+              **kwargs):
+        """Plot parsed projected bands 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
             species (Union[Sequence[Any], Dict[Any, List[int]], Dict[Any, Dict[int, List[int]]]], optional): list of atomic species(index or atom index) or dict of atomic species(index or atom index) and its angular momentum list. Defaults to [].
-            keyname (str): the keyword that extracts the PDOS. Allowed values: 'index', 'atom_index', 'species'
+            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.
+            keyname (str, optional): the keyword that extracts the PBANDS. Defaults to ''.
+
+        Returns:
+            BandPlot object: for manually plotting picture with bandplot.ax 
+        """
+
+        def _seg_plot(bandplot, lc, index, 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)
+            cbar.remove()
+            plt.cla()
+
+            return bandplot
+
+        wei, totnum = parse_projected_data(self.orbitals, species, keyname)
+        energy = self._shift_energy(energy, efermi, shift)
+        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)
+
+            return bandplot
+
+        if isinstance(species, (list, tuple)):
+            bandplots = []
+            for i, elem in enumerate(wei.keys()):
+                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)
+                    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}')
+                bandplots.append(bandplot)
+            return bandplots
+
+        elif isinstance(species, dict):
+            bandplots = []
+            for i, elem in enumerate(wei.keys()):
+                elem_file_dir = file_dir/f"{keyname}-{elem}"
+                elem_file_dir.mkdir(exist_ok=True)
+                for ang in wei[elem].keys():
+                    l_index = int(ang)
+                    if isinstance(wei[elem][ang], dict):
+                        for mag in wei[elem][ang].keys():
+                            bandplot = BandPlot(fig, ax, **kwargs)
+                            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)
+                                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}')
+                            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)
+                            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}')
+                        bandplots.append(bandplot)
+
+            return bandplots
+
+        plt.clf()
+
+    def plot(self,
+             fig: Figure,
+             ax: Union[axes.Axes, Sequence[axes.Axes]],
+             index: Union[Sequence[int], Dict[int, List[int]],
+                          Dict[int, Dict[int, List[int]]]] = [],
+             atom_index: Union[Sequence[int], Dict[int, List[int]],
+                               Dict[int, Dict[int, List[int]]]] = [],
+             species: Union[Sequence[str], Dict[str, List[int]],
+                            Dict[str, Dict[int, List[int]]]] = [],
+             efermi: Union[float, Sequence[float]] = [],
+             energy_range: Sequence[float] = [],
+             shift: bool = False,
+             outdir: PathLike = './',
+             cmapname='jet',
+             **kwargs):
+        """Plot parsed partial dos 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 [].
+            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.
+
+        Returns:
+            BandPlot object: for manually plotting picture with bandplot.ax 
         """
+        nums = len(self.bandfile)
+
+        if isinstance(self.bandfile, list):
+            if not efermi:
+                efermi = [0.0 for i in range(nums)]
+            _linestyle = kwargs.pop(
+                'linestyle', ['solid' for i in range(nums)])
+
+            for i, band in enumerate(self.energy):
+                if not index and not atom_index and not species:
+                    bandplot = self._plot(fig=fig, ax=ax, energy=band, species=[
+                    ], efermi=efermi[i], energy_range=energy_range, shift=shift, keyname='', linestyle=_linestyle[i], outdir=outdir, out_index=i, cmapname=cmapname, **kwargs)
+                if index:
+                    bandplot = self._plot(fig=fig, ax=ax, energy=band, species=index, efermi=efermi[i],
+                                          energy_range=energy_range, shift=shift, keyname='index', linestyle=_linestyle[i], outdir=outdir, out_index=i, cmapname=cmapname, **kwargs)
+                if atom_index:
+                    bandplot = self._plot(fig=fig, ax=ax, energy=band, species=atom_index, efermi=efermi[i],
+                                          energy_range=energy_range, shift=shift, keyname='atom_index', linestyle=_linestyle[i], outdir=outdir, out_index=i, cmapname=cmapname, **kwargs)
+                if species:
+                    bandplot = self._plot(fig=fig, ax=ax, energy=band, species=species, efermi=efermi[i],
+                                          energy_range=energy_range, shift=shift, keyname='species', linestyle=_linestyle[i], outdir=outdir, out_index=i, cmapname=cmapname, **kwargs)
 
-        return 
+        else:
+            if not index and not atom_index and not species:
+                bandplot = self._plot(fig=fig, ax=ax, energy=self.energy, species=[
+                ], efermi=efermi, energy_range=energy_range, shift=shift, keyname='', outdir=outdir, out_index=1, **kwargs)
+            if index:
+                bandplot = self._plot(fig=fig, ax=ax, energy=self.energy, species=index, efermi=efermi,
+                                      energy_range=energy_range, shift=shift, keyname='index', outdir=outdir, out_index=1, **kwargs)
+            if atom_index:
+                bandplot = self._plot(fig=fig, ax=ax, energy=self.energy, species=atom_index, efermi=efermi,
+                                      energy_range=energy_range, shift=shift, keyname='atom_index', outdir=outdir, out_index=1, **kwargs)
+            if species:
+                bandplot = self._plot(fig=fig, ax=ax, energy=self.energy, species=species, efermi=efermi,
+                                      energy_range=energy_range, shift=shift, keyname='species', outdir=outdir, out_index=1, **kwargs)
+
+        return bandplot
 
 
 if __name__ == "__main__":
@@ -404,15 +746,16 @@ def _write(self, species: Union[Sequence[Any], Dict[Any, List[int]], Dict[Any, D
     # notes = {'s': '(b)'}
     # datafile = [path/"soc.dat", path/"non-soc.dat"]
     # kptfile = path/"KPT"
-    #fig, ax = plt.subplots(figsize=(12, 12))
+    fig, ax = plt.subplots(figsize=(12, 6))
     # label = ["with SOC", "without SOC"]
     # color = ["r", "g"]
     # linestyle = ["solid", "dashed"]
-    # band = BandPlot(fig, ax, notes=notes, label=label,
-    #                 color=color, linestyle=linestyle)
-    # energy_range = [-5, 6]
-    # efermi = [4.417301755850272, 4.920435541999894]
-    # shift = True
-    # band.multiplot(datafile, kptfile, efermi, energy_range, shift)
-    # fig.savefig("band.png")
+    energy_range = [-5, 6]
+    efermi = 4.417301755850272
+    shift = False
+    #species = {"Ag": [2], "Cl": [1], "In": [0]}
+    atom_index = {8: {2: [1, 2]}, 4: {2: [1, 2]}, 10: [1, 2]}
     pband = PBand(str(path))
+    pband.plot(fig, ax, atom_index=atom_index, efermi=efermi,
+               energy_range=energy_range, shift=shift)
+    pband.write(atom_index=atom_index)
diff --git a/tools/plot-tools/abacus_plot/dos.py b/tools/plot-tools/abacus_plot/dos.py
index c9ad5714f0..da1296008f 100644
--- a/tools/plot-tools/abacus_plot/dos.py
+++ b/tools/plot-tools/abacus_plot/dos.py
@@ -256,7 +256,7 @@ def _write(self, species: Union[Sequence[Any], Dict[Any, List[int]], Dict[Any, D
             keyname (str): the keyword that extracts the PDOS. Allowed values: 'index', 'atom_index', 'species'
         """
 
-        dos = parse_projected_data(self.orbitals, species, keyname)
+        dos, totnum = parse_projected_data(self.orbitals, species, keyname)
         fmt = ['%13.7f', '%15.8f'] if self._nsplit == 1 else [
             '%13.7f', '%15.8f', '%15.8f']
         file_dir = Path(f"{outdir}", "PDOS_FILE")
@@ -268,7 +268,7 @@ def _write(self, species: Union[Sequence[Any], Dict[Any, List[int]], Dict[Any, D
                 data = np.hstack((self.energy.reshape(-1, 1), dos[elem]))
                 with open(file_dir/f"{keyname}-{elem}.dat", 'w') as f:
                     header_list.append(
-                        f"\tpartial DOS for atom species: {elem}")
+                        f"Partial DOS for {keyname}: {elem}")
                     header_list.append('')
                     for orb in self.orbitals:
                         if orb[keyname] == elem:
@@ -295,7 +295,7 @@ def _write(self, species: Union[Sequence[Any], Dict[Any, List[int]], Dict[Any, D
                             m_index = int(mag)
                             with open(elem_file_dir/f"{keyname}-{elem}_{ang}_{mag}.dat", 'w') as f:
                                 header_list.append(
-                                    f"\tpartial DOS for atom species: {elem}")
+                                    f"Partial DOS for {keyname}: {elem}")
                                 header_list.append('')
                                 for orb in self.orbitals:
                                     if orb[keyname] == elem and orb["l"] == l_index and orb["m"] == m_index:
@@ -314,7 +314,7 @@ def _write(self, species: Union[Sequence[Any], Dict[Any, List[int]], Dict[Any, D
                             (self.energy.reshape(-1, 1), dos[elem][ang]))
                         with open(elem_file_dir/f"{keyname}-{elem}_{ang}.dat", 'w') as f:
                             header_list.append(
-                                f"\tpartial DOS for atom species: {elem}")
+                                f"Partial DOS for {keyname}: {elem}")
                             header_list.append('')
                             for orb in self.orbitals:
                                 if orb[keyname] == elem and orb["l"] == l_index:
@@ -392,7 +392,7 @@ def _parial_plot(self,
             DOSPlot object: for manually plotting picture with dosplot.ax 
         """
 
-        dos = parse_projected_data(self.orbitals, species, keyname)
+        dos, totnum = parse_projected_data(self.orbitals, species, keyname)
         energy_f, tdos = self._shift_energy(efermi, shift, prec)
 
         if not species:
@@ -403,7 +403,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)):
diff --git a/tools/plot-tools/abacus_plot/utils.py b/tools/plot-tools/abacus_plot/utils.py
index 7d5012cfa6..8e31a87c58 100644
--- a/tools/plot-tools/abacus_plot/utils.py
+++ b/tools/plot-tools/abacus_plot/utils.py
@@ -289,12 +289,13 @@ def parse_projected_data(orbitals, species: Union[Sequence[Any], Dict[Any, List[
             if count:
                 data[elem] = data_temp
 
-        return data
+        return data, len(elements)
 
     elif isinstance(species, dict):
         data = defaultdict(dict)
         elements = list(species.keys())
         l = list(species.values())
+        totnum = 0
         for i, elem in enumerate(elements):
             if isinstance(l[i], dict):
                 for ang, mag in l[i].items():
@@ -312,6 +313,7 @@ def parse_projected_data(orbitals, species: Union[Sequence[Any], Dict[Any, List[
                                 l_count += 1
                         if m_count:
                             l_data[m_index] = data_temp
+                            totnum += 1
                     if l_count:
                         data[elem][l_index] = l_data
 
@@ -326,5 +328,6 @@ def parse_projected_data(orbitals, species: Union[Sequence[Any], Dict[Any, List[
                             count += 1
                     if count:
                         data[elem][l_index] = data_temp
+                        totnum += 1
 
-        return data
\ No newline at end of file
+        return data, totnum
\ No newline at end of file

From ddf2d2446656844c164719b60af0986c07df7345 Mon Sep 17 00:00:00 2001
From: jiyuang <jiyuyang@mail.ustc.com>
Date: Wed, 27 Apr 2022 16:24:54 +0800
Subject: [PATCH 04/12] add projected band structure ouput

---
 source/module_base/constants.h |   9 +
 source/src_io/energy_dos.cpp   | 304 +++++++++++++++++++++++++++++----
 2 files changed, 283 insertions(+), 30 deletions(-)

diff --git a/source/module_base/constants.h b/source/module_base/constants.h
index f916fdd0f4..cda66ff041 100644
--- a/source/module_base/constants.h
+++ b/source/module_base/constants.h
@@ -95,6 +95,15 @@ const double RYDBERG_SI =  HARTREE_SI/2.0; //J
 // zero up to a given accuracy
 //const double epsr  = 1.0e-6;
 const double threshold_wg  = 1.0e-14;
+
+const std::string Name_Angular[5][11] = 
+{
+    {"s"},
+    {"px", "py", "pz"},
+    {"d3z^2-r^2", "dxy", "dxz", "dx^2-y^2", "dyz"},
+    {"f5z^2-3r^2", "f5xz^2-xr^2", "f5yz^2-yr^2", "fzx^2-zy^2", "fxyz", "fx^3-3*xy^2", "f3yx^2-y^3"},
+    {"g1", "g2", "g3", "g4", "g5", "g6", "g7", "g8", "g9"}
+};
 }
 
 #endif 
diff --git a/source/src_io/energy_dos.cpp b/source/src_io/energy_dos.cpp
index bb89f7fe3b..551c2a04b1 100644
--- a/source/src_io/energy_dos.cpp
+++ b/source/src_io/energy_dos.cpp
@@ -461,35 +461,8 @@ void energy::perform_dos(Local_Orbital_wfc &lowf, LCAO_Hamilt &uhm)
 			 out.close();
 		 }
 
-		 std::string Name_Angular[5][11];
 		 /* decomposed Mulliken charge */
 
-		 Name_Angular[0][0] = "s          ";
-		 Name_Angular[1][0] = "px         ";
-		 Name_Angular[1][1] = "py         ";
-		 Name_Angular[1][2] = "pz         ";
-		 Name_Angular[2][0] = "d3z^2-r^2  ";
-		 Name_Angular[2][1] = "dxy        ";
-		 Name_Angular[2][2] = "dxz        ";
-		 Name_Angular[2][3] = "dx^2-y^2   ";
-		 Name_Angular[2][4] = "dyz        ";
-		 Name_Angular[3][0] = "f5z^2-3r^2 ";
-		 Name_Angular[3][1] = "f5xz^2-xr^2";
-		 Name_Angular[3][2] = "f5yz^2-yr^2";
-		 Name_Angular[3][3] = "fzx^2-zy^2 ";
-		 Name_Angular[3][4] = "fxyz       ";
-		 Name_Angular[3][5] = "fx^3-3*xy^2";
-		 Name_Angular[3][6] = "f3yx^2-y^3 ";
-		 Name_Angular[4][0] = "g1         ";
-		 Name_Angular[4][1] = "g2         ";
-		 Name_Angular[4][2] = "g3         ";
-		 Name_Angular[4][3] = "g4         ";
-		 Name_Angular[4][4] = "g5         ";
-		 Name_Angular[4][5] = "g6         ";
-		 Name_Angular[4][6] = "g7         ";
-		 Name_Angular[4][7] = "g8         ";
-		 Name_Angular[4][8] = "g9         ";
-
 		 {std::stringstream as;
 			 as << GlobalV::global_out_dir << "PDOS";
 			 std::ofstream out(as.str().c_str());
@@ -580,7 +553,7 @@ void energy::perform_dos(Local_Orbital_wfc &lowf, LCAO_Hamilt &uhm)
 					 const int m1 = atom1->iw2m[j];
 					 out <<std::setw(5) << i << std::setw(8) 
 						<< GlobalC::ucell.atoms[t].label <<std::setw(5)
-							<<L1<<std::setw(5) <<m1<<std::setw(5)<<N1+1<<std::setw(15)<< Name_Angular[L1][m1] << std::endl;
+							<<L1<<std::setw(5) <<m1<<std::setw(5)<<N1+1<<std::setw(15)<< ModuleBase::Name_Angular[L1][m1] << std::endl;
 				 }
 			 }
 			 out <<std::endl<<std::endl;
@@ -669,15 +642,286 @@ void energy::perform_dos(Local_Orbital_wfc &lowf, LCAO_Hamilt &uhm)
 			nks = GlobalC::kv.nkstot/2;
 		}
 
+		ModuleBase::ComplexMatrix weightk;
+		ModuleBase::matrix weight;
+		int NUM = 0;
+		if(GlobalV::GAMMA_ONLY_LOCAL)
+		{
+			NUM=GlobalV::NLOCAL*GlobalV::NBANDS*nspin0;
+			weightk.create(nspin0, GlobalV::NBANDS*GlobalV::NLOCAL, true);
+			weight.create(nspin0, GlobalV::NBANDS*GlobalV::NLOCAL, true);
+		}
+		else
+		{
+			NUM=GlobalV::NLOCAL*GlobalV::NBANDS*GlobalC::kv.nks;
+			weightk.create(GlobalC::kv.nks, GlobalV::NBANDS*GlobalV::NLOCAL, true);
+			weight.create(GlobalC::kv.nks, GlobalV::NBANDS*GlobalV::NLOCAL, true);
+		}
+
 		for(int is=0; is<nspin0; is++)
 		{
 			std::stringstream ss2;
 			ss2 << GlobalV::global_out_dir << "BANDS_" << is+1 << ".dat";
 			GlobalV::ofs_running << "\n Output bands in file: " << ss2.str() << std::endl;
 			Dos::nscf_band(is, ss2.str(), nks, GlobalV::NBANDS, this->ef*0, GlobalC::wf.ekb);
-		}
 
-	}
+			// Projeced band structure added by jiyy-2022-4-20
+			if(GlobalV::GAMMA_ONLY_LOCAL)
+			{
+				std::vector<ModuleBase::matrix>   Mulk;
+				Mulk.resize(1);
+				Mulk[0].create(pv->ncol,pv->nrow);
+
+
+				ModuleBase::matrix Dwf = lowf.wfc_gamma[is];
+				for (int i=0; i<GlobalV::NBANDS; ++i)		  
+				{  
+					const int NB= i+1;
+
+					const double one_float=1.0, zero_float=0.0;
+					const int one_int=1;
+
+				#ifdef __MPI
+					const char T_char='T';		
+					pdgemv_(
+							&T_char,
+							&GlobalV::NLOCAL,&GlobalV::NLOCAL,
+							&one_float,
+							uhm.LM->Sloc.data(), &one_int, &one_int, pv->desc,
+							Dwf.c, &one_int, &NB, pv->desc, &one_int,
+							&zero_float,
+							Mulk[0].c, &one_int, &NB, pv->desc,
+							&one_int);
+				#endif
+
+					for (int j=0; j<GlobalV::NLOCAL; ++j)
+					{
+
+						if ( pv->in_this_processor(j,i) )
+						{
+
+							const int ir = pv->trace_loc_row[j];
+							const int ic = pv->trace_loc_col[i];
+							weightk(is, i*GlobalV::NLOCAL+j) = Mulk[0](ic,ir)*lowf.wfc_gamma[is](ic,ir);
+						}
+					} 
+				}//ib
+			}//if
+			else
+			{
+				GlobalV::SEARCH_RADIUS = atom_arrange::set_sr_NL(
+					GlobalV::ofs_running,
+					GlobalV::OUT_LEVEL,
+					GlobalC::ORB.get_rcutmax_Phi(), 
+					GlobalC::ucell.infoNL.get_rcutmax_Beta(), 
+					GlobalV::GAMMA_ONLY_LOCAL);
+
+				atom_arrange::search(
+					GlobalV::SEARCH_PBC,
+					GlobalV::ofs_running,
+					GlobalC::GridD, 
+					GlobalC::ucell, 
+					GlobalV::SEARCH_RADIUS, 
+					GlobalV::test_atom_input);//qifeng-2019-01-21
+
+				uhm.LM->allocate_HS_R(pv->nnr);
+				uhm.LM->zeros_HSR('S');
+				uhm.genH.calculate_S_no();
+				uhm.genH.build_ST_new('S', false, GlobalC::ucell);
+				std::vector<ModuleBase::ComplexMatrix> Mulk;
+				Mulk.resize(1);
+				Mulk[0].create(pv->ncol,pv->nrow);
+
+
+				for(int ik=0;ik<GlobalC::kv.nks;ik++)
+				{
+
+					if(is == GlobalC::kv.isk[ik])
+					{
+						uhm.LM->allocate_HS_k(pv->nloc);
+						uhm.LM->zeros_HSk('S');
+						uhm.LM->folding_fixedH(ik);
+
+
+						ModuleBase::ComplexMatrix Dwfc = conj(lowf.wfc_k[ik]);
+
+						for (int i=0; i<GlobalV::NBANDS; ++i)		  
+						{     
+							const int NB= i+1;
+
+							const double one_float[2]={1.0, 0.0}, zero_float[2]={0.0, 0.0};
+							const int one_int=1;
+							//   const int two_int=2;
+							const char T_char='T';		// N_char='N',U_char='U'
+
+						#ifdef __MPI
+							pzgemv_(
+									&T_char,
+									&GlobalV::NLOCAL,&GlobalV::NLOCAL,
+									&one_float[0],
+									uhm.LM->Sloc2.data(), &one_int, &one_int, pv->desc,
+									Dwfc.c, &one_int, &NB, pv->desc, &one_int,
+									&zero_float[0],
+									Mulk[0].c, &one_int, &NB, pv->desc,
+									&one_int);
+						#endif
+
+
+							for (int j=0; j<GlobalV::NLOCAL; ++j)
+							{
+
+								if ( pv->in_this_processor(j,i) )
+								{
+
+									const int ir = pv->trace_loc_row[j];
+									const int ic = pv->trace_loc_col[i];
+
+									weightk(ik, i*GlobalV::NLOCAL+j) = Mulk[0](ic,ir)*lowf.wfc_k[ik](ic,ir);
+								}
+							}                             
+
+						}//ib
+
+					}//if                       
+				}//ik
+#ifdef __MPI
+				atom_arrange::delete_vector(
+					GlobalV::ofs_running,
+					GlobalV::SEARCH_PBC, 
+					GlobalC::GridD, 
+					GlobalC::ucell, 
+					GlobalV::SEARCH_RADIUS, 
+					GlobalV::test_atom_input);
+#endif
+			}//else
+		#ifdef __MPI
+		MPI_Reduce(weightk.real().c, weight.c , NUM , MPI_DOUBLE , MPI_SUM, 0, MPI_COMM_WORLD);
+		#endif
+
+		if(GlobalV::MY_RANK == 0)
+		{
+			std::stringstream ps2;
+			ps2 << GlobalV::global_out_dir << "PBANDS_" << is+1;
+			GlobalV::ofs_running << "\n Output projected bands in file: " << ps2.str() << std::endl;
+			std::ofstream out(ps2.str().c_str());
+
+			out << "<"<<"pband"<<">" <<std::endl;
+			out << "<"<<"nspin"<<">" << GlobalV::NSPIN<< "<"<<"/"<<"nspin"<<">"<< std::endl;
+			if (GlobalV::NSPIN==4)
+				out << "<"<<"norbitals"<<">" <<std::setw(2) <<GlobalV::NLOCAL/2<< "<"<<"/"<<"norbitals"<<">"<< std::endl;
+			else
+				out << "<"<<"norbitals"<<">" <<std::setw(2) <<GlobalV::NLOCAL<< "<"<<"/"<<"norbitals"<<">"<< std::endl;
+			out << "<"<<"band_structure nkpoints="<<"\""<<nks<<"\""<<" nbands="<<"\""<<GlobalV::NBANDS<<"\""<<" units=\"eV\""<<">" <<std::endl;
+			if(GlobalV::GAMMA_ONLY_LOCAL)
+			{
+				for(int ib = 0; ib < GlobalV::NBANDS; ib++)
+						out << " " << (GlobalC::wf.ekb[is*nks][ib]) * ModuleBase::Ry_to_eV;
+				out << std::endl;
+			}
+			else
+			{
+				for(int ik=0;ik<nks;ik++)
+				{
+					for(int ib = 0; ib < GlobalV::NBANDS; ib++)
+						out << " " << (GlobalC::wf.ekb[ik+is*nks][ib]) * ModuleBase::Ry_to_eV;
+					out << std::endl;
+				}
+			}
+			out << "<"<<"/"<<"band_structure"<<">"<<std::endl;
+
+			for (int i=0; i<GlobalC::ucell.nat; i++)
+			{   
+				int a = GlobalC::ucell.iat2ia[i];
+				int t = GlobalC::ucell.iat2it[i];
+				Atom* atom1 = &GlobalC::ucell.atoms[t];
+				const int s0 = GlobalC::ucell.itiaiw2iwt(t, a, 0);
+				for(int j=0; j<atom1->nw; ++j)
+				{
+					const int L1 = atom1->iw2l[j];
+					const int N1 = atom1->iw2n[j];
+					const int m1 = atom1->iw2m[j];
+					const int w = GlobalC::ucell.itiaiw2iwt(t, a, j);
+
+					//out << "<"<<"/"<<"energy"<<"_"<<"values"<<">" <<std::endl;
+					out << "<"<<"orbital" <<std::endl;
+					out <<std::setw(6)<< "index"<<"="<<"\""<<std::setw(40) <<w+1<<"\""<<std::endl;
+					out <<std::setw(5)<< "atom"<<"_"<<"index"<<"="<<"\""<<std::setw(40) <<i+1<<"\""<<std::endl;
+					out <<std::setw(8)<< "species"<<"="<<"\""<<GlobalC::ucell.atoms[t].label<<"\""<<std::endl;
+					out<<std::setw(2)<< "l"<<"="<<"\""<<std::setw(40)<<L1<<"\""<<std::endl;
+					out <<std::setw(2)<< "m"<<"="<<"\""<<std::setw(40)<<m1<<"\""<<std::endl;
+					out <<std::setw(2)<< "z"<<"="<<"\""<<std::setw(40)<<N1+1<<"\""<<std::endl;
+					out << ">" <<std::endl;
+					out << "<"<<"data"<<">" <<std::endl;
+					if(GlobalV::GAMMA_ONLY_LOCAL)
+					{
+						for(int ib=0;ib<GlobalV::NBANDS;ib++)
+						{
+							if (GlobalV::NSPIN==1 || GlobalV::NSPIN==2)
+								out <<std::setw(13)<< weight(is, ib*GlobalV::NLOCAL+w);
+							else if (GlobalV::NSPIN==4)
+							{
+								int w0 = w-s0;
+								out <<std::setw(13)<< weight(is, ib*GlobalV::NLOCAL+s0+2*w0)+weight(is, ib*GlobalV::NLOCAL+s0+2*w0+1);
+							}
+						}
+						out <<std::endl;
+					}
+					else
+					{
+						for(int ik=0;ik<nks;ik++)
+						{
+							for(int ib=0;ib<GlobalV::NBANDS;ib++)
+							{
+								if (GlobalV::NSPIN==1)
+									out <<std::setw(13)<< weight(ik, ib*GlobalV::NLOCAL+w);
+								else if(GlobalV::NSPIN==2)
+									out <<std::setw(13)<< weight(ik+nks*is, ib*GlobalV::NLOCAL+w);
+								else if (GlobalV::NSPIN==4)
+								{
+									int w0 = w-s0;
+									out <<std::setw(13)<< weight(ik, ib*GlobalV::NLOCAL+s0+2*w0)+weight(ik, ib*GlobalV::NLOCAL+s0+2*w0+1);
+								}
+							}
+							out <<std::endl;
+						}
+					}
+					out << "<"<<"/"<<"data"<<">" <<std::endl;
+					out << "<"<<"/"<<"orbital"<<">" <<std::endl;
+				}//j
+			}//i
+
+			out << "<"<<"/"<<"pband"<<">" <<std::endl;
+			out.close();}
+		}//is
+		{  std::stringstream os;
+			 os<<GlobalV::global_out_dir<<"Orbital";
+			 std::ofstream out(os.str().c_str());
+			 out<< std::setw(5)<<"io"<< std::setw(8) <<"spec" <<std::setw(5)<<"l"<<std::setw(5)<<"m"<<std::setw(5)<<"z"<<std::setw(5)<<"sym"<<std::endl;
+
+
+			 for (int i=0; i<GlobalC::ucell.nat; i++)
+			 {
+				 int   t = GlobalC::ucell.iat2it[i];
+				 Atom* atom1 = &GlobalC::ucell.atoms[t];  
+				 for(int j=0; j<atom1->nw; ++j)
+				 {
+					 const int L1 = atom1->iw2l[j];
+					 const int N1 = atom1->iw2n[j];
+					 const int m1 = atom1->iw2m[j];
+					 out <<std::setw(5) << i << std::setw(8) 
+						<< GlobalC::ucell.atoms[t].label <<std::setw(5)
+							<<L1<<std::setw(5) <<m1<<std::setw(5)<<N1+1<<std::setw(15)<< ModuleBase::Name_Angular[L1][m1] << std::endl;
+				 }
+			 }
+			 out <<std::endl<<std::endl;
+			 out <<std::setw(5)<< "io"<<std::setw(2)<<"="<<std::setw(2)<<"Orbital index in supercell"<<std::endl;
+			 out <<std::setw(5)<< "spec"<<std::setw(2)<<"="<<std::setw(2)<<"Atomic species label"<<std::endl;
+			 out <<std::setw(5)<< "l"<<std::setw(2)<<"="<<std::setw(2)<<"Angular mumentum quantum number"<<std::endl;
+			 out <<std::setw(5)<< "m"<<std::setw(2)<<"="<<std::setw(2)<<"Magnetic quantum number"<<std::endl;
+			 out <<std::setw(5)<< "z"<<std::setw(2)<<"="<<std::setw(2)<<"Zeta index of orbital"<<std::endl;
+			 out <<std::setw(5)<< "sym"<<std::setw(2)<<"="<<std::setw(2)<<"Symmetry name of real orbital"<<std::endl;
+			 out.close();}
+		}//out_band
 	return;
 }
 

From 3039557f9671f66140da974e25ea71d1a1e696db Mon Sep 17 00:00:00 2001
From: jiyuang <jiyuyang@mail.ustc.com>
Date: Wed, 27 Apr 2022 16:24:54 +0800
Subject: [PATCH 05/12] modify README for plot-tools

---
 tools/plot-tools/README.md           | 11 ++--
 tools/plot-tools/abacus_plot/band.py | 12 +++--
 tools/plot-tools/abacus_plot/main.py | 80 +++++++++++++++++++---------
 3 files changed, 69 insertions(+), 34 deletions(-)

diff --git a/tools/plot-tools/README.md b/tools/plot-tools/README.md
index 97db50fb4d..04eb088124 100644
--- a/tools/plot-tools/README.md
+++ b/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
diff --git a/tools/plot-tools/abacus_plot/band.py b/tools/plot-tools/abacus_plot/band.py
index 41d30b52e3..5b867f0ef7 100644
--- a/tools/plot-tools/abacus_plot/band.py
+++ b/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 
diff --git a/tools/plot-tools/abacus_plot/main.py b/tools/plot-tools/abacus_plot/main.py
index c90293755b..aea7bef259 100644
--- a/tools/plot-tools/abacus_plot/main.py
+++ b/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()

From c7740d34a1e65b3f6e87b53a0f5ec5a1dac72de6 Mon Sep 17 00:00:00 2001
From: jiyuyang <1041176461@qq.com>
Date: Wed, 27 Apr 2022 22:54:31 +0800
Subject: [PATCH 06/12] Update band-struc.md

add description for projected band structure
---
 docs/examples/band-struc.md | 31 ++++++++++++++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/docs/examples/band-struc.md b/docs/examples/band-struc.md
index 43791d5710..ad3354401a 100644
--- a/docs/examples/band-struc.md
+++ b/docs/examples/band-struc.md
@@ -58,4 +58,33 @@ points.
 Run the program, and you will see a file named BANDS_1.dat in the output directory. Plot it
 to get energy band structure.
 
-[back to top](#band-structure)
\ No newline at end of file
+Along with the BANDS_1.dat file, we also produce the projected band structure in a file called PBAND_1 in xml format.
+
+The PDOS file starts with number of atomic orbitals in the system, the text contents of element <band structure> is the same as data in the BANDS_1.dat file, such as:
+```
+<pband>
+<nspin>1</nspin>
+<norbitals>153</norbitals>
+<band_structure nkpoints="96" nbands="50" units="eV">
+...
+
+```
+
+The rest of the files arranged in sections, each section with a header such as below:
+    
+```
+<orbital
+ index="                                       1"
+ atom_index="                                       1"
+ species="Si"
+ l="                                       0"
+ m="                                       0"
+ z="                                       1"
+>
+<data>
+...
+</data>
+
+The shape of text contents of element <data> is (Number of k-points, Number of bands) 
+ 
+[back to top](#band-structure)

From aaaa8c9979c3d444a9be93b9ec2cf523af1ed5d7 Mon Sep 17 00:00:00 2001
From: jiyuyang <1041176461@qq.com>
Date: Wed, 27 Apr 2022 22:55:57 +0800
Subject: [PATCH 07/12] Update band-struc.md

---
 docs/examples/band-struc.md | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/docs/examples/band-struc.md b/docs/examples/band-struc.md
index ad3354401a..cc5dbd4b94 100644
--- a/docs/examples/band-struc.md
+++ b/docs/examples/band-struc.md
@@ -85,6 +85,8 @@ The rest of the files arranged in sections, each section with a header such as b
 ...
 </data>
 
+```
+
 The shape of text contents of element <data> is (Number of k-points, Number of bands) 
  
 [back to top](#band-structure)

From 6f9f9f706804e536d039484086175d2c41eb7813 Mon Sep 17 00:00:00 2001
From: jiyuang <jiyuyang@mail.ustc.com>
Date: Fri, 6 May 2022 10:00:20 +0800
Subject: [PATCH 08/12] fix energy_dos.cpp compiling bug

---
 source/src_io/energy_dos.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/source/src_io/energy_dos.cpp b/source/src_io/energy_dos.cpp
index 42f3deeb2f..95503863d8 100644
--- a/source/src_io/energy_dos.cpp
+++ b/source/src_io/energy_dos.cpp
@@ -726,8 +726,8 @@ void energy::perform_dos(Local_Orbital_wfc &lowf, LCAO_Hamilt &uhm)
 
 				uhm.LM->allocate_HS_R(pv->nnr);
 				uhm.LM->zeros_HSR('S');
-				uhm.genH.calculate_S_no();
-				uhm.genH.build_ST_new('S', false, GlobalC::ucell);
+				uhm.genH.calculate_S_no(uhm.LM->SlocR.data());
+				uhm.genH.build_ST_new('S', false, GlobalC::ucell, uhm.LM->SlocR.data());
 				std::vector<ModuleBase::ComplexMatrix> Mulk;
 				Mulk.resize(1);
 				Mulk[0].create(pv->ncol,pv->nrow);

From 0a9da9cd01d6c0cd7a68f8f5361fa684df300241 Mon Sep 17 00:00:00 2001
From: jiyuang <jiyuyang@mail.ustc.com>
Date: Sat, 7 May 2022 21:41:36 +0800
Subject: [PATCH 09/12] set default value of 'kkbeta' with 'size'

---
 source/module_cell/read_pp_upf201.cpp | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/source/module_cell/read_pp_upf201.cpp b/source/module_cell/read_pp_upf201.cpp
index 50bb7c2cc2..ecc894b700 100644
--- a/source/module_cell/read_pp_upf201.cpp
+++ b/source/module_cell/read_pp_upf201.cpp
@@ -197,7 +197,9 @@ int Pseudopot_upf::read_pseudo_upf201(std::ifstream &ifs)
 		for(int ip = 0 ; ip < nparameter; ++ip)
 		{
 			if(name[ip]=="type"){}
-			else if(name[ip]=="size"){}
+			else if(name[ip]=="size"){
+				kkbeta[ib] = atoi (val[ip].c_str());
+			}
 			else if(name[ip]=="columns"){}
 			else if(name[ip]=="index"){}
 			else if(name[ip]=="label"){}

From b11ff731caa5604a7668749293c5873c3e4f6638 Mon Sep 17 00:00:00 2001
From: jiyuang <jiyuyang@mail.ustc.com>
Date: Sat, 7 May 2022 22:38:13 +0800
Subject: [PATCH 10/12] Revert "set default value of 'kkbeta' with 'size'"

This reverts commit 0a9da9cd01d6c0cd7a68f8f5361fa684df300241.
---
 source/module_cell/read_pp_upf201.cpp | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/source/module_cell/read_pp_upf201.cpp b/source/module_cell/read_pp_upf201.cpp
index ecc894b700..50bb7c2cc2 100644
--- a/source/module_cell/read_pp_upf201.cpp
+++ b/source/module_cell/read_pp_upf201.cpp
@@ -197,9 +197,7 @@ int Pseudopot_upf::read_pseudo_upf201(std::ifstream &ifs)
 		for(int ip = 0 ; ip < nparameter; ++ip)
 		{
 			if(name[ip]=="type"){}
-			else if(name[ip]=="size"){
-				kkbeta[ib] = atoi (val[ip].c_str());
-			}
+			else if(name[ip]=="size"){}
 			else if(name[ip]=="columns"){}
 			else if(name[ip]=="index"){}
 			else if(name[ip]=="label"){}

From 6e479cad77ec2fc74af8e1ba349b7602b418c530 Mon Sep 17 00:00:00 2001
From: jiyuang <jiyuyang@mail.ustc.com>
Date: Wed, 11 May 2022 09:52:17 +0800
Subject: [PATCH 11/12] move Name_Angular from constant.h to energy.h

---
 source/module_base/constants.h | 9 ---------
 source/src_io/energy_dos.cpp   | 4 ++--
 source/src_pw/energy.h         | 8 ++++++++
 3 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/source/module_base/constants.h b/source/module_base/constants.h
index cda66ff041..f916fdd0f4 100644
--- a/source/module_base/constants.h
+++ b/source/module_base/constants.h
@@ -95,15 +95,6 @@ const double RYDBERG_SI =  HARTREE_SI/2.0; //J
 // zero up to a given accuracy
 //const double epsr  = 1.0e-6;
 const double threshold_wg  = 1.0e-14;
-
-const std::string Name_Angular[5][11] = 
-{
-    {"s"},
-    {"px", "py", "pz"},
-    {"d3z^2-r^2", "dxy", "dxz", "dx^2-y^2", "dyz"},
-    {"f5z^2-3r^2", "f5xz^2-xr^2", "f5yz^2-yr^2", "fzx^2-zy^2", "fxyz", "fx^3-3*xy^2", "f3yx^2-y^3"},
-    {"g1", "g2", "g3", "g4", "g5", "g6", "g7", "g8", "g9"}
-};
 }
 
 #endif 
diff --git a/source/src_io/energy_dos.cpp b/source/src_io/energy_dos.cpp
index 95503863d8..86297835ee 100644
--- a/source/src_io/energy_dos.cpp
+++ b/source/src_io/energy_dos.cpp
@@ -553,7 +553,7 @@ void energy::perform_dos(Local_Orbital_wfc &lowf, LCAO_Hamilt &uhm)
 					 const int m1 = atom1->iw2m[j];
 					 out <<std::setw(5) << i << std::setw(8) 
 						<< GlobalC::ucell.atoms[t].label <<std::setw(5)
-							<<L1<<std::setw(5) <<m1<<std::setw(5)<<N1+1<<std::setw(15)<< ModuleBase::Name_Angular[L1][m1] << std::endl;
+							<<L1<<std::setw(5) <<m1<<std::setw(5)<<N1+1<<std::setw(15)<< GlobalC::en.Name_Angular[L1][m1] << std::endl;
 				 }
 			 }
 			 out <<std::endl<<std::endl;
@@ -910,7 +910,7 @@ void energy::perform_dos(Local_Orbital_wfc &lowf, LCAO_Hamilt &uhm)
 					 const int m1 = atom1->iw2m[j];
 					 out <<std::setw(5) << i << std::setw(8) 
 						<< GlobalC::ucell.atoms[t].label <<std::setw(5)
-							<<L1<<std::setw(5) <<m1<<std::setw(5)<<N1+1<<std::setw(15)<< ModuleBase::Name_Angular[L1][m1] << std::endl;
+							<<L1<<std::setw(5) <<m1<<std::setw(5)<<N1+1<<std::setw(15)<< GlobalC::en.Name_Angular[L1][m1] << std::endl;
 				 }
 			 }
 			 out <<std::endl<<std::endl;
diff --git a/source/src_pw/energy.h b/source/src_pw/energy.h
index d9ea8be4bd..bbc0ae7609 100644
--- a/source/src_pw/energy.h
+++ b/source/src_pw/energy.h
@@ -58,6 +58,14 @@ class energy
 
 	int out_dos;			// control dos calculation
 	int out_band;                    // control band calculation  pengfei 2014-10-13
+	const std::string Name_Angular[5][11] = 
+	{
+    	{"s"},
+    	{"px", "py", "pz"},
+    	{"d3z^2-r^2", "dxy", "dxz", "dx^2-y^2", "dyz"},
+    	{"f5z^2-3r^2", "f5xz^2-xr^2", "f5yz^2-yr^2", "fzx^2-zy^2", "fxyz", "fx^3-3*xy^2", "f3yx^2-y^3"},
+    	{"g1", "g2", "g3", "g4", "g5", "g6", "g7", "g8", "g9"}
+	};          // name of atomic orbital    jiyy add 2022-05-10
 
 	double dos_emin_ev;
 	double dos_emax_ev;

From bd415d234d6af933ebf35c8f19981660c21d372e Mon Sep 17 00:00:00 2001
From: jiyuang <jiyuyang@mail.ustc.com>
Date: Wed, 11 May 2022 21:34:07 +0800
Subject: [PATCH 12/12] add a new parameter named "out_proj_band"

---
 docs/examples/band-struc.md   |   5 +-
 docs/input-main.md            |   8 ++-
 source/input.cpp              |   8 +++
 source/input.h                |   1 +
 source/input_conv.cpp         |   1 +
 source/src_io/energy_dos.cpp  | 122 ++++++++++++++++------------------
 source/src_io/write_input.cpp |   1 +
 source/src_pw/energy.h        |   3 +
 8 files changed, 83 insertions(+), 66 deletions(-)

diff --git a/docs/examples/band-struc.md b/docs/examples/band-struc.md
index cc5dbd4b94..f1a31fe814 100644
--- a/docs/examples/band-struc.md
+++ b/docs/examples/band-struc.md
@@ -27,6 +27,7 @@ pw_diag_thr 1.0e-7
 #Parameters (File)
 init_chg file
 out_band 1
+out_proj_band 1
 
 #Parameters (Smearing)
 smearing_method gaussian
@@ -58,9 +59,9 @@ points.
 Run the program, and you will see a file named BANDS_1.dat in the output directory. Plot it
 to get energy band structure.
 
-Along with the BANDS_1.dat file, we also produce the projected band structure in a file called PBAND_1 in xml format.
+If "out_proj_band" set 1, it will also produce the projected band structure in a file called PBAND_1 in xml format.
 
-The PDOS file starts with number of atomic orbitals in the system, the text contents of element <band structure> is the same as data in the BANDS_1.dat file, such as:
+The PBAND_1 file starts with number of atomic orbitals in the system, the text contents of element <band structure> is the same as data in the BANDS_1.dat file, such as:
 ```
 <pband>
 <nspin>1</nspin>
diff --git a/docs/input-main.md b/docs/input-main.md
index 99717ce4ad..1a75667352 100644
--- a/docs/input-main.md
+++ b/docs/input-main.md
@@ -32,7 +32,7 @@
 
 - [Variables related to output information](#variables-related-to-output-information)
 
-    [out_force](#out_force) | [out_mul](#out_mul) | [out_chg](#out_chg) | [out_pot](#out_pot) | [out_dm](#out-dm) | [out_wfc_pw](#out_wfc_pw) | [out_wfc_r](#out_wfc_r) | [out_wfc_lcao](#out_wfc_lcao) | [out_dos](#out-dos) | [out_band](#out-band) | [out_stru](#out-stru) | [out_level](#out_level) | [out_alllog](#out-alllog) | [out_mat_hs](#out_mat_hs) | [out_mat_r](#out_mat_r) | [out_mat_hs2](#out_mat_hs2) | [out_element_info](#out-element-info) | [restart_save](#restart_save) | [restart_load](#restart_load)
+    [out_force](#out_force) | [out_mul](#out_mul) | [out_chg](#out_chg) | [out_pot](#out_pot) | [out_dm](#out-dm) | [out_wfc_pw](#out_wfc_pw) | [out_wfc_r](#out_wfc_r) | [out_wfc_lcao](#out_wfc_lcao) | [out_dos](#out-dos) | [out_band](#out-band) | [out_proj_band](#out-proj-band) | [out_stru](#out-stru) | [out_level](#out_level) | [out_alllog](#out-alllog) | [out_mat_hs](#out_mat_hs) | [out_mat_r](#out_mat_r) | [out_mat_hs2](#out_mat_hs2) | [out_element_info](#out-element-info) | [restart_save](#restart_save) | [restart_load](#restart_load)
 
 - [Density of states](#density-of-states)
 
@@ -723,6 +723,12 @@ This part of variables are used to control the output of properties.
 - **Description**: Controls whether to output the band structure. For mroe information, refer to the [worked example](examples/band-struc.md)
 - **Default**: 0
 
+#### out_proj_band
+
+- **Type**: Integer
+- **Description**: Controls whether to output the projected band structure. For mroe information, refer to the [worked example](examples/band-struc.md)
+- **Default**: 0
+
 #### out_stru
 
 - **Type**: Boolean
diff --git a/source/input.cpp b/source/input.cpp
index 50dfe3603c..09a3088fe5 100644
--- a/source/input.cpp
+++ b/source/input.cpp
@@ -266,6 +266,7 @@ void Input::Default(void)
     out_wfc_r = 0;
     out_dos = 0;
     out_band = 0;
+    out_proj_band = 0;
     out_mat_hs = 0;
     out_mat_hs2 = 0; // LiuXh add 2019-07-15
     out_mat_r = 0; // jingan add 2019-8-14
@@ -985,6 +986,10 @@ bool Input::Read(const std::string &fn)
         {
             read_value(ifs, out_band);
         }
+        else if (strcmp("out_proj_band", word) == 0)
+        {
+            read_value(ifs, out_proj_band);
+        }
 
         else if (strcmp("out_mat_hs", word) == 0)
         {
@@ -1929,6 +1934,7 @@ void Input::Bcast()
     Parallel_Common::bcast_int(out_wfc_r);
     Parallel_Common::bcast_int(out_dos);
     Parallel_Common::bcast_int(out_band);
+    Parallel_Common::bcast_int(out_proj_band);
     Parallel_Common::bcast_int(out_mat_hs);
     Parallel_Common::bcast_int(out_mat_hs2); // LiuXh add 2019-07-15
     Parallel_Common::bcast_int(out_mat_r); // jingan add 2019-8-14
@@ -2215,6 +2221,7 @@ void Input::Check(void)
         out_stru = 0;
         out_dos = 0;
         out_band = 0;
+        out_proj_band = 0;
         cal_force = 0;
         init_wfc = "file";
         init_chg = "atomic"; // useless,
@@ -2236,6 +2243,7 @@ void Input::Check(void)
         out_stru = 0;
         out_dos = 0;
         out_band = 0;
+        out_proj_band = 0;
         cal_force = 0;
         init_wfc = "file";
         init_chg = "atomic";
diff --git a/source/input.h b/source/input.h
index 50c43852bd..aa68a1fda6 100644
--- a/source/input.h
+++ b/source/input.h
@@ -204,6 +204,7 @@ class Input
     int out_wfc_r; // 0: no; 1: yes
     int out_dos; // dos calculation. mohan add 20090909
     int out_band; // band calculation pengfei 2014-10-13
+    int out_proj_band; // projected band structure calculation jiyy add 2022-05-11
     int out_mat_hs; // output H matrix and S matrix in local basis.
     int out_mat_hs2; // LiuXh add 2019-07-16, output H(R) matrix and S(R) matrix in local basis.
     int out_mat_r; // jingan add 2019-8-14, output r(R) matrix.
diff --git a/source/input_conv.cpp b/source/input_conv.cpp
index 8a9fa8ec42..7427e0abfe 100644
--- a/source/input_conv.cpp
+++ b/source/input_conv.cpp
@@ -432,6 +432,7 @@ void Input_Conv::Convert(void)
     GlobalC::wf.out_wfc_r = INPUT.out_wfc_r;
     GlobalC::en.out_dos = INPUT.out_dos;
     GlobalC::en.out_band = INPUT.out_band;
+    GlobalC::en.out_proj_band = INPUT.out_proj_band;
 #ifdef __LCAO
     Local_Orbital_Charge::out_dm = INPUT.out_dm;
     Pdiag_Double::out_mat_hs = INPUT.out_mat_hs;
diff --git a/source/src_io/energy_dos.cpp b/source/src_io/energy_dos.cpp
index 86297835ee..85884cc78a 100644
--- a/source/src_io/energy_dos.cpp
+++ b/source/src_io/energy_dos.cpp
@@ -29,7 +29,7 @@ void energy::perform_dos(Local_Orbital_wfc &lowf, LCAO_Hamilt &uhm)
 
     const Parallel_Orbitals* pv = uhm.LM->ParaV;
 
-    if(out_dos !=0 || out_band !=0)
+    if(out_dos !=0 || out_band !=0 || out_proj_band !=0)
     {
         GlobalV::ofs_running << "\n\n\n\n";
         GlobalV::ofs_running << " >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>" << std::endl;
@@ -536,34 +536,7 @@ void energy::perform_dos(Local_Orbital_wfc &lowf, LCAO_Hamilt &uhm)
 
 			 out << "<"<<"/"<<"pdos"<<">" <<std::endl;
 			 out.close();}
-		 {  std::stringstream os;
-			 os<<GlobalV::global_out_dir<<"Orbital";
-			 std::ofstream out(os.str().c_str());
-			 out<< std::setw(5)<<"io"<< std::setw(8) <<"spec" <<std::setw(5)<<"l"<<std::setw(5)<<"m"<<std::setw(5)<<"z"<<std::setw(5)<<"sym"<<std::endl;
-
-
-			 for (int i=0; i<GlobalC::ucell.nat; i++)
-			 {
-				 int   t = GlobalC::ucell.iat2it[i];
-				 Atom* atom1 = &GlobalC::ucell.atoms[t];  
-				 for(int j=0; j<atom1->nw; ++j)
-				 {
-					 const int L1 = atom1->iw2l[j];
-					 const int N1 = atom1->iw2n[j];
-					 const int m1 = atom1->iw2m[j];
-					 out <<std::setw(5) << i << std::setw(8) 
-						<< GlobalC::ucell.atoms[t].label <<std::setw(5)
-							<<L1<<std::setw(5) <<m1<<std::setw(5)<<N1+1<<std::setw(15)<< GlobalC::en.Name_Angular[L1][m1] << std::endl;
-				 }
-			 }
-			 out <<std::endl<<std::endl;
-			 out <<std::setw(5)<< "io"<<std::setw(2)<<"="<<std::setw(2)<<"Orbital index in supercell"<<std::endl;
-			 out <<std::setw(5)<< "spec"<<std::setw(2)<<"="<<std::setw(2)<<"Atomic species label"<<std::endl;
-			 out <<std::setw(5)<< "l"<<std::setw(2)<<"="<<std::setw(2)<<"Angular mumentum quantum number"<<std::endl;
-			 out <<std::setw(5)<< "m"<<std::setw(2)<<"="<<std::setw(2)<<"Magnetic quantum number"<<std::endl;
-			 out <<std::setw(5)<< "z"<<std::setw(2)<<"="<<std::setw(2)<<"Zeta index of orbital"<<std::endl;
-			 out <<std::setw(5)<< "sym"<<std::setw(2)<<"="<<std::setw(2)<<"Symmetry name of real orbital"<<std::endl;
-			 out.close();}
+		this->print_orbital_file();
 
 	 }       
 	 delete[] pdos;
@@ -642,6 +615,28 @@ void energy::perform_dos(Local_Orbital_wfc &lowf, LCAO_Hamilt &uhm)
 			nks = GlobalC::kv.nkstot/2;
 		}
 
+		for(int is=0; is<nspin0; is++)
+		{
+			std::stringstream ss2;
+			ss2 << GlobalV::global_out_dir << "BANDS_" << is+1 << ".dat";
+			GlobalV::ofs_running << "\n Output bands in file: " << ss2.str() << std::endl;
+			Dos::nscf_band(is, ss2.str(), nks, GlobalV::NBANDS, this->ef*0, GlobalC::wf.ekb);
+
+		}
+	}//out_band
+
+	if(this->out_proj_band) // Projeced band structure added by jiyy-2022-4-20
+	{
+		int nks=0;
+		if(nspin0==1) 
+		{
+			nks = GlobalC::kv.nkstot;
+		}
+		else if(nspin0==2) 
+		{
+			nks = GlobalC::kv.nkstot/2;
+		}
+
 		ModuleBase::ComplexMatrix weightk;
 		ModuleBase::matrix weight;
 		int NUM = 0;
@@ -660,12 +655,7 @@ void energy::perform_dos(Local_Orbital_wfc &lowf, LCAO_Hamilt &uhm)
 
 		for(int is=0; is<nspin0; is++)
 		{
-			std::stringstream ss2;
-			ss2 << GlobalV::global_out_dir << "BANDS_" << is+1 << ".dat";
-			GlobalV::ofs_running << "\n Output bands in file: " << ss2.str() << std::endl;
-			Dos::nscf_band(is, ss2.str(), nks, GlobalV::NBANDS, this->ef*0, GlobalC::wf.ekb);
 
-			// Projeced band structure added by jiyy-2022-4-20
 			if(GlobalV::GAMMA_ONLY_LOCAL)
 			{
 				std::vector<ModuleBase::matrix>   Mulk;
@@ -893,35 +883,41 @@ void energy::perform_dos(Local_Orbital_wfc &lowf, LCAO_Hamilt &uhm)
 			out << "<"<<"/"<<"pband"<<">" <<std::endl;
 			out.close();}
 		}//is
-		{  std::stringstream os;
-			 os<<GlobalV::global_out_dir<<"Orbital";
-			 std::ofstream out(os.str().c_str());
-			 out<< std::setw(5)<<"io"<< std::setw(8) <<"spec" <<std::setw(5)<<"l"<<std::setw(5)<<"m"<<std::setw(5)<<"z"<<std::setw(5)<<"sym"<<std::endl;
-
-
-			 for (int i=0; i<GlobalC::ucell.nat; i++)
-			 {
-				 int   t = GlobalC::ucell.iat2it[i];
-				 Atom* atom1 = &GlobalC::ucell.atoms[t];  
-				 for(int j=0; j<atom1->nw; ++j)
-				 {
-					 const int L1 = atom1->iw2l[j];
-					 const int N1 = atom1->iw2n[j];
-					 const int m1 = atom1->iw2m[j];
-					 out <<std::setw(5) << i << std::setw(8) 
-						<< GlobalC::ucell.atoms[t].label <<std::setw(5)
-							<<L1<<std::setw(5) <<m1<<std::setw(5)<<N1+1<<std::setw(15)<< GlobalC::en.Name_Angular[L1][m1] << std::endl;
-				 }
-			 }
-			 out <<std::endl<<std::endl;
-			 out <<std::setw(5)<< "io"<<std::setw(2)<<"="<<std::setw(2)<<"Orbital index in supercell"<<std::endl;
-			 out <<std::setw(5)<< "spec"<<std::setw(2)<<"="<<std::setw(2)<<"Atomic species label"<<std::endl;
-			 out <<std::setw(5)<< "l"<<std::setw(2)<<"="<<std::setw(2)<<"Angular mumentum quantum number"<<std::endl;
-			 out <<std::setw(5)<< "m"<<std::setw(2)<<"="<<std::setw(2)<<"Magnetic quantum number"<<std::endl;
-			 out <<std::setw(5)<< "z"<<std::setw(2)<<"="<<std::setw(2)<<"Zeta index of orbital"<<std::endl;
-			 out <<std::setw(5)<< "sym"<<std::setw(2)<<"="<<std::setw(2)<<"Symmetry name of real orbital"<<std::endl;
-			 out.close();}
-		}//out_band
+		this->print_orbital_file();
+	}//out_proj_band
 	return;
 }
 
+void energy::print_orbital_file(void)
+{  
+	std::stringstream os;
+	os<<GlobalV::global_out_dir<<"Orbital";
+	std::ofstream out(os.str().c_str());
+	out<< std::setw(5)<<"io"<< std::setw(8) <<"spec" <<std::setw(5)<<"l"<<std::setw(5)<<"m"<<std::setw(5)<<"z"<<std::setw(5)<<"sym"<<std::endl;
+
+
+	for (int i=0; i<GlobalC::ucell.nat; i++)
+	{
+		int   t = GlobalC::ucell.iat2it[i];
+		Atom* atom1 = &GlobalC::ucell.atoms[t];  
+		for(int j=0; j<atom1->nw; ++j)
+		{
+			const int L1 = atom1->iw2l[j];
+			const int N1 = atom1->iw2n[j];
+			const int m1 = atom1->iw2m[j];
+			out <<std::setw(5) << i << std::setw(8) 
+				<< GlobalC::ucell.atoms[t].label <<std::setw(5)
+					<<L1<<std::setw(5) <<m1<<std::setw(5)<<N1+1<<std::setw(15)<< GlobalC::en.Name_Angular[L1][m1] << std::endl;
+		}
+	}
+	out <<std::endl<<std::endl;
+	out <<std::setw(5)<< "io"<<std::setw(2)<<"="<<std::setw(2)<<"Orbital index in supercell"<<std::endl;
+	out <<std::setw(5)<< "spec"<<std::setw(2)<<"="<<std::setw(2)<<"Atomic species label"<<std::endl;
+	out <<std::setw(5)<< "l"<<std::setw(2)<<"="<<std::setw(2)<<"Angular mumentum quantum number"<<std::endl;
+	out <<std::setw(5)<< "m"<<std::setw(2)<<"="<<std::setw(2)<<"Magnetic quantum number"<<std::endl;
+	out <<std::setw(5)<< "z"<<std::setw(2)<<"="<<std::setw(2)<<"Zeta index of orbital"<<std::endl;
+	out <<std::setw(5)<< "sym"<<std::setw(2)<<"="<<std::setw(2)<<"Symmetry name of real orbital"<<std::endl;
+	out.close();
+
+	return;
+}
\ No newline at end of file
diff --git a/source/src_io/write_input.cpp b/source/src_io/write_input.cpp
index d8ef77b508..9f5283fe85 100644
--- a/source/src_io/write_input.cpp
+++ b/source/src_io/write_input.cpp
@@ -97,6 +97,7 @@ void Input::Print(const std::string &fn) const
     ModuleBase::GlobalFunc::OUTP(ofs, "out_wfc_r", out_wfc_r, "output wave functions in realspace");
     ModuleBase::GlobalFunc::OUTP(ofs, "out_dos", out_dos, "output energy and dos");
     ModuleBase::GlobalFunc::OUTP(ofs, "out_band", out_band, "output energy and band structure");
+    ModuleBase::GlobalFunc::OUTP(ofs, "out_proj_band", out_proj_band, "output projected band structure");
     ModuleBase::GlobalFunc::OUTP(ofs, "restart_save", restart_save, "print to disk every step for restart");
     ModuleBase::GlobalFunc::OUTP(ofs, "restart_load", restart_load, "restart from disk");
     ModuleBase::GlobalFunc::OUTP(ofs, "read_file_dir", read_file_dir, "directory of files for reading");
diff --git a/source/src_pw/energy.h b/source/src_pw/energy.h
index bbc0ae7609..c9083aac4f 100644
--- a/source/src_pw/energy.h
+++ b/source/src_pw/energy.h
@@ -58,6 +58,7 @@ class energy
 
 	int out_dos;			// control dos calculation
 	int out_band;                    // control band calculation  pengfei 2014-10-13
+	int out_proj_band;                    // control projected band calculation  jiyy add 2022-05-11
 	const std::string Name_Angular[5][11] = 
 	{
     	{"s"},
@@ -89,6 +90,8 @@ class energy
 
 	void print_format(const std::string &name, const double &value);
 
+	void print_orbital_file();
+
 	void calculate_harris(const int &flag); //mohan add 2012-06-05
 
     double delta_e(void);