feat(librpa): allow keyword arguments in reading stdout qpe

mushroom/core/bs.py

@@ -1260,39 +1260,30 @@ def _decode_itrans_string(s):
         s (list of str): either "ivk:ick", "ivck:ivb:icb" or "ivk:ick:ivb:icb"
     """
     itrans = [x.strip() for x in s.split(":")]
+
+    def to_int(s, return_s=True):
+        try:
+            return int(s)
+        except ValueError:
+            if return_s:
+                return s
+            return None
+
     if len(itrans) == 2:
         ivk, ick = list(map(int, itrans))
         ivb = None
         icb = None
     elif len(itrans) == 3:
         ivk = int(itrans[0])
         ick = ivk
-        try:
-            ivb = int(itrans[1])
-        except ValueError:
-            ivb = itrans[1]
-        try:
-            icb = int(itrans[2])
-        except ValueError:
-            icb = itrans[2]
+        ivb = to_int(itrans[1])
+        icb = to_int(itrans[2])
     else:
         ivk, ick = itrans[:2]
-        try:
-            ivk = int(ivk)
-        except ValueError:
-            ivk = None
-        try:
-            ick = int(ick)
-        except ValueError:
-            ick = None
-        try:
-            ivb = int(itrans[2])
-        except ValueError:
-            ivb = itrans[2]
-        try:
-            icb = int(itrans[3])
-        except ValueError:
-            icb = itrans[3]
+        ivk = to_int(ivk, False)
+        ick = to_int(ick, False)
+        ivb = to_int(itrans[2])
+        icb = to_int(itrans[3])
     return ivk, ick, ivb, icb
 
 

mushroom/librpa.py

@@ -45,11 +45,20 @@ def read_self_energy_imagfreq(fn: str):
     return omegas, data
 
 
-def read_quasi_particle_energies_stdout(fn: Union[str, os.PathLike] = "librpa.out", kind: str = "e_qp"):
+def read_quasi_particle_energies_stdout(fn: Union[str, os.PathLike] = "librpa.out",
+                                        kind: str = "qp",
+                                        **kwargs):
     """Read QP energies from standard output
 
     Args:
-        fn (str): """
+        fn (str): path to LibRPA standard output
+        kind (str): any of "qp", "ks", and "exx"
+
+        Other keywords arguments are parsed to the ``BandStructure`` class.
+
+    Returns:
+        BandStructure object, k-points (in fractional coordinates)
+    """
     nspins = None
     nkpts = None
     nstates = None
@@ -105,11 +114,11 @@ def read_quasi_particle_energies_stdout(fn: Union[str, os.PathLike] = "librpa.ou
     vex = np.reshape(vex, (nspins, nkpts, nstates))
     eqp = np.reshape(eqp, (nspins, nkpts, nstates))
     if kind in ["e_qp", "qp"]:
-        bs = BandStructure(eqp, occ)
+        bs = BandStructure(eqp, occ, **kwargs)
     elif kind in ["e_ks", "e_mf", "ks", "mf"]:
-        bs = BandStructure(eks, occ)
+        bs = BandStructure(eks, occ, **kwargs)
     elif kind in ["e_hf", "hf", "exx"]:
-        bs = BandStructure(eks - vxc + vex, occ)
+        bs = BandStructure(eks - vxc + vex, occ, **kwargs)
     else:
         raise ValueError("kind %s not supported, use any of following: ks, qp, exx" % kind)
     return bs, kpoints

mushroom/test/test_librpa.py

@@ -62,11 +62,15 @@ def test_read_quasi_particle_energies(self):
    18          0.00000         -2.54792        -12.66794         -7.80805         -2.53696         -0.04424         -0.22524
 
 """
-        bs, kpts = librpa.read_quasi_particle_energies_stdout(StringIO(stdout))
+        bs, kpts = librpa.read_quasi_particle_energies_stdout(StringIO(stdout), kind="ks")
         self.assertEqual(bs.nspins, 1)
         self.assertEqual(bs.nkpts, 1)
         self.assertEqual(bs.nbands, 18)
-        self.assertEqual(bs.fund_gap(), 3.7043)
+        self.assertAlmostEqual(bs.fund_gap(), 2.97139)
+        bs, kpts = librpa.read_quasi_particle_energies_stdout(StringIO(stdout), kind="exx", use_occ_only=True)
+        self.assertAlmostEqual(bs.fund_gap(), 9.01068)
+        bs, kpts = librpa.read_quasi_particle_energies_stdout(StringIO(stdout), kind="exx")
+        self.assertAlmostEqual(bs.fund_gap(), 8.75884)
 
 if __name__ == "__main__":
     ut.main()