Skip to content

Refactor trajectory handling to explicit frame-index execution#344

Open
harryswift01 wants to merge 17 commits into
mainfrom
306-frame-index-refactor
Open

Refactor trajectory handling to explicit frame-index execution#344
harryswift01 wants to merge 17 commits into
mainfrom
306-frame-index-refactor

Conversation

@harryswift01
Copy link
Copy Markdown
Member

Summary

This PR refactors CodeEntropy’s trajectory handling to use explicit frame-index driven execution throughout the main analysis workflow.

Previously, several parts of the codebase relied on implicit MDAnalysis trajectory cursor state, physical frame slicing, or local reduced-trajectory frame indices. This made frame ownership difficult to reason about and could contribute to unstable or hard-to-debug regression differences.

This PR introduces a single explicit frame contract using FrameSelection and FrameSource, removes physical trajectory frame slicing from the main workflow path, and updates the DAG, neighbour search, conformational/dihedral analysis, and MDAnalysis utility code to follow that contract.

This prepares the codebase for future Dask frame-level parallelisation without implementing Dask in this PR.

Changes

Explicit frame abstraction

  • Added FrameSelection to define the selected absolute trajectory frame indices.
  • Added FrameSource as the single owner of MDAnalysis trajectory seeking.
  • Updated workflow shared data to include:
    • frame_selection
    • frame_source
    • explicit frame index aliases for compatibility.
  • Removed the old TrajectorySlice migration object.

Workflow frame handling

  • Removed physical trajectory frame slicing from the main entropy workflow.
  • _build_reduced_universe() now applies atom selection only.
  • Selected frames are now accessed explicitly through FrameSource.
  • Water entropy now receives frame bounds from FrameSelection.

Frame DAG execution

  • Updated LevelDAG to iterate over explicit frame indices from FrameSource.
  • Updated FrameGraph so frame-local execution owns trajectory seeking via frame_source.seek(frame_index).
  • Added clearer error handling for out-of-bounds frame access.
  • Preserved deterministic serial reduction of frame-local outputs.

Neighbour and search logic

  • Updated neighbour analysis to consume FrameSource and explicit frame indices.
  • Removed reliance on universe.trajectory.ts.frame.
  • Updated search caching to key off the explicit frame index instead of hidden trajectory cursor state.
  • Updated RAD and grid neighbour search paths to seek frames explicitly.

Conformational and dihedral analysis

  • Threaded FrameSelection through conformational-state construction.
  • Updated dihedral analysis to use explicit frame bounds.
  • Preserved correct local indexing for dihedral_results.results.angles.
  • Clarified the distinction between:
    • absolute MDAnalysis trajectory frame indices
    • local analysis result-array indices.

MDAnalysis utility cleanup

  • Removed AnalysisFromFunction usage from workflow-relevant utility paths.
  • Replaced implicit extraction with explicit indexed frame loops.
  • Added explicit reduced-universe construction from frame index lists.
  • Added validation for empty frame selections and invalid timeseries kinds.

Unit tests

  • Updated stale unit tests to use the new FrameSelection / FrameSource contract.
  • Added atomic tests for:
    • frame selection semantics
    • frame source seeking
    • empty frame selections
    • frame DAG out-of-bounds errors
    • dihedral frame-selection errors
    • neighbour empty-frame behaviour
    • MDAnalysis utility validation branches.

Impact

  • Makes trajectory frame ownership explicit and consistent.
  • Removes hidden dependency on MDAnalysis trajectory cursor state from the main workflow.
  • Reduces the risk of frame drift between workflow, DAG, neighbour, and dihedral calculations.
  • Makes regression behaviour more deterministic and easier to debug.
  • Provides a cleaner foundation for future Dask frame-level parallelisation.

@harryswift01 harryswift01 added this to the 2.3.0 milestone May 22, 2026
@harryswift01 harryswift01 self-assigned this May 22, 2026
@harryswift01 harryswift01 added the feature request New feature or request label May 22, 2026
@harryswift01 harryswift01 requested review from jimboid and skfegan May 22, 2026 10:19
Copy link
Copy Markdown
Member

@jimboid jimboid left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A solid improvement to how the frames are being handled in CE. Instead of relying on MDA frame pointers to implicitly address frames this now explicitly sets the frame index at many different parts of the code which will eliminate issues around which frame is actually being used in different parts of the code. This should tighten the numerical certainty of the code.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

feature request New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants