fix(dftbplus): parse all atoms instead of truncating to 4

dpdata/formats/dftbplus/output.py

@@ -40,8 +40,12 @@ def read_dftb_plus(
     energy = None
     with open_file(fn_1) as f:
         flag = 0
+        n_atoms = 0
         for line in f:
             if flag == 1:
+                # Header line: "N C" where N is atom count, C is coordinate type
+                header = line.split()
+                n_atoms = int(header[0])
                 flag += 1
             elif flag == 2:
                 components = line.split()
@@ -50,25 +54,39 @@ def read_dftb_plus(
                 flag = 1
                 coord = []
                 symbols = []
-            elif flag in (3, 4, 5, 6):
+                n_atoms = 0
+            elif flag == 3 and n_atoms > 0:
+                # Parse exactly n_atoms coordinate rows
                 s = line.split()
                 components_num = int(s[1])
                 symbols.append(components[components_num - 1])
                 coord.append([float(s[2]), float(s[3]), float(s[4])])
-                flag += 1
-                if flag == 7:
+                n_atoms -= 1
+                if n_atoms == 0:
                     flag = 0
     with open_file(fn_2) as f:
         flag = 0
+        n_forces = 0
         for line in f:
             if line.startswith("Total Forces"):
                 flag = 8
                 forces = []
-            elif flag in (8, 9, 10, 11):
+                n_forces = 0
+            elif flag == 8:
+                # First line after "Total Forces" header contains force data
+                # We don't know the count yet, so parse until we hit a non-force line
                 s = line.split()
-                forces.append([float(s[1]), float(s[2]), float(s[3])])
-                flag += 1
-                if flag == 12:
+                if len(s) >= 4:
+                    try:
+                        # Try to parse as force row: "index fx fy fz"
+                        int(s[0])
+                        float(s[1])
+                        forces.append([float(s[1]), float(s[2]), float(s[3])])
+                        n_forces += 1
+                    except (ValueError, IndexError):
+                        # Not a force row, we're done
+                        flag = 0
+                else:
                     flag = 0
             elif line.startswith("Total energy:"):
                 s = line.split()

tests/dftbplus/detailed_6atoms.out

@@ -0,0 +1,11 @@
+SCC converged
+
+Total Forces
+    1      0.010000000000      0.002000000000      0.003000000000
+    2     -0.015000000000      0.001000000000      0.002000000000
+    3      0.005000000000     -0.008000000000      0.001000000000
+    4      0.003000000000      0.004000000000     -0.006000000000
+    5     -0.002000000000      0.003000000000      0.005000000000
+    6      0.001000000000     -0.001000000000     -0.002000000000
+
+Total energy:                       -12.3456789012 H         -335.9999 eV

tests/dftbplus/dftb_pin_6atoms.hsd

@@ -0,0 +1,14 @@
+Geometry = GenFormat {
+6 C
+C H O N S Cl
+1     1        1.000000        0.000000        0.000000
+2     2        1.500000        0.000000        0.000000
+3     3        2.000000        1.000000        0.000000
+4     4        2.500000        0.000000        1.000000
+5     5        3.000000        1.000000        1.000000
+6     6        3.500000        0.500000        0.500000
+}
+Driver = {}
+Hamiltonian = DFTB {
+  SCC = Yes
+}

tests/test_dftbplus_6atoms.py

@@ -0,0 +1,66 @@
+"""Test that DFTB+ parser handles systems with more than 4 atoms."""
+
+from __future__ import annotations
+
+import os
+import sys
+import unittest
+
+import numpy as np
+
+# Add parent directory to path for imports
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from dpdata.formats.dftbplus.output import read_dftb_plus
+
+
+class TestDftbPlusSixAtoms(unittest.TestCase):
+    """Test DFTB+ parser with a 6-atom system (beyond the original 4-atom limit)."""
+
+    def test_six_atoms_parsed_correctly(self):
+        """Verify that all 6 atoms are parsed, not truncated to 4."""
+        symbols, coord, energy, forces = read_dftb_plus(
+            "dftbplus/dftb_pin_6atoms.hsd",
+            "dftbplus/detailed_6atoms.out",
+        )
+
+        # Should have 6 atoms, not 4
+        self.assertEqual(len(symbols), 6)
+        self.assertEqual(coord.shape, (6, 3))
+        self.assertEqual(forces.shape, (6, 3))
+
+        # Check atom symbols
+        expected_symbols = ["C", "H", "O", "N", "S", "Cl"]
+        np.testing.assert_array_equal(symbols, expected_symbols)
+
+        # Check coordinates
+        expected_coord = np.array(
+            [
+                [1.0, 0.0, 0.0],
+                [1.5, 0.0, 0.0],
+                [2.0, 1.0, 0.0],
+                [2.5, 0.0, 1.0],
+                [3.0, 1.0, 1.0],
+                [3.5, 0.5, 0.5],
+            ]
+        )
+        np.testing.assert_array_almost_equal(coord, expected_coord)
+
+        # Check forces
+        expected_forces = np.array(
+            [
+                [0.01, 0.002, 0.003],
+                [-0.015, 0.001, 0.002],
+                [0.005, -0.008, 0.001],
+                [0.003, 0.004, -0.006],
+                [-0.002, 0.003, 0.005],
+                [0.001, -0.001, -0.002],
+            ]
+        )
+        np.testing.assert_array_almost_equal(forces, expected_forces)
+
+        # Check energy
+        self.assertAlmostEqual(energy, -12.3456789012)
+
+
+if __name__ == "__main__":
+    unittest.main()