Add test for tb.dat

Author: qiaojunfeng

examples/silicon/.gitignore

@@ -1 +1,2 @@
 build
+build_tb

examples/silicon/read_tb.py

@@ -1,8 +1,5 @@
 #!/usr/bin/env python
-
-# (c) 2015-2018, ETH Zurich, Institut fuer Theoretische Physik
-# Author: Dominik Gresch <[email protected]>
-
+# Construct a Model from wannier90 tb.dat file.
 import os
 import shutil
 import subprocess
@@ -12,29 +9,30 @@
 
 if __name__ == "__main__":
     WANNIER90_COMMAND = os.path.expanduser("~/git/wannier90/wannier90.x")
-    BUILD_DIR = "./build"
-
-    # shutil.rmtree(BUILD_DIR, ignore_errors=True)
-    # shutil.copytree("./input", BUILD_DIR)
-    # subprocess.call([WANNIER90_COMMAND, "silicon"], cwd=BUILD_DIR)
+    BUILD_DIR = "./build_tb"
 
-    # model = tb.Model.from_wannier_folder(BUILD_DIR, prefix="silicon")
-    # print(model)
-    # print(model.hop[(2, 1, -3)])
+    if not os.path.exists(BUILD_DIR):
+        shutil.copytree("./input", BUILD_DIR)
+        subprocess.call([WANNIER90_COMMAND, "silicon"], cwd=BUILD_DIR)
 
-    BUILD_DIR = "./build2"
-    model = tb.Model.from_wannier_tb_file(
-        tb_file=f'{BUILD_DIR}/silicon_tb.dat',
-        wsvec_file=f'{BUILD_DIR}/silicon_wsvec.dat'
+    model = tb.Model.from_wannier_tb_files(
+        tb_file=f"{BUILD_DIR}/silicon_tb.dat",
+        wsvec_file=f"{BUILD_DIR}/silicon_wsvec.dat",
     )
     print(model)
-    # print(model.hop[(2, 1, -3)])
-    # exit()
 
+    # Compute band structure along an arbitrary kpath
     theta = 37 / 180 * np.pi
     phi = 43 / 180 * np.pi
     rlist = np.linspace(0, 2, 20)
-    klist = [[r*np.sin(theta)*np.cos(phi), r*np.sin(theta)*np.sin(phi), r*np.cos(theta)] for r in rlist]
+    klist = [
+        [
+            r * np.sin(theta) * np.cos(phi),
+            r * np.sin(theta) * np.sin(phi),
+            r * np.cos(theta),
+        ]
+        for r in rlist
+    ]
 
     eigvals = model.eigenval(klist)
 

tbmodels/_tb_model.py

@@ -538,7 +538,7 @@ def _read_tb(iterator, ignore_orbital_order=False):
 
         lattice = np.zeros((3, 3))
         for i in range(3):
-            lattice[i, :] = np.fromstring(next(iterator), sep=' ')
+            lattice[i, :] = np.fromstring(next(iterator), sep=" ")
 
         num_wann = int(next(iterator))
 
@@ -554,28 +554,32 @@ def _read_tb(iterator, ignore_orbital_order=False):
         # <0n|H|Rm>
         hop_list = []
         for ir in range(nrpts):
-            next(iterator) # skip empty
+            next(iterator)  # skip empty
             r_vec = [int(_) for _ in next(iterator).strip().split()]
             for j in range(num_wann):
                 for i in range(num_wann):
                     line = next(iterator).strip().split()
                     iw, jw = [int(_) for _ in line[:2]]
                     if not ignore_orbital_order and (iw != i + 1 or jw != j + 1):
-                        raise ValueError(f"Inconsistent orbital numbers in line '{line}'")
+                        raise ValueError(
+                            f"Inconsistent orbital numbers in line '{line}'"
+                        )
                     ham = (float(line[2]) + 1j * float(line[3])) / deg_pts[ir]
                     hop_list.append([ham, i, j, r_vec])
 
         # <0n|r|Rm>
         r_list = []
         for ir in range(nrpts):
-            next(iterator) # skip empty
+            next(iterator)  # skip empty
             r_vec = [int(_) for _ in next(iterator).strip().split()]
             for j in range(num_wann):
                 for i in range(num_wann):
                     line = next(iterator).strip().split()
                     iw, jw = [int(_) for _ in line[:2]]
                     if not ignore_orbital_order and (iw != i + 1 or jw != j + 1):
-                        raise ValueError(f"Inconsistent orbital numbers in line '{line}'")
+                        raise ValueError(
+                            f"Inconsistent orbital numbers in line '{line}'"
+                        )
                     r = np.array([float(_) for _ in line[2:]])
                     r = r[::2] + 1j * r[1::2]
                     r_list.append([r, i, j, r_vec])
@@ -769,7 +773,7 @@ def remap_hoppings(hop_entries):
             return cls.from_hop_list(size=num_wann, hop_list=hop_entries, **kwargs)
 
     @classmethod  # noqa: MC0001
-    def from_wannier_tb_file(  # pylint: disable=too-many-locals
+    def from_wannier_tb_files(  # pylint: disable=too-many-locals
         cls,
         *,
         tb_file: str,
@@ -803,19 +807,19 @@ def from_wannier_tb_file(  # pylint: disable=too-many-locals
 
         with open(tb_file) as f:
             lattice, num_wann, nrpts, deg_pts, hop_list, r_list = cls._read_tb(f)
-        
+
         kwargs["uc"] = lattice
 
-        def get_centers(r_list: list) -> list:
-            centers = [None for _ in range(num_wann)]
+        def get_centers(r_list: ty.List[ty.Any]) -> ty.List[npt.NDArray[np.float_]]:
+            centers = [np.zeros(3) for _ in range(num_wann)]
             for r, i, j, r_vec in r_list:
                 if r_vec != [0, 0, 0]:
                     continue
                 if i != j:
                     continue
                 r = np.array(r)
                 if not np.allclose(np.abs(r.imag), 0):
-                    raise ValueError(f'Center should be real: WF {i+1}, center = {r}')
+                    raise ValueError(f"Center should be real: WF {i+1}, center = {r}")
                 centers[i] = r.real
             return centers
 
@@ -829,21 +833,15 @@ def get_centers(r_list: list) -> list:
         hop_entries = hop_list
 
         with open(wsvec_file) as f:
-            wsvec_generator = cls._async_parse(
-                cls._read_wsvec(f), chunksize=num_wann
-            )
+            wsvec_generator = cls._async_parse(cls._read_wsvec(f), chunksize=num_wann)
 
             def remap_hoppings(hop_entries):
                 for t, orbital_1, orbital_2, R in hop_entries:
                     # Step _async_parse to where it accepts
                     # a new key.
                     # The _async_parse does not raise StopIteration
-                    next(  # pylint: disable=stop-iteration-return
-                        wsvec_generator
-                    )
-                    T_list = wsvec_generator.send(
-                        (orbital_1, orbital_2, tuple(R))
-                    )
+                    next(wsvec_generator)  # pylint: disable=stop-iteration-return
+                    T_list = wsvec_generator.send((orbital_1, orbital_2, tuple(R)))
                     N = len(T_list)
                     for T in T_list:
                         # not using numpy here increases performance
@@ -856,9 +854,7 @@ def remap_hoppings(hop_entries):
 
             hop_entries = remap_hoppings(hop_entries)
 
-            return cls.from_hop_list(
-                size=num_wann, hop_list=hop_entries, **kwargs
-            )
+            return cls.from_hop_list(size=num_wann, hop_list=hop_entries, **kwargs)
 
     @staticmethod
     def _async_parse(iterator, chunksize=1):