Skip to content

Cache customised united-atom axes topology for frame covariance#373

Merged
harryswift01 merged 3 commits into
365-add-dask-execution-for-conformationdagfrom
372-cache-customised-united-atom-axes
Jun 24, 2026
Merged

Cache customised united-atom axes topology for frame covariance#373
harryswift01 merged 3 commits into
365-add-dask-execution-for-conformationdagfrom
372-cache-customised-united-atom-axes

Conversation

@harryswift01

Copy link
Copy Markdown
Member

Summary

This PR improves customised united-atom covariance performance by caching static axes-topology metadata before frame execution. The frame-local covariance loop can now reuse cached atom-index relationships instead of repeatedly rediscovering the same bonded atom topology with MDAnalysis selections.

Changes

Cached united-atom axes topology:

  • Added static axes-topology metadata for customised united-atom axes.
  • Added BuildAxesTopologyNode to construct shared_data["axes_topology"].
  • Cached only frame-independent topology data, including UA heavy atom indices, bonded heavy/light atom indices, UA atom membership, residue heavy atom indices, and residue UA masses.

Cached axes execution path:

  • Added cached axes methods in AxesCalculator for computing customised UA axes from cached topology.
  • Updated FrameCovarianceNode to use the cached topology path when available.
  • Preserved the existing get_UA_axes(...) path as a fallback.
  • Kept all frame-dependent values, such as positions, forces, centres, axes, torques, moments of inertia, and covariance matrices, computed per frame.

DAG integration and tests:

  • Registered BuildAxesTopologyNode in LevelDAG after BuildBeadsNode.
  • Added atomic unit tests for the axes-topology node.
  • Added coverage for cached UA axes and cached bonded-axes paths.
  • Updated covariance-node and LevelDAG tests to reflect the new cached topology dependency.
  • Confirmed existing entropy output remains unchanged for the profiled benzene benchmark.

Impact

  • Reduces repeated MDAnalysis topology selection work inside the frame covariance loop.
  • Preserves existing covariance, force, torque, entropy, Dask, and scheduling behaviour.
  • Provides a safe fallback to the previous axes path when cached topology is unavailable.

@harryswift01 harryswift01 added this to the 2.3.0 milestone Jun 24, 2026
@harryswift01 harryswift01 requested a review from jimboid June 24, 2026 09:28
@harryswift01 harryswift01 self-assigned this Jun 24, 2026
@harryswift01 harryswift01 added the feature request New feature or request label Jun 24, 2026

@jimboid jimboid left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

A nifty optimisation. PR is complete with all required changes clear and matching the description in the PR summary.

@harryswift01 harryswift01 merged commit 763079c into 365-add-dask-execution-for-conformationdag Jun 24, 2026
42 of 45 checks passed
@harryswift01 harryswift01 deleted the 372-cache-customised-united-atom-axes branch June 24, 2026 10:10
@harryswift01 harryswift01 linked an issue Jun 24, 2026 that may be closed by this pull request
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.

Cache customised united-atom axes topology for frame covariance

2 participants