fix(dftbplus): parse all geometry and force rows

dpdata/formats/dftbplus/output.py

@@ -38,42 +38,52 @@ def read_dftb_plus(
     symbols = None
     forces = None
     energy = None
+    natoms = None
     with open_file(fn_1) as f:
-        flag = 0
-        for line in f:
-            if flag == 1:
-                flag += 1
-            elif flag == 2:
-                components = line.split()
-                flag += 1
-            elif line.startswith("Geometry"):
-                flag = 1
-                coord = []
-                symbols = []
-            elif flag in (3, 4, 5, 6):
-                s = line.split()
-                components_num = int(s[1])
-                symbols.append(components[components_num - 1])
+        lines = iter(f)
+        for line in lines:
+            if not line.startswith("Geometry"):
+                continue
+            # GenFormat declares the atom count on the first line after the
+            # opening brace and the element table on the next line.  The old
+            # state machine hard-coded four geometry rows, so ammonia-sized
+            # fixtures passed while larger molecules were truncated.
+            count_line = next(lines).split()
+            natoms = int(count_line[0])
+            coordinate_mode = count_line[1].upper()
+            components = next(lines).split()
+            coord = []
+            symbols = []
+            for _ in range(natoms):
+                s = next(lines).split()
+                symbols.append(components[int(s[1]) - 1])
                 coord.append([float(s[2]), float(s[3]), float(s[4])])
-                flag += 1
-                if flag == 7:
-                    flag = 0
+            if coordinate_mode == "F":
+                # Fractional GenFormat coordinates are expressed in the
+                # lattice-vector basis that follows the atom records.
+                origin = np.array([float(value) for value in next(lines).split()])
+                lattice = np.array(
+                    [[float(value) for value in next(lines).split()] for _ in range(3)]
+                )
+                coord = (np.asarray(coord) @ lattice + origin).tolist()
+            elif coordinate_mode not in {"C", "S"}:
+                raise ValueError(
+                    f"unsupported GenFormat coordinate mode: {coordinate_mode}"
+                )
+            break
+    if natoms is None:
+        raise ValueError("GenFormat Geometry block not found in DFTB+ input")
     with open_file(fn_2) as f:
-        flag = 0
-        for line in f:
+        lines = iter(f)
+        for line in lines:
             if line.startswith("Total Forces"):
-                flag = 8
                 forces = []
-            elif flag in (8, 9, 10, 11):
-                s = line.split()
-                forces.append([float(s[1]), float(s[2]), float(s[3])])
-                flag += 1
-                if flag == 12:
-                    flag = 0
+                for _ in range(natoms):
+                    s = next(lines).split()
+                    forces.append([float(s[1]), float(s[2]), float(s[3])])
             elif line.startswith("Total energy:"):
                 s = line.split()
                 energy = float(s[2])
-                flag = 0
 
     symbols = np.array(symbols)
     forces = np.array(forces)

tests/test_dftbplus.py

@@ -1,11 +1,14 @@
 from __future__ import annotations
 
 import unittest
+from io import StringIO
 
 import numpy as np
 from comp_sys import CompLabeledSys, IsNoPBC
 from context import dpdata
 
+from dpdata.formats.dftbplus.output import read_dftb_plus
+
 
 class TestDeepmdLoadAmmonia(unittest.TestCase, CompLabeledSys, IsNoPBC):
     def setUp(self):
@@ -55,6 +58,68 @@ def setUp(self):
         self.f_places = 6
         self.v_places = 6
 
+    def test_parser_reads_more_than_four_atoms(self):
+        input_text = """Geometry = GenFormat {
+5 C
+C H
+1 1 0.0 0.0 0.0
+2 2 1.0 0.0 0.0
+3 2 0.0 1.0 0.0
+4 2 0.0 0.0 1.0
+5 1 1.0 1.0 1.0
+}
+"""
+        output_text = """Total energy: -1.0 H
+Total Forces
+1 0.1 0.2 0.3
+2 0.4 0.5 0.6
+3 0.7 0.8 0.9
+4 1.0 1.1 1.2
+5 1.3 1.4 1.5
+"""
+        symbols, coords, energy, forces = read_dftb_plus(
+            StringIO(input_text), StringIO(output_text)
+        )
+        self.assertEqual(len(symbols), 5)
+        self.assertEqual(coords.shape, (5, 3))
+        self.assertEqual(forces.shape, (5, 3))
+        self.assertEqual(energy, -1.0)
+
+    def test_parser_converts_fractional_genformat_coordinates(self):
+        input_text = """Geometry = GenFormat {
+2 F
+H O
+1 1 0.5 0.0 0.0
+2 2 0.0 0.5 0.0
+0.1 0.2 0.3
+2.0 0.0 0.0
+0.0 4.0 0.0
+0.0 0.0 6.0
+}
+"""
+        output_text = """Total energy: -1.0 H
+Total Forces
+1 0.1 0.2 0.3
+2 0.4 0.5 0.6
+"""
+
+        symbols, coords, _, _ = read_dftb_plus(
+            StringIO(input_text), StringIO(output_text)
+        )
+
+        np.testing.assert_array_equal(symbols, ["H", "O"])
+        np.testing.assert_allclose(
+            coords,
+            [[1.1, 0.2, 0.3], [0.1, 2.2, 0.3]],
+        )
+
+    def test_parser_rejects_input_without_genformat_geometry(self):
+        with self.assertRaisesRegex(ValueError, "Geometry block not found"):
+            read_dftb_plus(
+                StringIO("Hamiltonian = DFTB {}\n"),
+                StringIO("Total energy: -1.0 H\n"),
+            )
+
 
 if __name__ == "__main__":
     unittest.main()