Skip to content

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

Description

@harryswift01

Goal

Cache static united-atom axes topology so customised frame covariance does not repeatedly rediscover the same MDAnalysis topology inside the frame loop.

Background

After the conformational topology optimisation, profiling showed the next major bottleneck is customised UA axes construction inside FrameCovarianceNode.

The current covariance path already caches bead atom indices, but AxesCalculator.get_UA_axes() still repeatedly selects heavy atoms, bonded atoms, hydrogens, and UA masses for every frame. These relationships are static and can be cached once during LevelDAG setup.

Scope

  • Add a static BuildAxesTopologyNode.

  • Store cached topology in shared_data["axes_topology"].

  • Build the cache after BuildBeadsNode.

  • Cache only frame-independent UA topology:

    • heavy atom index
    • UA atom indices
    • bonded heavy atom indices
    • bonded hydrogen/light atom indices
    • residue heavy atom indices
    • residue UA masses
  • Add cached axes methods such as get_UA_axes_from_topology(...) and get_bonded_axes_from_topology(...).

  • Update FrameCovarianceNode to use the cached path when available.

  • Keep get_UA_axes(...) as the fallback.

  • Do not change force, torque, covariance, entropy, Dask, or scheduling logic.

  • Do not cache frame-dependent values such as positions, forces, axes, torques, or covariance matrices.

Metadata

Metadata

Assignees

Labels

Type

No fields configured for Task.

Projects

No projects

Relationships

None yet

Development

No branches or pull requests

Issue actions