Add pyatb parse mode.

Author: Lonya0

README.md

@@ -16,12 +16,13 @@ However, the user can always manage the dependency themselves, here are the pack
 Current:
 
 | Package  | Structure | Eigenvalues | Hamiltonian | Density matrix | Overlap matrix |
-|  :----:  |  :----:   |   :----:    |    :----:   |     :----:     |     :----:     |
-| ABACUS   | √         | √           | √           | √              | √              |
-| RESCU    | √         |             | √           |                | √              |
-| SIESTA   | √         |             | √           | √              | √              |
-| Gaussian | √         |             | √           | √              | √              |
-| VASP     | √         | √           |             |                |                |
+|:--------:|:---------:|:-----------:|    :----:   |     :----:     |     :----:     |
+|  ABACUS  |     √     |      √      | √           | √              | √              |
+|  RESCU   |     √     |             | √           |                | √              |
+|  SIESTA  |     √     |             | √           | √              | √              |
+| Gaussian |     √     |             | √           | √              | √              |
+|   VASP   |     √     |      √      |             |                |                |
+|  PYATB   |     √     |      √      |             |                |                |
 
 Ongoing:
 
@@ -35,15 +36,15 @@ Ongoing:
 To parse the DFT output files into readable data format, user can follows:
 
 ```bash
-dftio parse [-h] [-ll {DEBUG,3,INFO,2,WARNING,1,ERROR,0}] [-lp LOG_PATH] [-m MODE] [-n NUM_WORKERS] [-r ROOT] [-p PREFIX] [-o OUTROOT] [-f FORMAT] [-ham] [-ovp] [-dm] [-eig]
+usage: dftio parse [-h] [-ll {DEBUG,3,INFO,2,WARNING,1,ERROR,0}] [-lp LOG_PATH] [-m MODE] [-n NUM_WORKERS] [-r ROOT] [-p PREFIX] [-o OUTROOT] [-f FORMAT] [-ham] [-ovp] [-dm] [-eig] [-min BAND_INDEX_MIN]
 
-optional arguments:
+options:
   -h, --help            show this help message and exit
   -ll {DEBUG,3,INFO,2,WARNING,1,ERROR,0}, --log-level {DEBUG,3,INFO,2,WARNING,1,ERROR,0}
                         set verbosity level by string or number, 0=ERROR, 1=WARNING, 2=INFO and 3=DEBUG (default: INFO)
   -lp LOG_PATH, --log-path LOG_PATH
                         set log file to log messages to disk, if not specified, the logs will only be output to console (default: None)
-  -m MODE, --mode MODE  The name of the DFT software. (default: abacus)
+  -m MODE, --mode MODE  The name of the DFT software, currently support abacus/rescu/siesta/gaussian/pyatb (default: abacus)
   -n NUM_WORKERS, --num_workers NUM_WORKERS
                         The number of workers used to parse the dataset. (For n>1, we use the multiprocessing to accelerate io.) (default: 1)
   -r ROOT, --root ROOT  The root directory of the DFT files. (default: ./)
@@ -52,12 +53,14 @@ optional arguments:
   -o OUTROOT, --outroot OUTROOT
                         The output root directory. (default: ./)
   -f FORMAT, --format FORMAT
-                        The output root directory. (default: dat)
+                        The output file format, should be dat, ase or lmdb. (default: dat)
   -ham, --hamiltonian   Whether to parse the Hamiltonian matrix. (default: False)
   -ovp, --overlap       Whether to parse the Overlap matrix (default: False)
   -dm, --density_matrix
                         Whether to parse the Density matrix (default: False)
   -eig, --eigenvalue    Whether to parse the kpoints and eigenvalues (default: False)
+  -min BAND_INDEX_MIN, --band_index_min BAND_INDEX_MIN
+                        The initial band index for eigenvalues to save.(0-band_index_min) bands will be ignored! (default: 0)
 ```
 
 ## Call for Contributors

dftio/__main__.py

@@ -77,7 +77,7 @@ def main_parser() -> argparse.ArgumentParser:
         "--mode",
         type=str,
         default="abacus",
