diff --git a/castep_outputs/parsers/bands_file_parser.py b/castep_outputs/parsers/bands_file_parser.py index 3d76338..5395b3e 100644 --- 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 index 7aa828f..976004d 100644 --- 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": {} } \ No newline at end of file diff --git a/castep_outputs/test/bands.yaml b/castep_outputs/test/bands.yaml index 1a17020..dc45ce7 100644 --- 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 index 2330bb9..2647e9c 100644 --- 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