feat: read aims self energy restart file

mushroom/aims/gw.py

@@ -2,14 +2,17 @@
 """Utilites related to GW in FHI-aims"""
 import os
 import re
-import numpy as np
+import struct
 from typing import Callable, Union
 from concurrent.futures import ProcessPoolExecutor, as_completed
 
+import numpy as np
+
 from mushroom.core.logger import loggers
 
 __all__ = [
-    "read_aims_self_energy_dir"
+    "read_aims_self_energy_dir",
+    "read_aims_self_energy_restart_file"
 ]
 
 _logger = loggers["aims"]
@@ -300,3 +303,23 @@ def _read_aims_single_specfunc_dat(specfuncdat_fn, unit: str = "ev"):
         specfunc /= conv
 
     return eKS, eQP_wo_c, omegas, real_part, imag_part, specfunc
+
+
+def read_aims_self_energy_restart_file(fn: Union[str, os.PathLike] = "self_energy_grid.dat"):
+    """Read restart file containing self-energy on imaginary frequency axis"""
+    with open(fn, "rb") as h:
+        data = h.read()
+
+    header = struct.unpack("i" * 4, data[:16])
+    nspin, nkpts, nstates, nfreq = header
+
+    begin, end = 16, 16 + 8 * nfreq
+    omega_imag = struct.unpack('d' * nfreq, data[begin:end])
+    begin, end = end, 16 + 8 * nfreq + 16 * nfreq * (2 + nspin * nkpts * nstates)
+    data = struct.unpack('d' * nspin * nkpts * nstates * nfreq * 2,
+                         data[begin:end])
+    # convert to numpy array
+    omega_imag = np.array(omega_imag)
+    data = np.array(data[0::2]) + np.array(data[1::2]) * 1.0j
+    data = np.reshape(data, (nspin, nkpts, nstates, nfreq))
+    return omega_imag, data

mushroom/aims/test/data/test_self_energy_restart_file.json

@@ -0,0 +1,9 @@
+{
+  "self_energy_grid_1.dat": {
+    "_comment": "H2O",
+    "nspins": 1,
+    "nkpts": 1,
+    "nstates": 9,
+    "nfreqs": 32
+  }
+}

mushroom/aims/test/test_gw.py

@@ -9,7 +9,7 @@
 import numpy as np
 
 from mushroom.aims.gw import read_aims_self_energy_dir, _read_aims_single_sigc_dat, \
-    get_nsbk_filename, get_nsbk_filename_pattern
+    get_nsbk_filename, get_nsbk_filename_pattern, read_aims_self_energy_restart_file
 
 
 class test_read_self_energy_directory(ut.TestCase):
@@ -115,5 +115,21 @@ def test_real_cases(self):
             self.assertEqual(np.shape(sigc_bands_merged)[2], sum(verify["nkpts_band"]))
 
 
+class test_read_self_energy_restart_file(ut.TestCase):
+
+    def test_real_cases(self):
+        datadir = pathlib.Path(__file__).parent / "data"
+        testcases_json = datadir / "test_self_energy_restart_file.json"
+        with testcases_json.open('r') as h:
+            verifies = json.load(h)
+
+        for case, verify in verifies.items():
+            sefile = datadir / case
+            omega, sigc = read_aims_self_energy_restart_file(sefile)
+            print("Testing self_energy_restart_file: {} - {}".format(case, verify["_comment"]))
+            self.assertEqual(len(omega), verify["nfreqs"])
+            self.assertTupleEqual(sigc.shape, tuple(verify[x] for x in ["nspins", "nkpts", "nstates", "nfreqs"]))
+
+
 if __name__ == '__main__':
     ut.main()