-        help="The name of the DFT software, currently support abacus/rescu/siesta/gaussian",
+        help="The name of the DFT software, currently support abacus/rescu/siesta/gaussian/pyatb",
     )
 
     parser_parse.add_argument(
@@ -244,9 +244,6 @@ def main():
             with Pool(args.num_workers) as p:
                 list(tqdm(p.imap(wapper(dict_args), range(len(parser))), total=len(parser), desc="Parsing the DFT files: "))
         else:
-            parser = ParserRegister(
-                **dict_args
-            )
             for i in tqdm(range(len(parser)), desc="Parsing the DFT files: "):
                 parser.write(idx=i, **dict_args)
 

dftio/io/__init__.py

@@ -3,12 +3,14 @@
 from dftio.io.gaussian.gaussian_parser import GaussianParser
 from dftio.io.siesta.siesta_parser import SiestaParser
 from dftio.io.vasp.vasp_parser import VASPParser
+from dftio.io.pyatb.pyatb_parser import PyatbParser
 
 
 __all__ = [
     "AbacusParser",
     "RescuParser",
     "GaussianParser",
     "SiestaParser",
-    "VASPParser"
+    "VASPParser",
+    "PyatbParser"
 ]

dftio/io/pyatb/pyatb_parser.py

@@ -0,0 +1,76 @@
+import ase
+import dpdata
+import os
+import numpy as np
+from dftio.io.parse import Parser, ParserRegister
+from dftio.register import Register
+from dftio.data import _keys
+
+@ParserRegister.register("pyatb")
+class PyatbParser(Parser):
+    """
+    PYATB parser is commonly used for normal input file generated by 'pyatb_input ... --band' and running 'pyatb'
+    locally in generated 'pyatb' directory.
+
+    Which should be like:
+    pyatb/
+    |-- Input
+    |-- Out
+    |   |-- Band_Structure
+    |   |   |-- band.dat *
+    |   |   |-- band.pdf
+    |   |   |-- band_info.dat
+    |   |   |-- high_symmetry_kpoint.dat
+    |   |   |-- kpt.dat *
+    |   |   |-- plot_band.py
+    |   |   `-- x_coor_array.dat
+    |   |-- input.json
+    |   `-- running.log
+    |-- STRU *
+    `-- get_Energy.out
+
+    * means essential, but normally you will get other output files as well.
+    """
+
+    def __init__(
+            self,
+            root,
+            prefix,
+            **kwargs
+            ):
+        super(PyatbParser, self).__init__(root, prefix)
+        self.raw_sys = [
+            dpdata.System(os.path.join(self.raw_datas[idx], 'pyatb', "STRU"), fmt="abacus/stru") for
+            idx in range(len(self.raw_datas))]
+
+    # essential
+    def get_structure(self, idx):
+        sys = self.raw_sys[idx]
+        
+        structure = {
+            _keys.ATOMIC_NUMBERS_KEY: np.array([ase.atom.atomic_numbers[i] for i in sys.data["atom_names"]], dtype=np.int32)[sys.data["atom_types"]],
+            _keys.PBC_KEY: np.array([True, True, True]) # abacus does not allow non-pbc structure
+        }
+        structure[_keys.POSITIONS_KEY] = sys.data["coords"].astype(np.float32)
+        structure[_keys.CELL_KEY] = sys.data["cells"].astype(np.float32)
+
+        return structure
+    
+    # essential
+    def get_eigenvalue(self, idx, band_index_min=0):
+        path = self.raw_datas[idx]
+
+        assert os.path.exists(os.path.join(path, 'pyatb', "Out", "Band_Structure", "band.dat"))
+        eigs = np.loadtxt(os.path.join(path, 'pyatb', "Out", "Band_Structure", "band.dat"))[np.newaxis, :, band_index_min:]
+        assert os.path.exists(os.path.join(path, 'pyatb', "Out", "Band_Structure", "kpt.dat"))
+        kpts = np.loadtxt(os.path.join(path, 'pyatb', "Out", "Band_Structure", "kpt.dat"))
+
+        return {_keys.ENERGY_EIGENVALUE_KEY: eigs.astype(np.float32), _keys.KPOINT_KEY: kpts.astype(np.float32)}
+    
+    # essential
+    def get_basis(self, idx):
+        raise NotImplementedError("PYATB does not support block parsing yet.")
+    
+    # essential
+    def get_blocks(self, idx, hamiltonian=True, overlap=False, density_matrix=False):
+        raise NotImplementedError("PYATB does not support block parsing yet.")

example/pyatb/README.md

@@ -0,0 +1,68 @@
+# How to use PYATB mode
+PYATB parser is commonly used for normal input file generated by `pyatb_input ... --band` without changing default config and running `pyatb`
+    locally in generated 'pyatb' directory.
+
+Which should be like:
+
+    pyatb/
+    |-- Input
+    |-- Out
+    |   |-- Band_Structure
+    |   |   |-- band.dat  *
+    |   |   |-- band.pdf
+    |   |   |-- band_info.dat
+    |   |   |-- high_symmetry_kpoint.dat
+    |   |   |-- kpt.dat  *
+    |   |   |-- plot_band.py
+    |   |   `-- x_coor_array.dat
+    |   |-- input.json
+    |   `-- running.log
+    |-- STRU  *
+    `-- get_Energy.out
+
+\* means essential, but normally you will get other output files as well.
+
+Then the root (or root + prefix) should be the directory that contain pyatb directory.
+
+# Example:
+
+I have a directory:
+
+    abacus_pyatb_test/
+    |-- task.001
+    |   |-- INPUT
+    |   |-- STRU
+    |   |-- OUT.ABACUS
+    |   |   |-- INPUT
+    |   |   |-- STRU.cif
+    |   |   |-- data-HR-sparse_SPIN0.csr
+    |   |   |-- data-SR-sparse_SPIN0.csr
+    |   |   |-- data-rR-sparse.csr
+    |   |   |-- kpoints
+    |   |   `-- running_scf.log
+    |   |-- Out
+    |   |   `-- running.log
+    |   `-- pyatb
+    |       |-- Input
+    |       |-- STRU  *
+    |       |-- get_Energy.out
+    |       `-- Out
+    |           |-- input.json
+    |           |-- running.log
+    |           `-- Band_Structure
+    |              |-- band.dat  *
+    |              |-- band.pdf
+    |              |-- band_info.dat
+    |              |-- high_symmetry_kpoint.dat
+    |              |-- kpt.dat  *
+    |              |-- plot_band.py
+    |              `-- x_coor_array.dat
+    |-- ...
+
+In `/task.001` run `pyatb_input -i . --band`, which generate `pyatb/Input`, `pyatb/STRU` and `pyatb/get_Energy.out`.
+
+Then in `/pyatb` run `pyatb`, which generate `Band_Structure/band.dat` and `Band_Structure/kpt.dat`.
+
+Without changing output directory name, you will get same directory structure!
+
+Files with * are to be read, if any is missing, you could add by yourself.