extract PDOS by index(atom_index) using plot-tools

Author: jiyuang

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                                        |
 

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)