From 376317dabaa14b84992fde5071cba8560c485337 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 23 Jul 2026 01:18:50 +0800 Subject: [PATCH 1/3] feat(lammps): load selected dump frames efficiently Add f_idx-based sparse frame loading for LAMMPS dump trajectories while preserving requested order and duplicate indices. Stop scanning after the final requested frame and validate invalid selections explicitly. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- dpdata/formats/lammps/dump.py | 151 +++++++++++++++++++++++------ dpdata/plugins/lammps.py | 9 +- tests/test_lammps_dump_skipload.py | 75 ++++++++++++++ 3 files changed, 205 insertions(+), 30 deletions(-) diff --git a/dpdata/formats/lammps/dump.py b/dpdata/formats/lammps/dump.py index 89e75e4de..257dd2c49 100644 --- a/dpdata/formats/lammps/dump.py +++ b/dpdata/formats/lammps/dump.py @@ -1,6 +1,7 @@ #!/usr/bin/env python3 from __future__ import annotations +import numbers import os import sys from typing import TYPE_CHECKING @@ -175,26 +176,127 @@ def box2dumpbox(orig, box): return bounds, tilt -def load_file(fname: FileType, begin=0, step=1): +def _iter_frames(fp): + """Yield one LAMMPS dump frame at a time without retaining skipped frames.""" + frame = [] + for raw_line in fp: + line = raw_line.rstrip("\n") + if "ITEM: TIMESTEP" in line: + if frame: + yield frame + frame = [line] + elif frame: + # Ignore any preamble before the first TIMESTEP marker, matching the + # historical loader behavior. + frame.append(line) + if frame: + yield frame + + +def _normalize_frame_indices(f_idx): + """Validate frame indices while preserving their order and duplicates.""" + if isinstance(f_idx, numbers.Integral) and not isinstance(f_idx, bool): + indices = [int(f_idx)] + else: + try: + indices = list(f_idx) + except TypeError as exc: + raise TypeError( + "f_idx must be an integer or an iterable of integers" + ) from exc + + if not indices: + raise ValueError("f_idx must not be empty") + + normalized = [] + for index in indices: + if not isinstance(index, numbers.Integral) or isinstance(index, bool): + raise TypeError("f_idx must contain only integers") + if index < 0: + raise ValueError("f_idx must contain only non-negative frame indices") + normalized.append(int(index)) + return normalized + + +def load_file( + fname: FileType, + begin=0, + step=1, + f_idx: int | list[int] | np.ndarray | None = None, +): + """Load selected frames from a LAMMPS dump file. + + Parameters + ---------- + fname : FileType + Dump file path or an open text file object. + begin : int, optional + First frame to load when ``f_idx`` is not provided. + step : int, optional + Interval between loaded frames when ``f_idx`` is not provided. + f_idx : int or array-like of int, optional + Specific non-negative frame indices to load. The requested order and + duplicate indices are preserved. This option cannot be combined with + non-default ``begin`` or ``step`` values. + + Returns + ------- + list[str] + Lines belonging to the selected frames. + + Raises + ------ + IndexError + If any requested frame index is outside the trajectory. + TypeError + If ``f_idx`` contains a non-integer value. + ValueError + If ``f_idx`` is empty, contains a negative value, or is combined with + non-default ``begin`` or ``step`` values. + """ + if f_idx is not None: + if begin != 0 or step != 1: + raise ValueError("f_idx cannot be combined with begin or step") + + frame_indices = _normalize_frame_indices(f_idx) + requested = set(frame_indices) + last_requested = max(requested) + selected = {} + nframes = 0 + + with open_file(fname) as fp: + for frame_index, frame in enumerate(_iter_frames(fp)): + nframes = frame_index + 1 + if frame_index in requested: + selected[frame_index] = frame + if frame_index == last_requested: + # The generator retains only the current frame, so stopping + # here avoids reading and parsing the remainder of the file. + break + + missing = sorted(requested.difference(selected)) + if missing: + raise IndexError( + f"Requested frame indices {missing} are out of range for " + f"a trajectory containing {nframes} frames" + ) + + lines = [] + for frame_index in frame_indices: + lines.extend(selected[frame_index]) + return lines + + if begin < 0: + raise ValueError("begin must be non-negative") + if step <= 0: + raise ValueError("step must be positive") + lines = [] - buff = [] - cc = -1 with open_file(fname) as fp: - while True: - line = fp.readline().rstrip("\n") - if not line: - if cc >= begin and (cc - begin) % step == 0: - lines += buff - buff = [] - cc += 1 - return lines - if "ITEM: TIMESTEP" in line: - if cc >= begin and (cc - begin) % step == 0: - lines += buff - buff = [] - cc += 1 - if cc >= begin and (cc - begin) % step == 0: - buff.append(line) + for frame_index, frame in enumerate(_iter_frames(fp)): + if frame_index >= begin and (frame_index - begin) % step == 0: + lines.extend(frame) + return lines def get_spin_keys(inputfile): @@ -342,17 +444,8 @@ def split_traj(dump_lines): 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 + frame_ends = marks[1:] + [len(dump_lines)] + return [dump_lines[start:end] for start, end in zip(marks, frame_ends)] def from_system_data(system, f_idx=0, timestep=0): diff --git a/dpdata/plugins/lammps.py b/dpdata/plugins/lammps.py index 9a4622659..fd5c543b6 100644 --- a/dpdata/plugins/lammps.py +++ b/dpdata/plugins/lammps.py @@ -142,6 +142,7 @@ def from_system( step: int = 1, unwrap: bool = False, input_file: str = None, + f_idx: int | list[int] | np.ndarray | None = None, **kwargs, ): """Read the data from a lammps dump file. @@ -160,13 +161,19 @@ def from_system( Whether to unwrap the coordinates input_file : str, optional The input file name + f_idx : int or array-like of int, optional + Specific non-negative frame indices to load. The requested order + and duplicate indices are preserved. Cannot be combined with + non-default ``begin`` or ``step`` values. Returns ------- dict The system data """ - lines = dpdata.formats.lammps.dump.load_file(file_name, begin=begin, step=step) + lines = dpdata.formats.lammps.dump.load_file( + file_name, begin=begin, step=step, f_idx=f_idx + ) data = dpdata.formats.lammps.dump.system_data( lines, type_map, unwrap=unwrap, input_file=input_file ) diff --git a/tests/test_lammps_dump_skipload.py b/tests/test_lammps_dump_skipload.py index 299e1db48..55f0b0283 100644 --- a/tests/test_lammps_dump_skipload.py +++ b/tests/test_lammps_dump_skipload.py @@ -1,5 +1,6 @@ from __future__ import annotations +import io import os import unittest @@ -7,6 +8,20 @@ from comp_sys import CompSys, IsPBC from context import dpdata +from dpdata.formats.lammps import dump + + +class CountingStringIO(io.StringIO): + """Track how many lines the trajectory reader consumes.""" + + def __init__(self, value): + super().__init__(value) + self.lines_read = 0 + + def __next__(self): + self.lines_read += 1 + return super().__next__() + class TestLmpDumpSkip(unittest.TestCase, CompSys, IsPBC): def setUp(self): @@ -20,3 +35,63 @@ def setUp(self): self.e_places = 6 self.f_places = 6 self.v_places = 4 + + +class TestLmpDumpFrameSelection(unittest.TestCase): + def setUp(self): + self.dump_file = os.path.join("poscars", "conf.5.dump") + self.type_map = ["O", "H"] + + def test_select_frames_preserves_order_and_duplicates(self): + all_frames = dpdata.System( + self.dump_file, fmt="lammps/dump", type_map=self.type_map + ) + frame_indices = np.array([4, 1, 4]) + + selected = dpdata.System( + self.dump_file, + fmt="lammps/dump", + type_map=self.type_map, + f_idx=frame_indices, + ) + expected = all_frames.sub_system(frame_indices) + + np.testing.assert_allclose(selected["coords"], expected["coords"]) + np.testing.assert_allclose(selected["cells"], expected["cells"]) + + def test_select_single_frame_by_integer(self): + selected = dpdata.System( + self.dump_file, + fmt="lammps/dump", + type_map=self.type_map, + f_idx=2, + ) + expected = dpdata.System( + self.dump_file, fmt="lammps/dump", type_map=self.type_map + )[2] + + np.testing.assert_allclose(selected["coords"], expected["coords"]) + np.testing.assert_allclose(selected["cells"], expected["cells"]) + + def test_file_object_is_read_once_and_stops_after_last_target(self): + with open(self.dump_file) as fp: + content = fp.read() + stream = CountingStringIO(content) + + lines = dump.load_file(stream, f_idx=[1]) + frames = dump.split_traj(lines) + + self.assertEqual(frames[0][1], "1") + self.assertLess(stream.lines_read, len(content.splitlines())) + + def test_invalid_frame_indices(self): + with self.assertRaisesRegex(ValueError, "must not be empty"): + dump.load_file(self.dump_file, f_idx=[]) + with self.assertRaisesRegex(ValueError, "non-negative"): + dump.load_file(self.dump_file, f_idx=[-1]) + with self.assertRaisesRegex(IndexError, "out of range"): + dump.load_file(self.dump_file, f_idx=[5]) + + def test_frame_indices_are_mutually_exclusive_with_slice(self): + with self.assertRaisesRegex(ValueError, "cannot be combined"): + dump.load_file(self.dump_file, begin=1, f_idx=[2]) From b68b96db615b9600ab5bc4940437bbd5d0902adb Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Mon, 27 Jul 2026 02:24:47 +0800 Subject: [PATCH 2/3] test(lammps): cover the f_idx TypeError paths Add the non-integer element case suggested in review, plus the bool rejection and the non-iterable scalar case, so both `TypeError` branches of `_normalize_frame_indices` are exercised. Co-Authored-By: Claude Opus 5 --- tests/test_lammps_dump_skipload.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/test_lammps_dump_skipload.py b/tests/test_lammps_dump_skipload.py index 55f0b0283..b333bcf25 100644 --- a/tests/test_lammps_dump_skipload.py +++ b/tests/test_lammps_dump_skipload.py @@ -91,6 +91,12 @@ def test_invalid_frame_indices(self): dump.load_file(self.dump_file, f_idx=[-1]) with self.assertRaisesRegex(IndexError, "out of range"): dump.load_file(self.dump_file, f_idx=[5]) + with self.assertRaisesRegex(TypeError, "only integers"): + dump.load_file(self.dump_file, f_idx=["a"]) + with self.assertRaisesRegex(TypeError, "only integers"): + dump.load_file(self.dump_file, f_idx=[True]) + with self.assertRaisesRegex(TypeError, "an iterable of integers"): + dump.load_file(self.dump_file, f_idx=1.5) def test_frame_indices_are_mutually_exclusive_with_slice(self): with self.assertRaisesRegex(ValueError, "cannot be combined"): From ce9dfa7f3b95ae393ee348c1745ccc246d235ce4 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Tue, 28 Jul 2026 21:16:05 +0800 Subject: [PATCH 3/3] fix(lammps): harden selective dump loading Ignore blank separators, clamp final-frame trailers, and strengthen the early-termination probe across read APIs. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- dpdata/formats/lammps/dump.py | 20 ++++++++++- tests/test_lammps_dump_skipload.py | 53 +++++++++++++++++++++++++++--- 2 files changed, 68 insertions(+), 5 deletions(-) diff --git a/dpdata/formats/lammps/dump.py b/dpdata/formats/lammps/dump.py index 257dd2c49..064621b77 100644 --- a/dpdata/formats/lammps/dump.py +++ b/dpdata/formats/lammps/dump.py @@ -185,6 +185,10 @@ def _iter_frames(fp): if frame: yield frame frame = [line] + elif not line: + # Blank separators are common in concatenated trajectories and do + # not belong to any numeric dump section. + continue elif frame: # Ignore any preamble before the first TIMESTEP marker, matching the # historical loader behavior. @@ -445,7 +449,21 @@ def split_traj(dump_lines): if len(marks) == 0: return None frame_ends = marks[1:] + [len(dump_lines)] - return [dump_lines[start:end] for start, end in zip(marks, frame_ends)] + 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 def from_system_data(system, f_idx=0, timestep=0): diff --git a/tests/test_lammps_dump_skipload.py b/tests/test_lammps_dump_skipload.py index b333bcf25..cda242a65 100644 --- a/tests/test_lammps_dump_skipload.py +++ b/tests/test_lammps_dump_skipload.py @@ -12,15 +12,41 @@ class CountingStringIO(io.StringIO): - """Track how many lines the trajectory reader consumes.""" + """Track every supported read strategy and whether it rewinds.""" def __init__(self, value): super().__init__(value) self.lines_read = 0 + self.max_position = 0 + self.seek_calls = 0 + + def _record_position(self): + self.max_position = max(self.max_position, self.tell()) def __next__(self): self.lines_read += 1 - return super().__next__() + value = super().__next__() + self._record_position() + return value + + def read(self, *args, **kwargs): + value = super().read(*args, **kwargs) + self._record_position() + return value + + def readline(self, *args, **kwargs): + value = super().readline(*args, **kwargs) + self._record_position() + return value + + def readlines(self, *args, **kwargs): + value = super().readlines(*args, **kwargs) + self._record_position() + return value + + def seek(self, *args, **kwargs): + self.seek_calls += 1 + return super().seek(*args, **kwargs) class TestLmpDumpSkip(unittest.TestCase, CompSys, IsPBC): @@ -73,7 +99,7 @@ def test_select_single_frame_by_integer(self): np.testing.assert_allclose(selected["coords"], expected["coords"]) np.testing.assert_allclose(selected["cells"], expected["cells"]) - def test_file_object_is_read_once_and_stops_after_last_target(self): + def test_file_object_stops_after_last_target_without_rewinding(self): with open(self.dump_file) as fp: content = fp.read() stream = CountingStringIO(content) @@ -82,7 +108,26 @@ def test_file_object_is_read_once_and_stops_after_last_target(self): frames = dump.split_traj(lines) self.assertEqual(frames[0][1], "1") - self.assertLess(stream.lines_read, len(content.splitlines())) + first_23_lines = "".join(content.splitlines(keepends=True)[:23]) + self.assertLessEqual(stream.max_position, len(first_23_lines)) + self.assertLess(stream.max_position, len(content)) + self.assertLessEqual(stream.lines_read, 23) + self.assertEqual(stream.seek_calls, 0) + + def test_blank_trailer_is_ignored_and_extra_text_is_clamped(self): + with open(self.dump_file) as fp: + content = fp.read().rstrip("\n") + "\n\n" + + lines = dump.load_file(io.StringIO(content)) + frames = dump.split_traj([*lines, "# end of run"]) + self.assertEqual(len(frames), 5) + self.assertFalse(any(not line for frame in frames for line in frame)) + self.assertNotIn("# end of run", frames[-1]) + + system = dpdata.System( + io.StringIO(content), fmt="lammps/dump", type_map=self.type_map + ) + self.assertEqual(system.get_nframes(), 5) def test_invalid_frame_indices(self): with self.assertRaisesRegex(ValueError, "must not be empty"):