feat(lammps): load selected dump frames efficiently#1040
Conversation
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
Merging this PR will not alter performance
|
dosubot
Bot
added
size:L
There was a problem hiding this comment.
Pull request overview
This PR adds efficient sparse frame loading for the lammps/dump format by introducing f_idx selection, enabling callers to load only specific non-negative frame indices without parsing and storing the full trajectory.
Changes:
- Added
f_idxsupport toLAMMPSDumpFormat.from_systemand propagated it down to the dump loader. - Implemented streaming frame iteration and index normalization/validation in
dpdata/formats/lammps/dump.pyto load only requested frames and stop early when possible. - Updated trajectory splitting to respect actual
ITEM: TIMESTEPboundaries, and added tests covering order/duplicates, scalar index input, and validation errors.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| tests/test_lammps_dump_skipload.py | Adds tests for f_idx selection semantics, validation, and early-stop behavior. |
| dpdata/plugins/lammps.py | Extends the lammps/dump plugin API to accept f_idx and forwards it to the loader. |
| dpdata/formats/lammps/dump.py | Implements streaming frame parsing, f_idx validation/selection, and robust frame splitting by TIMESTEP markers. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| def __next__(self): | ||
| self.lines_read += 1 | ||
| return super().__next__() |
📝 WalkthroughWalkthroughThe LAMMPS dump reader now supports explicit sparse frame selection, streaming reads with early termination, clearer index validation, and marker-based trajectory splitting. The LAMMPS plugin forwards the new selection parameter, with tests covering selection behavior and errors. ChangesLAMMPS frame selection
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant LAMMPSDumpFormat
participant load_file
participant _iter_frames
Caller->>LAMMPSDumpFormat: request f_idx
LAMMPSDumpFormat->>load_file: forward f_idx
load_file->>_iter_frames: read frames sequentially
_iter_frames-->>load_file: yield frame lines
load_file-->>LAMMPSDumpFormat: return selected frames
LAMMPSDumpFormat-->>Caller: return system data
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/test_lammps_dump_skipload.py`:
- Around line 14-23: Update CountingStringIO to track maximum stream position
across read, readline, readlines, and __next__, rather than only incrementing
lines_read in __next__. In the early-termination assertion around the affected
test, verify the recorded maximum position remains below len(content),
preserving the probe’s intended detection of any full-stream consumption.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: fc0e87d7-32b7-4dcf-bb42-971836541fe3
📒 Files selected for processing (3)
dpdata/formats/lammps/dump.pydpdata/plugins/lammps.pytests/test_lammps_dump_skipload.py
| 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__() |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Make the early-termination probe detect all stream consumption.
lines_read only tracks __next__; a full read()/readline() regression would still pass line 85. Track the maximum stream position from read, readline, readlines, and __next__, then assert it remains below len(content).
Also applies to: 76-85
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/test_lammps_dump_skipload.py` around lines 14 - 23, Update
CountingStringIO to track maximum stream position across read, readline,
readlines, and __next__, rather than only incrementing lines_read in __next__.
In the early-termination assertion around the affected test, verify the recorded
maximum position remains below len(content), preserving the probe’s intended
detection of any full-stream consumption.
Summary
f_idxsupport tolammps/dumpso callers can load arbitrary non-negative frame indices directlySystem.sub_systembegin/stepselectionsITEM: TIMESTEPboundaries instead of assuming a fixed block lengthExample:
Performance benchmark
The benchmark compares direct sparse loading against the existing workflow of loading the complete trajectory and then calling
sub_system(indices).System(..., f_idx=indices)System(...).sub_system(indices)Direct sparse loading was 14.38x faster and reduced peak RSS by 64.39% for this workload. Both methods returned the same 10 selected frames. Because the final frame was selected, both methods traversed the full file; the improvement comes from avoiding storage and parsing of unselected frames.
Validation
python -m unittest test_lammps_dump_skipload.py test_lammps_dump_to_system.py test_lammps_dump_unfold.py test_lammps_dump_shift_origin.py test_lammps_dump_idx.py test_lammps_read_from_trajs.py test_lammps_spin.py— 70 tests passedruff check dpdata/ tests/test_lammps_dump_skipload.py— passedruff format --check dpdata/ tests/test_lammps_dump_skipload.py— passeddpdata --helpanddpdata --version— passedparmeddependency is not installed in the environment, and 43 tests were skippedCloses #367.
Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh
Summary by CodeRabbit
New Features
Bug Fixes
Tests