From 376317dabaa14b84992fde5071cba8560c485337 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 23 Jul 2026 01:18:50 +0800 Subject: [PATCH] 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])