Merge pull request #104 from 1041176461/develop

modify plot-tools

Author: dyzheng

tools/plot-tools/README.md

@@ -1,7 +1,7 @@
 <!--
  * @Date: 2021-08-21 21:58:06
  * @LastEditors: jiyuyang
- * @LastEditTime: 2021-08-26 14:51:04
+ * @LastEditTime: 2022-01-03 17:21:08
  * @Mail: [email protected], [email protected]
 -->
 
@@ -31,7 +31,7 @@ First, prepare a json file e.g. band-input.json:
 ```
 |    Property    |           Type           |                              Note                              |
 | :------------: | :----------------------: | :------------------------------------------------------------: |
-|   *filename*   |   `str` or `List[str]`   |               Bands data file output from ABACUS               |
+|   *bandfile*   |   `str` or `List[str]`   |               Bands data file output from ABACUS               |
 |    *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'` |
@@ -89,7 +89,6 @@ First, prepare a json file e.g. dos-input.json:
 			]
 		},
 	"pdosfig": "pdos.png",
-	"tdosfig": "tdos.png"
 }
 ```
 If you only want to plot total DOS, you can modify `pdosfile` to `tdosfile` and do not set `species` and `pdosfig`.
@@ -104,7 +103,17 @@ If you only want to plot total DOS, you can modify `pdosfile` to `tdosfile` and
 |   *tdosfig*    |                                    `str`                                    |                                         Output picture of total DOS                                         |
 |   *pdosfig*    |                                    `str`                                    |                                        Output picture of partial DOS                                        |
 
-Then, the following command will plot both total DOS and partial DOS:
+Then, the following command will plot total DOS:
 ```shell
