diff --git a/dpdata/formats/lammps/dump.py b/dpdata/formats/lammps/dump.py index 064621b7..80f9451a 100644 --- a/dpdata/formats/lammps/dump.py +++ b/dpdata/formats/lammps/dump.py @@ -27,18 +27,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,7 +198,7 @@ def _iter_frames(fp): frame = [] for raw_line in fp: line = raw_line.rstrip("\n") - if "ITEM: TIMESTEP" in line: + if _is_item_header(line, "TIMESTEP"): if frame: yield frame frame = [line] @@ -385,10 +402,150 @@ def get_spin(lines, spin_keys): return None +def _atom_line_error(line, head): + """Return why an ATOMS payload row is unusable, or ``None`` if valid.""" + keys = head.split() + words = line.split() + ncols = len(keys) - 2 + if len(words) < ncols: + return f"truncated atom line {line.strip()!r}" + + try: + id_idx = keys.index("id") - 2 + type_idx = keys.index("type") - 2 + except ValueError: + return "ATOMS header is missing an 'id' or 'type' column" + + coord_tp_and_sf = get_coordtype_and_scalefactor(keys) + if coord_tp_and_sf is None: + return "ATOMS header does not contain atomic coordinates" + coordtype, _, _ = coord_tp_and_sf + + try: + int(words[id_idx]) + int(words[type_idx]) + for key in coordtype: + float(words[keys.index(key) - 2]) + except (ValueError, IndexError): + return f"unparsable atom line {line.strip()!r}" + return None + + +def _box_line_error(line, head): + """Return why a BOX BOUNDS row is unusable, or ``None`` if valid.""" + words = line.split() + ncols = 3 if "xy xz yz" in head else 2 + if len(words) < ncols: + return f"truncated box bound line {line.strip()!r}" + + try: + for word in words: + float(word) + except ValueError: + return f"unparsable box bound line {line.strip()!r}" + 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, box_head = _get_block(frame_lines, "BOX BOUNDS") + if len(box_blk) < 3: + return f"only {len(box_blk)} of 3 box bound lines" + for ii in box_blk[:3]: + error = _box_line_error(ii, box_head) + if error is not None: + return error + + atoms_blk, head = _get_block(frame_lines, "ATOMS") + if len(atoms_blk) != natoms: + return f"{len(atoms_blk)} atom lines for {natoms} atoms" + for ii in atoms_blk: + error = _atom_line_error(ii, head) + if error is not None: + return error + 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] + + head = frame_lines[atoms_header] + payload_lines = 0 + for idx in range(atoms_header + 1, len(frame_lines)): + if ( + _is_data_line(frame_lines[idx]) + and _atom_line_error(frame_lines[idx], head) is None + ): + 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) @@ -444,26 +601,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 - frame_ends = marks[1:] + [len(dump_lines)] - frames = [dump_lines[start:end] for start, end in zip(marks, frame_ends)] - for index, frame in enumerate(frames): - try: - natoms = get_natoms(frame) - atoms_header = next( - ii for ii, line in enumerate(frame) if "ITEM: ATOMS" in line - ) - except (IndexError, StopIteration, UnboundLocalError, ValueError): - # Leave malformed frames intact so the existing parsers can report - # their precise structural error. - continue - # A dump frame ends after its declared atom records. Ignore unrelated - # trailer text after the final frame instead of parsing it as an atom. - frames[index] = frame[: atoms_header + natoms + 1] - return frames + # 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): diff --git a/tests/test_lammps_dump_incomplete.py b/tests/test_lammps_dump_incomplete.py new file mode 100644 index 00000000..ff19d3c5 --- /dev/null +++ b/tests/test_lammps_dump_incomplete.py @@ -0,0 +1,193 @@ +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_last_frame_with_trailer_is_skipped(self): + atoms = [ + i for i, line in enumerate(self.lines) if line.startswith("ITEM: ATOMS") + ] + damaged = self.lines[: atoms[-1] + 2] + damaged.append("Loop time of 0.5 on 1 procs") + path = self._write("truncated_with_trailer.dump", damaged) + + 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) + and "unparsable atom line" in str(ii.message) + for ii in caught + ), + "a non-atom trailer must not complete a truncated frame", + ) + + 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_malformed_box_bounds_frame_is_skipped(self): + boxes = [ + i + for i, line in enumerate(self.lines) + if line.startswith("ITEM: BOX BOUNDS") + ] + damaged = list(self.lines) + damaged[boxes[-1] + 1 : boxes[-1] + 4] = [ + "low high tilt", + "low high tilt", + "low high tilt", + ] + path = self._write("malformed_box.dump", damaged) + + 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) + and "unparsable box bound line" in str(ii.message) + for ii in caught + ), + "a frame with non-numeric box bounds must be skipped", + ) + + 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()