Replace CASTEP_MD parser with castep_outputs

MDANSE/Src/MDANSE/Framework/Converters/CASTEP.py

@@ -116,11 +116,11 @@ def run_step(self, index: int) -> tuple[int, None]:
         # Read the information in the frame
         time_step = frame["time"]
 
-        unit_cell = UnitCell(np.vstack(frame["h"]))
-        coords = np.vstack(tuple(arr[1] for arr in frame["R"]))
+        unit_cell = UnitCell(frame["h"])
+        coords = frame["R"]
         variables = {
-            "velocities": np.vstack(tuple(arr[1] for arr in frame["V"])),
-            "gradients": np.vstack(tuple(arr[1] for arr in frame["F"])),
+            "velocities": frame["V"],
+            "gradients": frame["F"],
         }
 
         conf = PeriodicRealConfiguration(

MDANSE/Src/MDANSE/Framework/Parsers/CASTEP_MD.py

@@ -24,6 +24,7 @@
 import numpy as np
 import numpy.typing as npt
 from ase.io.trajectory import Trajectory as ASETrajectory
+from castep_outputs.tools.md_geom_parser import MDGeomParser, MDGeomTimestepInfo
 from more_itertools import first, one, split_at
 
 from MDANSE.Framework.AtomMapping import AtomLabel
@@ -38,27 +39,6 @@
 
 
 class CASTEPMDFile(Parser):
-    @staticmethod
-    def float_list_handler(x: str) -> npt.NDArray[float]:
-        return np.array(x.split(), dtype=np.float64)
-
-    @staticmethod
-    def atom_prop_handler(x: str) -> tuple[tuple[str, int], list[float]]:
-        x = x.split()
-        return (x[0], int(x[1])), np.array(x[2:], dtype=np.float64)
-
-    KEY_HANDLERS = {
-        "h": float_list_handler,
-        "hv": float_list_handler,
-        "E": float_list_handler,
-        "T": float,
-        "P": float,
-        "S": float_list_handler,
-        "R": atom_prop_handler,
-        "V": atom_prop_handler,
-        "F": atom_prop_handler,
-    }
-
     UNIT_CONV = defaultdict(
         lambda: 1.0,
         h=BOHR,
@@ -73,41 +53,28 @@ def atom_prop_handler(x: str) -> tuple[tuple[str, int], list[float]]:
 
     @property
     def frames(self) -> Iterator:
-        with open(self.filename, encoding="utf-8") as castep_file:
-            file = map(str.strip, castep_file)
-
-            # Skip header
-            file = dropwhile(lambda line: not line.startswith("END"), file)
-            next(file)
+        parser = MDGeomParser(self.filename)
 
-            file = split_at(file, lambda line: not line)
-            file = filter(None, file)
+        yield from map(self.parse_frame, parser)
 
-            yield from map(self.parse_frame, file)
-
-    @staticmethod
-    def parse_frame(frame: Iterable[str]) -> dict[str, Any]:
+    def parse_frame(self, frame: MDGeomTimestepInfo) -> dict[str, Any]:
         accum = defaultdict(list)
 
-        for line in frame:
-            data, *key = line.split("<--")
-
-            if not key:
-                accum["time"] = float(data) * AUT
-            else:
-                key = one(key).strip()
-                if key not in "RVF":
-                    accum[key].append(
-                        CASTEPMDFile.KEY_HANDLERS[key](data)
-                        * CASTEPMDFile.UNIT_CONV[key]
-                    )
-                else:
-                    specind, val = CASTEPMDFile.KEY_HANDLERS[key](data)
-                    accum[key].append((specind, val * CASTEPMDFile.UNIT_CONV[key]))
+        accum["time"] = frame["time"] * AUT
+        accum["h"] = np.array(frame["lattice_vectors"]) * self.UNIT_CONV["h"]
+        for prop in "RVF":
+            accum[prop] = (
+                np.array([ion[prop] for ion in frame["ions"].values()])
+                * self.UNIT_CONV[prop]
+            )
 
         return accum
 
+    @cached_property
+    def n_frames(self) -> int:
+        return len(MDGeomParser(self.filename))
+
     @cached_property
     def element_list(self) -> Collection[str]:
-        first_frame = first(self.frames)
-        return [spec for (spec, ind), _ in first_frame["R"]]
+        first_frame = MDGeomParser(self.filename)[0]
+        return [spec for spec, _ind in first_frame["ions"]]

MDANSE/pyproject.toml

@@ -43,6 +43,7 @@ dependencies = [
     "more-itertools",
     "typing_extensions",
     "setuptools-git-versioning>=3.0,<4",
+    "castep_outputs==0.3.3"
 ]
 
 [dependency-groups]