Simplify bands parser

oerc0122 · Feb 7, 2025 · 7336d29 · 7336d29
1 parent ee396eb
commit 7336d29
Show file tree

Hide file tree

Showing 4 changed files with 80 additions and 69 deletions.
diff --git a/castep_outputs/parsers/bands_file_parser.py b/castep_outputs/parsers/bands_file_parser.py
@@ -2,13 +2,12 @@
 from __future__ import annotations
 
 import re
-from collections import defaultdict
-from typing import TextIO, TypedDict
+from typing import Literal, TextIO, TypedDict
 
 from ..utilities import castep_res as REs
 from ..utilities.datatypes import ThreeVector
 from ..utilities.filewrapper import Block
-from ..utilities.utility import fix_data_types
+from ..utilities.utility import to_type
 from .parse_utilities import parse_regular_header
 
 
@@ -29,11 +28,15 @@ class BandsQData(TypedDict, total=False):
     weight: float
 
 
-class BandsFileInfo(TypedDict, total=False):
-    """Bands eigenvalue info of a bands calculation."""
-
-    #: Bands info in file.
-    bands: list[BandsQData]
+BandsFileInfo = TypedDict('BandsFileInfo', {
+    "eigenvalues": int,
+    "electrons": int,
+    "k-points": int,
+    "spin components": int,
+    "Fermi Energy": float,
+    "coords": dict,
+    "bands": list[BandsQData],
+    })
 
 
 def parse_bands_file(bands_file: TextIO) -> BandsFileInfo:
@@ -50,8 +53,10 @@ def parse_bands_file(bands_file: TextIO) -> BandsFileInfo:
     BandsFileInfo
         Parsed info.
     """
-    bands_info: BandsFileInfo = defaultdict(list)
-    qdata = {}
+    bands_info: BandsFileInfo = {"bands": []}
+    qdata: BandsQData = {}
+    accum: list[str] = []
+    current: Literal["band", "band_up", "band_dn"] = "band"
 
     block = Block.from_re("", bands_file, "", REs.THREEVEC_RE, n_end=3)
     data = parse_regular_header(block, ("Fermi energy",))
@@ -60,38 +65,28 @@ def parse_bands_file(bands_file: TextIO) -> BandsFileInfo:
     for line in bands_file:
         if line.startswith("K-point"):
             if qdata:
-                fix_data_types(qdata, {"qpt": float,
-                                       "weight": float,
-                                       "spin_comp": int,
-                                       "band": float,
-                                       "band_up": float,
-                                       "band_dn": float,
-                                       })
+                qdata[current] = to_type(accum, float)
+                qdata["spin_comp"] = "band_dn" in qdata
                 bands_info["bands"].append(qdata)
+                accum = []
+                current = "band"
             _, _, *qpt, weight = line.split()
-            qdata = {"qpt": qpt, "weight": weight, "spin_comp": None, "band": []}
+            qdata = {"qpt": to_type(qpt, float), "weight": float(weight)}
 
         elif line.startswith("Spin component"):
-            qdata["spin_comp"] = line.split()[2]
-            if qdata["spin_comp"] != "1":
-                qdata["band_up"] = qdata.pop("band")
-                if "band_dn" not in qdata:
-                    qdata["band_dn"] = []
+            spin_comp = int(line.split()[2])
+            if spin_comp != 1:
+                qdata["band_up"] = to_type(accum, float)
+                accum = []
+                current = "band_dn"
 
         elif re.match(rf"^\s*{REs.FNUMBER_RE}$", line.strip()):
-            if qdata["spin_comp"] != "1":
-                qdata["band_dn"].append(line)
-            else:
-                qdata["band"].append(line)
+            accum.append(line.strip())
+
 
     if qdata:
-        fix_data_types(qdata, {"qpt": float,
-                               "weight": float,
-                               "spin_comp": int,
-                               "band": float,
-                               "band_up": float,
-                               "band_dn": float,
-                               })
+        qdata[current] = to_type(accum, float)
+        qdata["spin_comp"] = "band_dn" in qdata
         bands_info["bands"].append(qdata)
 
     return bands_info
diff --git a/castep_outputs/test/bands.json b/castep_outputs/test/bands.json
@@ -1,29 +1,4 @@
 {
-  "k-points": 2,
-  "spin components": 1,
-  "electrons": 300,
-  "eigenvalues": 3,
-  "Fermi energy": [
-    0.323077
-  ],
-  "unit_cell": [
-    [
-      11.932676,
-      0.0,
-      0.0
-    ],
-    [
-      0.0,
-      11.932676,
-      0.0
-    ],
-    [
-      0.0,
-      0.0,
-      11.932676
-    ]
-  ],
-  "coords": {},
   "bands": [
     {
       "qpt": [
@@ -32,12 +7,15 @@
         -0.375
       ],
       "weight": 0.03125,
-      "spin_comp": 1,
       "band": [
+        -2.56339233,
+        -2.56319008,
+        -2.5626493,
         -2.56339233,
         -2.56319008,
         -2.5626493
-      ]
+      ],
+      "spin_comp": false
     },
     {
       "qpt": [
@@ -46,12 +24,42 @@
         -0.125
       ],
       "weight": 0.03125,
-      "spin_comp": 1,
-      "band": [
+      "band_up": [
         -2.56369745,
         -2.56331049,
         -2.56254
-      ]
+      ],
+      "band_dn": [
+        2.56369745,
+        2.56331049,
+        2.56254
+      ],
+      "spin_comp": true
     }
-  ]
+  ],
+  "k-points": 2,
+  "spin components": 1,
+  "electrons": 300,
+  "eigenvalues": 3,
+  "Fermi energy": [
+    0.323077
+  ],
+  "unit_cell": [
+    [
+      11.932676,
+      0.0,
+      0.0
+    ],
+    [
+      0.0,
+      11.932676,
+      0.0
+    ],
+    [
+      0.0,
+      0.0,
+      11.932676
+    ]
+  ],
+  "coords": {}
 }
diff --git a/castep_outputs/test/bands.yaml b/castep_outputs/test/bands.yaml
@@ -1,12 +1,13 @@
 Fermi energy: [0.323077]
 bands:
-- band: [-2.56339233, -2.56319008, -2.5626493]
+- band: [-2.56339233, -2.56319008, -2.5626493, -2.56339233, -2.56319008, -2.5626493]
   qpt: [-0.375, -0.375, -0.375]
-  spin_comp: 1
+  spin_comp: false
   weight: 0.03125
-- band: [-2.56369745, -2.56331049, -2.56254]
+- band_dn: [2.56369745, 2.56331049, 2.56254]
+  band_up: [-2.56369745, -2.56331049, -2.56254]
   qpt: [-0.375, -0.375, -0.125]
-  spin_comp: 1
+  spin_comp: true
   weight: 0.03125
 coords: {}
 eigenvalues: 3

diff --git a/castep_outputs/test/test.bands b/castep_outputs/test/test.bands
@@ -12,8 +12,15 @@ Spin component 1
    -2.56339233
    -2.56319008
    -2.56264930
+   -2.56339233
+   -2.56319008
+   -2.56264930
 K-point     2 -0.37500000 -0.37500000 -0.12500000  0.03125000
 Spin component 1
    -2.56369745
    -2.56331049
    -2.56254000
+Spin component 2
+   2.56369745
+   2.56331049
+   2.56254000