-abacus-plot -d dos-input.json
+abacus-plot -t tdos-input.json
 ```
+
+Then, the following command will plot partial DOS:
+```shell
+abacus-plot -p pdos-input.json
+```
+
+Then, the following command will output parsed partial DOS to directory `PDOS_FILE`:
+```shell
+abacus-plot -o pdos-input.json
+```

tools/plot-tools/abacus_plot/band.py

@@ -1,16 +1,15 @@
 '''
-Date: 2021-05-08 11:47:09
+Date: 2021-12-29 10:27:01
 LastEditors: jiyuyang
-LastEditTime: 2021-08-26 12:07:22
+LastEditTime: 2022-01-03 17:06:14
 Mail: [email protected], [email protected]
 '''
 
 from collections import OrderedDict, namedtuple
+import numpy as np
 from os import PathLike
 from typing import Sequence, Tuple
-
-import matplotlib.pyplot as plt
-import numpy as np
+from matplotlib.figure import Figure
 from matplotlib import axes
 
 from abacus_plot.utils import energy_minus_efermi, list_elem2str, read_kpt
@@ -19,6 +18,16 @@
 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
+
     @classmethod
     def set_vcband(cls, energy: Sequence) -> Tuple[namedtuple, namedtuple]:
         """Separate valence and conduct band
@@ -57,13 +66,12 @@ def read(cls, filename: PathLike) -> Tuple[np.ndarray, np.ndarray]:
         data = np.loadtxt(filename)
         X, y = np.split(data, (1, ), axis=1)
         x = X.flatten()
+
         return x, y
 
-    @classmethod
-    def _set_figure(cls, ax: axes.Axes, index: dict, range: Sequence):
+    def _set_figure(self, index: dict, range: Sequence):
         """set figure and axes for plotting
 
-        :params ax: matplotlib.axes.Axes object
         :params index: dict of label of points of x-axis and its index in data file. Range of x-axis based on index.value()
         :params range: range of y-axis
         """
@@ -79,131 +87,157 @@ def _set_figure(cls, ax: axes.Axes, index: dict, range: Sequence):
                 values.append(t)
 
         # x-axis
-        ax.set_xticks(values)
-        ax.set_xticklabels(keys)
-        ax.set_xlim(values[0], values[-1])
-        ax.set_xlabel("Wave Vector")
+        self.ax.set_xticks(values)
+        self.ax.set_xticklabels(keys)
+        self.ax.set_xlim(values[0], values[-1])
+        if "xlabel_params" in self.plot_params.keys():
+            self.ax.set_xlabel(
+                "Wave Vector", **self.plot_params["xlabel_params"])
+        else:
+            self.ax.set_xlabel("Wave Vector", size=25)
 
         # y-axis
         if range:
-            ax.set_ylim(range[0], range[1])
-        ax.set_ylabel(r"$E-E_{fermi}(eV)$")
+            self.ax.set_ylim(range[0], range[1])
+        if "ylabel_params" in self.plot_params.keys():
+            self.ax.set_ylabel(
+                "Energy(eV)", **self.plot_params["ylabel_params"])
+        else:
+            self.ax.set_ylabel("Energy(eV)", size=25)
+
+        # notes
+        if "notes" in self.plot_params.keys():
+            from matplotlib.offsetbox import AnchoredText
+            if "s" in self.plot_params["notes"].keys() and len(self.plot_params["notes"].keys()) == 1:
+                self.ax.add_artist(AnchoredText(self.plot_params["notes"]["s"], loc='upper left', prop=dict(size=25),
+                                                borderpad=0.2, frameon=False))
+            else:
+                self.ax.add_artist(AnchoredText(**self.plot_params["notes"]))
 
-        # others
-        ax.grid(axis='x', lw=1.2)
-        ax.axhline(0, linestyle="--", c='b', lw=1.0)
-        handles, labels = ax.get_legend_handles_labels()
-        by_label = OrderedDict(zip(labels, handles))
-        ax.legend(by_label.values(), by_label.keys())
+        # ticks
+        if "tick_params" in self.plot_params.keys():
+            self.ax.tick_params(**self.plot_params["tick_params"])
+        else:
+            self.ax.tick_params(labelsize=25)
+
+        # frame
+        bwidth = self._bwidth
+        self.ax.spines['top'].set_linewidth(bwidth)
+        self.ax.spines['right'].set_linewidth(bwidth)
+        self.ax.spines['left'].set_linewidth(bwidth)
+        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 "hline_params" in self.plot_params.keys():
+            self.ax.axhline(0, **self.plot_params["hline_params"])
+        else:
+            self.ax.axhline(0, linestyle="--", c='b', lw=1.0)
+
+        if self._label:
+            handles, labels = self.ax.get_legend_handles_labels()
+            by_label = OrderedDict(zip(labels, handles))
+            if "legend_prop" in self.plot_params.keys():
+                self.ax.legend(by_label.values(), by_label.keys(),
+                               prop=self.plot_params["legend_prop"])
+            else:
+                self.ax.legend(by_label.values(),
+                               by_label.keys(), prop={'size': 15})
 
-    @classmethod
-    def plot(cls, x: Sequence, y: Sequence, index: Sequence, efermi: float = 0, energy_range: Sequence[float] = [], label: str = None, color: str = None, outfile: PathLike = 'band.png'):
+    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
-        :params label: band label. Default: ''
-        :params color: band color. Default: 'black'
-        :params outfile: band picture file name. Default: 'band.png'
         """
 
-        fig, ax = plt.subplots()
-
-        if not color:
-            color = 'black'
+        if not self._color:
+            self._color = 'black'
 
         kpoints, energy = x, y
         energy = energy_minus_efermi(energy, efermi)
 
-        ax.plot(kpoints, energy, lw=0.8, color=color, label=label)
-        cls._set_figure(ax, index, energy_range)
+        self.ax.plot(kpoints, energy, lw=self._lw, color=self._color,
+                     label=self._label, linestyle=self._linestyle)
+        self._set_figure(index, energy_range)
 
-        plt.savefig(outfile)
-
-    @classmethod
-    def singleplot(cls, datafile: PathLike, kptfile: str = [], efermi: float = 0, energy_range: Sequence[float] = [], shift: bool = False, label: str = None, color: str = None, outfile: PathLike = 'band.png'):
+    def singleplot(self, datafile: PathLike, kptfile: str = '', efermi: float = 0, energy_range: Sequence[float] = [], shift: bool = False):
         """Plot band structure using data file
 
         :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 label: band label. Default: ''
-        :params color: band color. Default: 'black'
-        :params outfile: band picture file name. Default: 'band.png'
         """
 
-        fig, ax = plt.subplots()
         kpt = read_kpt(kptfile)
 
-        if not color:
-            color = 'black'
+        if not self._color:
+            self._color = 'black'
 
-        kpoints, energy = cls.read(datafile)
+        kpoints, energy = self.read(datafile)
+        energy = energy_minus_efermi(energy, efermi)
         if shift:
-            vb, cb = cls.set_vcband(energy_minus_efermi(energy, efermi))
-            ax.plot(kpoints, np.vstack((vb.band, cb.band)).T,
-                    lw=0.8, color=color, label=label)
-            cls.info(kpt.full_kpath, vb, cb)
+            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:
-            ax.plot(kpoints, energy_minus_efermi(energy, efermi),
-                    lw=0.8, color=color, label=label)
+            self.ax.plot(kpoints, energy,
+                         lw=self._lw, color=self._color, label=self._label, linestyle=self._linestyle)
         index = kpt.label_special_k
-        cls._set_figure(ax, index, energy_range)
+        self._set_figure(index, energy_range)
 
-        plt.savefig(outfile)
-
-    @classmethod
-    def multiplot(cls, datafile: Sequence[PathLike], kptfile: str = '', efermi: Sequence[float] = [], energy_range: Sequence[float] = [], shift: bool = True, label: Sequence[str] = None, color: Sequence[str] = None, outfile: PathLike = 'band.png'):
+    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
-        :params label: list of band labels, its length equals to `filename`
-        :params color: list of band colors, its length equals to `filename`
-        :params outfile: band picture file name. Default: 'band.png'
         """
 
-        fig, ax = plt.subplots()
         kpt = read_kpt(kptfile)
 
         if not efermi:
             efermi = [0.0 for i in range(len(datafile))]
-        if not label:
-            label = ['' for i in range(len(datafile))]
-        if not color:
-            color = ['black' 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 = cls.read(file)
+            kpoints, energy = self.read(file)
             if shift:
-                vb, cb = cls.set_vcband(energy_minus_efermi(energy, efermi[i]))
+                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
 
-                ax.plot(kpoints, np.vstack((vb.band, cb.band)).T,
-                        lw=0.8, color=color[i], label=label[i])
-                cls.info(kpt.full_kpath, vb, cb)
+                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:
-                ax.plot(kpoints, energy_minus_efermi(energy, efermi[i]),
-                        lw=0.8, color=color[i], label=label[i])
+                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
-        cls._set_figure(ax, index, energy_range)
-
-        plt.savefig(outfile)
+        self._set_figure(index, energy_range)
 
     @classmethod
     def bandgap(cls, vb: namedtuple, cb: namedtuple):
@@ -254,3 +288,25 @@ def band_type(vbm_x, cbm_x):
         for i, j in enumerate(cbm_k):
             if i != 0:
                 print(f"{''.ljust(30)}{' '.join(list_elem2str(j))}", flush=True)
+
+
+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"
+    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")