feat(lammps): load selected dump frames efficiently

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):

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
         )

tests/test_lammps_dump_skipload.py

@@ -1,12 +1,27 @@
 from __future__ import annotations
 
+import io
 import os
 import unittest
 
 import numpy as np
 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,69 @@ 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])
+        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"):
+            dump.load_file(self.dump_file, begin=1, f_idx=[2])