fix(lammps): tolerate incomplete and commented dump frames

dpdata/formats/lammps/dump.py

@@ -26,18 +26,35 @@ class UnwrapWarning(UserWarning):
 warnings.simplefilter("once", UnwrapWarning)
 
 
+def _is_data_line(line):
+    """Tell payload lines apart from blanks and ``#`` comments.
+
+    LAMMPS itself never writes either, but concatenated or post-processed
+    trajectories often carry them, and every consumer of a block parses its
+    lines as numbers.
+    """
+    stripped = line.strip()
+    return bool(stripped) and not stripped.startswith("#")
+
+
+def _is_item_header(line, key=None):
+    """Return whether a non-comment line starts with a LAMMPS ITEM header."""
+    prefix = "ITEM:" if key is None else f"ITEM: {key}"
+    return _is_data_line(line) and line.lstrip().startswith(prefix)
+
+
 def _get_block(lines, key):
     for idx in range(len(lines)):
-        if ("ITEM: " + key) in lines[idx]:
+        if _is_item_header(lines[idx], key):
             break
     idx_s = idx + 1
     for idx in range(idx_s, len(lines)):
-        if ("ITEM: ") in lines[idx]:
+        if _is_item_header(lines[idx]):
             break
     idx_e = idx
     if idx_e == len(lines) - 1:
         idx_e += 1
-    return lines[idx_s:idx_e], lines[idx_s - 1]
+    return [ii for ii in lines[idx_s:idx_e] if _is_data_line(ii)], lines[idx_s - 1]
 
 
 def get_atype(lines, type_idx_zero=False):
@@ -181,14 +198,15 @@ def load_file(fname: FileType, begin=0, step=1):
     cc = -1
     with open_file(fname) as fp:
         while True:
-            line = fp.readline().rstrip("\n")
-            if not line:
+            raw_line = fp.readline()
+            if raw_line == "":
                 if cc >= begin and (cc - begin) % step == 0:
                     lines += buff
                     buff = []
                 cc += 1
                 return lines
-            if "ITEM: TIMESTEP" in line:
+            line = raw_line.rstrip("\n")
+            if _is_item_header(line, "TIMESTEP"):
                 if cc >= begin and (cc - begin) % step == 0:
                     lines += buff
                     buff = []
@@ -279,10 +297,98 @@ def get_spin(lines, spin_keys):
         return None
 
 
+def _describe_incomplete_frame(frame_lines):
+    """Return why a dump frame is unusable, or ``None`` when it is intact.
+
+    A run killed mid-write leaves a final frame whose sections are missing or
+    short. Reporting that here keeps the failure legible instead of surfacing
+    as a ragged-array error deep inside the coordinate readers.
+    """
+    for key in ("NUMBER OF ATOMS", "BOX BOUNDS", "ATOMS"):
+        if not any(_is_item_header(line, key) for line in frame_lines):
+            return f"missing the 'ITEM: {key}' section"
+
+    natoms_blk, _ = _get_block(frame_lines, "NUMBER OF ATOMS")
+    if not natoms_blk:
+        return "empty 'ITEM: NUMBER OF ATOMS' section"
+    try:
+        natoms = int(natoms_blk[0])
+    except ValueError:
+        return f"unparsable atom count {natoms_blk[0].strip()!r}"
+
+    box_blk, _ = _get_block(frame_lines, "BOX BOUNDS")
+    if len(box_blk) < 3:
+        return f"only {len(box_blk)} of 3 box bound lines"
+
+    atoms_blk, head = _get_block(frame_lines, "ATOMS")
+    if len(atoms_blk) != natoms:
+        return f"{len(atoms_blk)} atom lines for {natoms} atoms"
+    ncols = len(head.split()) - 2
+    for ii in atoms_blk:
+        if len(ii.split()) < ncols:
+            return f"truncated atom line {ii.strip()!r}"
+    return None
+
+
+def _drop_incomplete_frames(array_lines):
+    """Keep the intact frames, warning once per frame that is dropped."""
+    kept = []
+    for idx, frame_lines in enumerate(array_lines):
+        reason = _describe_incomplete_frame(frame_lines)
+        if reason is None:
+            kept.append(frame_lines)
+        else:
+            warnings.warn(
+                f"incomplete frame {idx} in the dump file ({reason}); it is ignored"
+            )
+    return kept
+
+
+def _clamp_after_atom_payload(frame_lines):
+    """Discard non-ITEM trailers after the declared atom payload.
+
+    Blank and comment lines may appear inside a post-processed ATOMS block,
+    so the physical line count cannot locate its end. Count only payload lines
+    and leave short frames untouched for ``_describe_incomplete_frame`` to
+    diagnose.
+    """
+    natoms_blk, _ = _get_block(frame_lines, "NUMBER OF ATOMS")
+    if not natoms_blk:
+        return frame_lines
+    try:
+        natoms = int(natoms_blk[0])
+    except ValueError:
+        return frame_lines
+
+    atoms_header = next(
+        (idx for idx, line in enumerate(frame_lines) if _is_item_header(line, "ATOMS")),
+        None,
+    )
+    if atoms_header is None:
+        return frame_lines
+    if natoms == 0:
+        return frame_lines[: atoms_header + 1]
+
+    payload_lines = 0
+    for idx in range(atoms_header + 1, len(frame_lines)):
+        if _is_data_line(frame_lines[idx]):
+            payload_lines += 1
+            if payload_lines == natoms:
+                return frame_lines[: idx + 1]
+    return frame_lines
+
+
 def system_data(
     lines, type_map=None, type_idx_zero=True, unwrap=False, input_file=None
 ):
     array_lines = split_traj(lines)
