Introduction of bond-derived axes for residue level rotation/united atom level translation

CodeEntropy/levels/nodes/covariance.py

@@ -79,7 +79,6 @@ def run(self, ctx: FrameCtx) -> dict[str, Any]:
         beads = shared["beads"]
         args = shared["args"]
         axes_manager = shared.get("axes_manager")
-        axes_topology = shared.get("axes_topology")
 
         fp = float(args.force_partitioning)
         combined = bool(getattr(args, "combined_forcetorque", False))
@@ -111,7 +110,6 @@ def run(self, ctx: FrameCtx) -> dict[str, Any]:
                         group_id=group_id,
                         beads=beads,
                         axes_manager=axes_manager,
-                        axes_topology=axes_topology,
                         box=box,
                         force_partitioning=fp,
                         customised_axes=customised_axes,
@@ -129,7 +127,6 @@ def run(self, ctx: FrameCtx) -> dict[str, Any]:
                         group_id=group_id,
                         beads=beads,
                         axes_manager=axes_manager,
-                        axes_topology=axes_topology,
                         box=box,
                         customised_axes=customised_axes,
                         force_partitioning=fp,
@@ -175,7 +172,6 @@ def _process_united_atom(
         group_id: int,
         beads: dict[Any, list[Any]],
         axes_manager: Any,
-        axes_topology: Any | None,
         box: np.ndarray | None,
         force_partitioning: float,
         customised_axes: bool,
@@ -193,7 +189,6 @@ def _process_united_atom(
             group_id: Molecule-group identifier used for within-frame averaging.
             beads: Mapping of bead keys to reduced-universe atom-index arrays.
             axes_manager: Axes helper used to build translation and rotation axes.
-            axes_topology: Optional cached axes topology generated during static setup.
             box: Optional periodic box vector.
             force_partitioning: Force partitioning factor for highest-level vectors.
             customised_axes: Whether customised UA axes should be used.
@@ -202,7 +197,43 @@ def _process_united_atom(
             out_torque: Frame-local torque second-moment accumulator, mutated in place.
             molcount: Per-residue group sample counters, mutated in place.
         """
+
         for local_res_i, res in enumerate(mol.residues):
+            # local_res_i is relative to first atom in molecule
+
+            if len(mol.residues) > 1:
+                # there are multiple residues in the molecule
+                # build residue group here
+                if local_res_i == 0:
+                    # first residue
+                    relative_id = res.resindex
+                    # index relative to first in mda universe
+                    res_position = -1
+                    residue_group = mol.select_atoms(
+                        f"resindex {local_res_i + relative_id} "
+                        f"or resindex {local_res_i + 1 + relative_id}"
+                    ).residues
+
+                elif local_res_i == len(mol.residues) - 1:
+                    # last residue
+                    res_position = 1
+                    residue_group = mol.select_atoms(
+                        f"resindex {local_res_i - 1 + relative_id} "
+                        f"or resindex {local_res_i + relative_id}"
+                    ).residues
+
+                else:
+                    res_position = 0
+                    residue_group = mol.select_atoms(
+                        f"resindex {local_res_i - 1 + relative_id} "
+                        f"or resindex {local_res_i + relative_id} "
+                        f"or resindex {local_res_i + 1 + relative_id}"
+                    ).residues
+
+            else:
+                # only one residue
+                res_position = None
+                residue_group = res
             bead_key = (mol_id, "united_atom", local_res_i)
             bead_idx_list = beads.get(bead_key, [])
             if not bead_idx_list:
@@ -213,17 +244,14 @@ def _process_united_atom(
                 continue
 
             force_vecs, torque_vecs = self._build_ua_vectors(
-                u=u,
-                mol_id=mol_id,
-                local_res_i=local_res_i,
-                residue_atoms=res.atoms,
+                residue_group=residue_group.atoms,
                 bead_groups=bead_groups,
                 axes_manager=axes_manager,
-                axes_topology=axes_topology,
                 box=box,
                 force_partitioning=force_partitioning,
                 customised_axes=customised_axes,
                 is_highest=is_highest,
+                res_position=res_position,
             )
 
             F, T = self._ft.compute_frame_covariance(force_vecs, torque_vecs)
@@ -243,7 +271,6 @@ def _process_residue(
         group_id: int,
         beads: dict[Any, list[Any]],
         axes_manager: Any,
-        axes_topology: Any | None,
         box: np.ndarray | None,
         customised_axes: bool,
         force_partitioning: float,
@@ -263,7 +290,6 @@ def _process_residue(
             group_id: Molecule-group identifier used for within-frame averaging.
             beads: Mapping of bead keys to reduced-universe atom-index arrays.
             axes_manager: Axes helper used to build translation and rotation axes.
-            axes_topology: Optional cached axes topology generated during static setup.
             box: Optional periodic box vector.
             customised_axes: Whether customised residue axes should be used.
             force_partitioning: Force partitioning factor for highest-level vectors.
@@ -284,12 +310,9 @@ def _process_residue(
             return
 
         force_vecs, torque_vecs = self._build_residue_vectors(
-            u=u,
             mol=mol,
-            mol_id=mol_id,
             bead_groups=bead_groups,
             axes_manager=axes_manager,
-            axes_topology=axes_topology,
             box=box,
             customised_axes=customised_axes,
             force_partitioning=force_partitioning,
@@ -402,32 +425,27 @@ def _process_polymer(
     def _build_ua_vectors(
         self,
         *,
-        u: Any,
-        mol_id: int,
-        local_res_i: int,
         bead_groups: list[Any],
-        residue_atoms: Any,
+        residue_group: Any,
         axes_manager: Any,
-        axes_topology: Any | None,
         box: np.ndarray | None,
         force_partitioning: float,
         customised_axes: bool,
         is_highest: bool,
+        res_position: int,
     ) -> tuple[list[np.ndarray], list[np.ndarray]]:
         """Build force and torque vectors for united-atom beads.
 
         Args:
-            u: Universe-like object used to resolve cached atom indices.
-            mol_id: Molecule index used in axes-topology lookup keys.
-            local_res_i: Local residue index used in axes-topology lookup keys.
-            bead_groups: Atom groups representing UA beads in a residue.
-            residue_atoms: Atom group for the parent residue.
-            axes_manager: Axes helper used to select axes, centres, and moments.
-            axes_topology: Optional cached axes topology generated during static setup.
-            box: Optional periodic box vector.
-            force_partitioning: Force partitioning factor for highest-level vectors.
-            customised_axes: Whether customised UA axes should be used.
-            is_highest: Whether UA is the highest active level.
+            bead_groups: List of UA bead AtomGroups for the residue.
+            residue_group: AtomGroup for the residue group atoms.
+            axes_manager: Axes manager used to determine axes/centers/MOI.
+            box: Optional box vector used for PBC-aware displacements.
+            force_partitioning: Force scaling factor applied at highest level.
+            customised_axes: Whether to use customised axes methods when available.
+            is_highest: Whether UA level is the highest level for the molecule.
+            res_position: Where the residue is in the residue group
+
 
         Returns:
             A tuple containing lists of force vectors and torque vectors.
@@ -437,28 +455,22 @@ def _build_ua_vectors(
 
         for ua_i, bead in enumerate(bead_groups):
             if customised_axes:
-                ua_topology = None
-                if axes_topology is not None:
-                    ua_topology = axes_topology.ua.get((mol_id, local_res_i, ua_i))
-
-                if ua_topology is not None:
-                    trans_axes, rot_axes, center, moi = (
-                        axes_manager.get_UA_axes_from_topology(
-                            u=u,
-                            residue_atoms=residue_atoms,
-                            topology=ua_topology,
-                            box=box,
-                        )
-                    )
-                else:
-                    trans_axes, rot_axes, center, moi = axes_manager.get_UA_axes(
-                        residue_atoms, ua_i
-                    )
+                trans_axes, rot_axes, center, moi = axes_manager.get_UA_axes(
+                    residue_group, ua_i, res_position
+                )
             else:
-                make_whole(residue_atoms)
+                make_whole(residue_group)
                 make_whole(bead)
-
-                trans_axes = residue_atoms.principal_axes()
+                if res_position == -1:
+                    # first residue in group
+                    residue = residue_group.residues[0]
+                elif res_position == 0 or res_position == 1:
+                    # middle or last residue => second in group
+                    residue = residue_group.residues[1]
+                else:
+                    # res_position is None bc there is only one residue
+                    residue = residue_group
+                trans_axes = residue.atoms.principal_axes()
                 rot_axes, moi = axes_manager.get_vanilla_axes(bead)
                 center = bead.center_of_mass(unwrap=True)
 
@@ -487,12 +499,9 @@ def _build_ua_vectors(
     def _build_residue_vectors(
         self,
         *,
-        u: Any,
         mol: Any,
-        mol_id: int,
         bead_groups: list[Any],
         axes_manager: Any,
-        axes_topology: Any | None,
         box: np.ndarray | None,
         customised_axes: bool,
         force_partitioning: float,
@@ -501,12 +510,9 @@ def _build_residue_vectors(
         """Build force and torque vectors for residue beads.
 
         Args:
-            u: Universe-like object used to resolve cached atom indices.
             mol: Molecule fragment containing residues and atoms.
-            mol_id: Molecule index used in axes-topology lookup keys.
             bead_groups: Atom groups representing residue beads.
             axes_manager: Axes helper used to select axes, centres, and moments.
-            axes_topology: Optional cached axes topology generated during static setup.
             box: Optional periodic box vector.
             customised_axes: Whether customised residue axes should be used.
             force_partitioning: Force partitioning factor for highest-level vectors.
@@ -518,16 +524,15 @@ def _build_residue_vectors(
         force_vecs: list[np.ndarray] = []
         torque_vecs: list[np.ndarray] = []
 
+        relative_res_i = mol.residues[0].resindex
+
         for local_res_i, bead in enumerate(bead_groups):
             trans_axes, rot_axes, center, moi = self._get_residue_axes(
-                u=u,
                 mol=mol,
-                mol_id=mol_id,
                 bead=bead,
+                relative_res_i=relative_res_i,
                 local_res_i=local_res_i,
                 axes_manager=axes_manager,
-                axes_topology=axes_topology,
-                box=box,
                 customised_axes=customised_axes,
             )
 
@@ -556,48 +561,30 @@ def _build_residue_vectors(
     def _get_residue_axes(
         self,
         *,
-        u: Any,
         mol: Any,
-        mol_id: int,
         bead: Any,
         local_res_i: int,
+        relative_res_i: int,
         axes_manager: Any,
-        axes_topology: Any | None,
-        box: np.ndarray | None,
         customised_axes: bool,
     ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
         """Return axes, centre, and inertia data for a residue bead.
 
         Args:
-            u: Universe-like object used to resolve cached atom indices.
             mol: Molecule fragment containing residues and atoms.
-            mol_id: Molecule index used in axes-topology lookup keys.
             bead: Atom group representing the residue bead.
             local_res_i: Residue index local to ``mol``.
             axes_manager: Axes helper used to select axes, centres, and moments.
-            axes_topology: Optional cached axes topology generated during static setup.
-            box: Optional periodic box vector.
             customised_axes: Whether customised residue axes should be used.
 
         Returns:
             A tuple of translation axes, rotation axes, centre, and moments of inertia.
         """
         if customised_axes:
             res = mol.residues[local_res_i]
-            residue_topology = None
-            if axes_topology is not None:
-                residue_topology = axes_topology.residue.get((mol_id, local_res_i))
-
-            if residue_topology is not None:
-                return axes_manager.get_residue_axes_from_topology(
-                    u=u,
-                    mol=mol,
-                    residue_atoms=res.atoms,
-                    topology=residue_topology,
-                    box=box,
-                )
-
-            return axes_manager.get_residue_axes(mol, local_res_i, residue=res.atoms)
+            return axes_manager.get_residue_axes(
+                mol, local_res_i, relative_res_i, residue=res.atoms
+            )
 
         make_whole(mol.atoms)
         make_whole(bead)

docs/science.rst

@@ -72,9 +72,9 @@ For the polymer level, the translational and rotational axes are defined as the
 
 For the residue level, there are two situations.
 When the residue is not bonded to any other residues, the translational and rotational axes are the principal axes of the molecule.
-When the residue is part of a larger polymer, the translational axes are the principal axes of the polymer, and the rotational axes are defined from the average position of the bonds to neighbouring residues.
+When the residue is part of a larger polymer, the translational axes are the principal axes of the polymer, and the rotational axes are defined from the two heavy atoms bonded to neighbour residues(E1,E2) and the average position of all other backbone atoms in the residue (C). The backbone of a residue is defined as the shortest path between the two edge atoms of the residue, i.e. the two heavy atoms bonded to neighbour residues.The centre of rotation is located at the point where the perpendicular from C meets the E1-E2 vector.
 
-For the united atom level, the translational axes are defined as the principal axes of the residue and the rotational axes are defined from the average position of the bonds to neighbouring heavy atoms.
+For the united atom level, the translational axes are defined as the residue rotational axes and the rotational axes are defined from the average position of the bonds to neighbouring heavy atoms.
 If there are no bonds to other heavy atoms, the principal axes of the molecule are used.
 
 Conformational Entropy

tests/regression/baselines/benzaldehyde/combined_forcetorque_false.json

@@ -2,15 +2,15 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 0.07119323721997475,
+        "united_atom:Transvibrational": 0.08982962903796131,
         "united_atom:Rovibrational": 49.68669738152346,
         "residue:Transvibrational": 69.48692941204929,
         "residue:Rovibrational": 68.46147102540942,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 20.481571492615355
       },
-      "total": 208.1878625488175
+      "total": 208.2064989406355
     }
   }
 }

tests/regression/baselines/benzaldehyde/frame_window.json

@@ -2,15 +2,15 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 40.30267601961045,
+        "united_atom:Transvibrational": 41.75648599130188,
         "united_atom:Rovibrational": 38.21906858443615,
         "residue:FTmat-Transvibrational": 73.41098578352612,
         "residue:FTmat-Rovibrational": 57.881504393660364,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 20.481571492615355
       },
-      "total": 230.29580627384846
+      "total": 231.74961624553984
     }
   }
 }

tests/regression/baselines/benzaldehyde/selection_subset.json

@@ -2,15 +2,15 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 0.07119323721997475,
+        "united_atom:Transvibrational": 0.08982962903796131,
         "united_atom:Rovibrational": 49.68669738152346,
         "residue:FTmat-Transvibrational": 87.43527331108173,
         "residue:FTmat-Rovibrational": 61.67126452972779,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 20.481571492615355
       },
-      "total": 219.34599995216834
+      "total": 219.3646363439863
     }
   }
 }