From da13fd89d9b68b090be3c84903ace9321ed416ba Mon Sep 17 00:00:00 2001
From: Jacob Wilkins <jacob.wilkins@stfc.ac.uk>
Date: Sat, 24 Feb 2024 14:00:39 +0000
Subject: [PATCH] Add deltascf parsing

---
 castep_outputs/parsers/castep_file_parser.py | 46 +++++++++++++++++++-
 castep_outputs/test/test_castep_parser.py    | 19 ++++++--
 castep_outputs/utilities/datatypes.py        | 22 ++++++++++
 3 files changed, 82 insertions(+), 5 deletions(-)

diff --git a/castep_outputs/parsers/castep_file_parser.py b/castep_outputs/parsers/castep_file_parser.py
index 8b3507c..14e98da 100644
--- a/castep_outputs/parsers/castep_file_parser.py
+++ b/castep_outputs/parsers/castep_file_parser.py
@@ -26,6 +26,7 @@
     CellInfo,
     CharTable,
     ConstraintsReport,
+    DeltaSCFReport,
     DipoleTable,
     ElasticProperties,
     FinalConfig,
@@ -87,6 +88,7 @@ class Filters(Flag):
     BS = auto()
     CELL = auto()
     CHEM_SHIELDING = auto()
+    DELTA_SCF = auto()
     DIPOLE = auto()
     ELASTIC = auto()
     ELF = auto()
@@ -119,7 +121,7 @@ class Filters(Flag):
            ELASTIC | ELF | FINAL_CONFIG | FORCE |
            MD_SUMMARY | OPTICS | POPN_ANALYSIS |
            POSITION | SOLVATION | SPECIES_PROPS | SPIN |
-           STRESS | TDDFT | THERMODYNAMICS | TSS)
+           STRESS | TDDFT | THERMODYNAMICS | TSS | DELTA_SCF)
 
     MEDIUM = LOW | PARAMETERS | GEOM_OPT | MD | PHONON
 
@@ -1245,6 +1247,18 @@ def parse_castep_file(castep_file_in: TextIO,
 
             curr_run["elastic"].update(_process_elastic_properties(block))
 
+        # DeltaSCF
+
+        elif block := Block.from_re(line, castep_file,
+                                    "Calculating MODOS weights", r"^\s*$"):
+
+            if Filters.DELTA_SCF not in to_parse:
+                continue
+
+            logger("Found delta SCF data")
+
+            curr_run["delta_scf"] = _process_delta_scf(block)
+
         # --- Extra blocks for testing
 
         # Hugoniot data
@@ -2394,3 +2408,33 @@ def _process_elastic_properties(block: Block) -> ElasticProperties:
                                            )
 
     return accum
+
+
+def _process_delta_scf(block: Block) -> DeltaSCFReport:
+    """Process MODOS delta SCF block."""
+    accum: DeltaSCFReport = {"states": []}
+
+    for line in block:
+        if line.startswith("Taking band from"):
+            accum["file"] = line.split()[-1]
+        elif "MODOS state" in line:
+            accum["states"].append({})
+        elif "nr." in line:
+            accum["states"][-1]["band"] = int(get_numbers(line)[0])
+        elif "spin" in line:
+            accum["states"][-1]["spin"] = int(get_numbers(line)[0])
+        elif "Population of state" in line:
+            numbers = get_numbers(line)
+            band, spin, pop = int(numbers[0]), int(numbers[1]), float(numbers[2])
+
+            for state in accum["states"]:
+                if state["band"] == band and state["spin"] == spin:
+                    state["pop"] = pop
+
+        elif "Writing file" in line:
+            band, spin = map(int, get_numbers(line)[-2:])
+            for state in accum["states"]:
+                if state["band"] == band and state["spin"] == spin:
+                    state["file"] = line.split()[-1]
+
+    return accum
diff --git a/castep_outputs/test/test_castep_parser.py b/castep_outputs/test/test_castep_parser.py
index 881132b..45a6a41 100644
--- a/castep_outputs/test/test_castep_parser.py
+++ b/castep_outputs/test/test_castep_parser.py
@@ -901,7 +901,8 @@ def test_get_target_stress(self):
 
     def test_delta_scf(self):
         test_text = io.StringIO("""
-+-------------------INPUT PARAMETERS-------------------+
+               Calculating MODOS weights
+ +-------------------INPUT PARAMETERS-------------------+
 Taking band from model   N2-base.check
   MODOS state               1
   MODOS band  nr.           5
@@ -911,12 +912,22 @@ def test_delta_scf(self):
   MODOS band has spin       1
 |DeltaSCF| Population of state:     5     1   1.000000
 |DeltaSCF| Population of state:     6     1   0.000000
+ Writing file   N2-modos.modos_state_5_1
+ Writing file   N2-modos.modos_state_6_1
 
     """)
-        self.skipTest("Not implemented yet")
+
         test_dict = parse_castep_file(test_text)[0]
-        pprint.pprint(test_dict)
-        self.assertEqual(test_dict, {})
+
+        self.assertEqual(test_dict, {'delta_scf': {'file': 'N2-base.check',
+                                                   'states': [{'band': 5,
+                                                               'file': 'N2-modos.modos_state_5_1',
+                                                               'pop': 1.0,
+                                                               'spin': 1},
+                                                              {'band': 6,
+                                                               'file': 'N2-modos.modos_state_6_1',
+                                                               'pop': 0.0,
+                                                               'spin': 1}]}})
 
     def test_get_scf(self):
         test_text = io.StringIO("""
diff --git a/castep_outputs/utilities/datatypes.py b/castep_outputs/utilities/datatypes.py
index 7f2b2af..99344f7 100644
--- a/castep_outputs/utilities/datatypes.py
+++ b/castep_outputs/utilities/datatypes.py
@@ -308,6 +308,28 @@ class SCFReport(TypedDict, total=False):
     debug_info: SCFDebugInfo
 
 
+class DeltaState(TypedDict):
+    """Delta SCF state information."""
+
+    #: Band checked.
+    band: int
+    #: Original source.
+    file: str
+    #: State population.
+    pop: float
+    #: Spin state.
+    spin: int
+
+
+class DeltaSCFReport(TypedDict, total=False):
+    """Delta SCF total report."""
+
+    #: Source file
+    file: str
+    #: Checked states.
+    states: list[DeltaState]
+
+
 # Bonds
 
 class BondInfo(TypedDict):