+    if array_lines is None:
+        raise RuntimeError(
+            "no 'ITEM: TIMESTEP' marker found; this is not a LAMMPS dump file"
+        )
+    array_lines = _drop_incomplete_frames(array_lines)
+    if not array_lines:
+        raise RuntimeError("no complete frame found in the dump file")
     lines = array_lines[0]
     system = {}
     system["atom_numbs"] = get_natoms_vec(lines)
@@ -338,21 +444,18 @@ def system_data(
 def split_traj(dump_lines):
     marks = []
     for idx, ii in enumerate(dump_lines):
-        if "ITEM: TIMESTEP" in ii:
+        if _is_item_header(ii, "TIMESTEP"):
             marks.append(idx)
     if len(marks) == 0:
         return None
-    elif len(marks) == 1:
-        return [dump_lines]
-    else:
-        block_size = marks[1] - marks[0]
-        ret = []
-        for ii in marks:
-            ret.append(dump_lines[ii : ii + block_size])
-        # for ii in range(len(marks)-1):
-        #     assert(marks[ii+1] - marks[ii] == block_size)
-        return ret
-    return None
+    # Slice every frame at the next marker instead of reusing the first
+    # frame's length. A truncated final frame, or any extra line anywhere in
+    # the file, otherwise shifts each later frame out of alignment.
+    bounds = [*marks, len(dump_lines)]
+    return [
+        _clamp_after_atom_payload(dump_lines[bounds[ii] : bounds[ii + 1]])
+        for ii in range(len(marks))
+    ]
 
 
 def from_system_data(system, f_idx=0, timestep=0):

tests/test_lammps_dump_incomplete.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+
+import os
+import shutil
+import tempfile
+import unittest
+import warnings
+
+import numpy as np
+from context import dpdata
+
+SOURCE = os.path.join("poscars", "conf.5.dump")
+
+
+class TestLmpDumpIncomplete(unittest.TestCase):
+    """A dump file may be truncated mid-frame or carry stray comment lines."""
+
+    def setUp(self):
+        with open(SOURCE) as fp:
+            self.lines = fp.read().split("\n")
+        self.tmp_dir = tempfile.mkdtemp()
+
+    def tearDown(self):
+        shutil.rmtree(self.tmp_dir, ignore_errors=True)
+
+    def _write(self, name, lines):
+        path = os.path.join(self.tmp_dir, name)
+        with open(path, "w") as fp:
+            fp.write("\n".join(lines))
+        return path
+
+    def _load(self, path, **kwargs):
+        return dpdata.System(path, fmt="lammps/dump", type_map=["O", "H"], **kwargs)
+
+    def test_complete_file_is_unchanged(self):
+        self.assertEqual(self._load(SOURCE).get_nframes(), 5)
+
+    def test_truncated_last_frame_is_skipped(self):
+        # A run killed while writing leaves the final frame without its atom
+        # lines; the earlier frames are still usable.
+        path = self._write("truncated.dump", self.lines[:-3])
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            system = self._load(path)
+        self.assertEqual(system.get_nframes(), 4)
+        self.assertTrue(
+            any("incomplete frame 4" in str(ii.message) for ii in caught),
+            "dropping a frame must be reported",
+        )
+
+    def test_truncated_frame_does_not_shift_later_frames(self):
+        # Frames are sliced at their own markers, so a short frame in the
+        # middle must not push the remaining frames out of alignment.
+        starts = [i for i, ll in enumerate(self.lines) if "ITEM: TIMESTEP" in ll]
+        damaged = self.lines[: starts[1] + 4] + self.lines[starts[2] :]
+        path = self._write("short_middle.dump", damaged)
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            system = self._load(path)
+        reference = self._load(SOURCE)
+        self.assertEqual(system.get_nframes(), 4)
+        self.assertTrue(any("incomplete frame 1" in str(ii.message) for ii in caught))
+        # Frame 1 is the damaged one; frames 2..4 must survive intact.
+        np.testing.assert_allclose(
+            system["coords"][1:], reference["coords"][2:], atol=1e-10
+        )
+
+    def test_comment_lines_are_ignored(self):
+        atoms = [i for i, ll in enumerate(self.lines) if ll.startswith("ITEM: ATOMS")]
+        annotated = list(self.lines)
+        for offset, idx in enumerate(atoms):
+            annotated.insert(idx + 1 + offset, "##### restart marker")
+        path = self._write("commented.dump", annotated)
+        system = self._load(path)
+        reference = self._load(SOURCE)
+        self.assertEqual(system.get_nframes(), reference.get_nframes())
+        np.testing.assert_allclose(system["coords"], reference["coords"], atol=1e-10)
+
+    def test_blank_lines_are_ignored(self):
+        atoms = [i for i, ll in enumerate(self.lines) if ll.startswith("ITEM: ATOMS")]
+        separated = list(self.lines)
+        for offset, idx in enumerate(atoms):
+            separated.insert(idx + 1 + offset, "")
+        path = self._write("blank_lines.dump", separated)
+        system = self._load(path)
+        reference = self._load(SOURCE)
+        self.assertEqual(system.get_nframes(), reference.get_nframes())
+        np.testing.assert_allclose(system["coords"], reference["coords"], atol=1e-10)
+
+    def test_atom_comments_and_non_item_trailer_preserve_last_frame(self):
+        atoms = [
+            i for i, line in enumerate(self.lines) if line.startswith("ITEM: ATOMS")
+        ]
+        annotated = list(self.lines)
+        annotated[atoms[-1] + 1 : atoms[-1] + 1] = ["", "# final frame annotation"]
+        annotated.extend(["Loop time of 0.5 on 1 procs", "END OF RUN"])
+        path = self._write("annotated_with_trailer.dump", annotated)
+
+        system = self._load(path)
+        reference = self._load(SOURCE)
+        self.assertEqual(system.get_nframes(), reference.get_nframes())
+        np.testing.assert_allclose(system["coords"], reference["coords"], atol=1e-10)
+
+    def test_item_like_comment_is_not_a_frame_boundary(self):
+        atoms = next(
+            i for i, line in enumerate(self.lines) if line.startswith("ITEM: ATOMS")
+        )
+        annotated = list(self.lines)
+        annotated.insert(atoms + 1, "# ITEM: TIMESTEP is documentation, not a frame")
+        path = self._write("commented_timestep.dump", annotated)
+
+        system = self._load(path, begin=1, step=2)
+        reference = self._load(SOURCE, begin=1, step=2)
+        self.assertEqual(system.get_nframes(), reference.get_nframes())
+        np.testing.assert_allclose(system["coords"], reference["coords"], atol=1e-10)
+
+    def test_item_like_comment_before_atoms_header_is_ignored(self):
+        atoms = next(
+            i for i, line in enumerate(self.lines) if line.startswith("ITEM: ATOMS")
+        )
+        annotated = list(self.lines)
+        annotated.insert(atoms, "# ITEM: ATOMS id type x y z")
+        path = self._write("commented_atoms_header.dump", annotated)
+
+        system = self._load(path)
+        reference = self._load(SOURCE)
+        self.assertEqual(system.get_nframes(), reference.get_nframes())
+        np.testing.assert_allclose(system["coords"], reference["coords"], atol=1e-10)
+
+    def test_no_usable_frame_raises(self):
+        starts = [i for i, ll in enumerate(self.lines) if "ITEM: TIMESTEP" in ll]
+        path = self._write("headers_only.dump", self.lines[: starts[0] + 4])
+        with self.assertRaisesRegex(RuntimeError, "no complete frame"):
+            self._load(path)
+
+    def test_not_a_dump_file_raises(self):
+        path = self._write("garbage.dump", ["not a dump file", "1 2 3"])
+        with self.assertRaisesRegex(RuntimeError, "not a LAMMPS dump file"):
+            self._load(path)
+
+
+if __name__ == "__main__":
+    unittest.main()