Residue Bonded Axes

CodeEntropy/levels/nodes/covariance.py

@@ -200,7 +200,30 @@ def _process_united_atom(
         Returns:
             None. Mutates out_force/out_torque and molcount in-place.
         """
+
         for local_res_i, res in enumerate(mol.residues):
+            if len(mol.residues) > 1:
+                # there are multiple residues in the molecule
+                # build residue group here
+                if local_res_i == 0:
+                    # first residue
+                    res_position = -1
+                    res_next = mol.residues[1]
+                    residue_group = res + res_next
+                elif local_res_i == len(mol.residues) - 1:
+                    # last residue
+                    res_position = 1
+                    res_prev = mol.residues[-2]
+                    residue_group = res + res_prev
+                else:
+                    res_position = 0
+                    res_prev = mol.residues[local_res_i - 1]
+                    res_next = mol.residues[local_res_i + 1]
+                    residue_group = res_prev + res + res_next
+            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:
@@ -211,13 +234,14 @@ def _process_united_atom(
                 continue
 
             force_vecs, torque_vecs = self._build_ua_vectors(
-                residue_atoms=res.atoms,
+                residue_group=residue_group.atoms,
                 bead_groups=bead_groups,
                 axes_manager=axes_manager,
                 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)
@@ -413,23 +437,26 @@ def _build_ua_vectors(
         self,
         *,
         bead_groups: list[Any],
-        residue_atoms: Any,
+        residue_group: Any,
         axes_manager: Any,
         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/torque vectors for UA-level beads of one residue.
 
         Args:
             bead_groups: List of UA bead AtomGroups for the residue.
-            residue_atoms: AtomGroup for the residue atoms (used for axes when vanilla).
+            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 (force_vecs, torque_vecs), each a list of (3,) vectors ordered
@@ -441,13 +468,21 @@ def _build_ua_vectors(
         for ua_i, bead in enumerate(bead_groups):
             if customised_axes:
                 trans_axes, rot_axes, center, moi = axes_manager.get_UA_axes(
-                    residue_atoms, ua_i
+                    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)
 

tests/regression/baselines/dna/combined_forcetorque_off.json

@@ -5,14 +5,14 @@
         "united_atom:Transvibrational": 0.0,
         "united_atom:Rovibrational": 0.002160679012128457,
         "residue:Transvibrational": 0.0,
-        "residue:Rovibrational": 3.376800684085249,
+        "residue:Rovibrational": 6.832779629985204,
         "polymer:Transvibrational": 21.18266215491188,
         "polymer:Rovibrational": 12.837576042626923,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "polymer:Orientational": 4.758905336627712
       },
-      "total": 42.1581048972639
+      "total": 45.61408384316385
     },
     "1": {
       "components": {

tests/regression/baselines/dna/default.json

@@ -5,14 +5,14 @@
         "united_atom:Transvibrational": 0.0,
         "united_atom:Rovibrational": 0.002160679012128457,
         "residue:Transvibrational": 0.0,
-        "residue:Rovibrational": 3.376800684085249,
+        "residue:Rovibrational": 6.832779629985204,
         "polymer:FTmat-Transvibrational": 12.341104347192612,
         "polymer:FTmat-Rovibrational": 0.0,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "polymer:Orientational": 4.758905336627712
       },
-      "total": 20.478971046917703
+      "total": 23.934949992817657
     },
     "1": {
       "components": {

tests/regression/baselines/dna/frame_window.json

@@ -2,17 +2,17 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 0.0,
+        "united_atom:Transvibrational": 2.3245522868433963e-06,
         "united_atom:Rovibrational": 1.5821720528374943,
         "residue:Transvibrational": 0.0,
-        "residue:Rovibrational": 27.397449238560412,
+        "residue:Rovibrational": 36.126942876174176,
         "polymer:FTmat-Transvibrational": 48.62026970762269,
         "polymer:FTmat-Rovibrational": 0.0,
         "united_atom:Conformational": 10.584542990557836,
         "residue:Conformational": 0.0,
         "polymer:Orientational": 4.758905336627712
       },
-      "total": 92.94333932620614
+      "total": 101.6728352883722
     },
     "1": {
       "components": {

tests/regression/baselines/dna/grouping_each.json

@@ -5,14 +5,14 @@
         "united_atom:Transvibrational": 0.0,
         "united_atom:Rovibrational": 0.002160679012128457,
         "residue:Transvibrational": 0.0,
-        "residue:Rovibrational": 3.376800684085249,
+        "residue:Rovibrational": 6.832779629985204,
         "polymer:FTmat-Transvibrational": 12.341104347192612,
         "polymer:FTmat-Rovibrational": 0.0,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "polymer:Orientational": 4.758905336627712
       },
-      "total": 20.478971046917703
+      "total": 23.934949992817657
     },
     "1": {
       "components": {

tests/regression/baselines/dna/selection_subset.json

@@ -5,14 +5,14 @@
         "united_atom:Transvibrational": 0.0,
         "united_atom:Rovibrational": 0.002160679012128457,
         "residue:Transvibrational": 0.0,
-        "residue:Rovibrational": 3.376800684085249,
+        "residue:Rovibrational": 6.832779629985204,
         "polymer:FTmat-Transvibrational": 12.341104347192612,
         "polymer:FTmat-Rovibrational": 0.0,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "polymer:Orientational": 4.758905336627712
       },
-      "total": 20.478971046917703
+      "total": 23.934949992817657
     },
     "1": {
       "components": {

tests/unit/CodeEntropy/levels/nodes/test_frame_covariance_node.py

@@ -401,7 +401,7 @@ def test_build_ua_vectors_customised_axes_true_calls_get_UA_axes():
     node = FrameCovarianceNode()
 
     bead = _BeadGroup(1)
-    residue_atoms = MagicMock()
+    residue_group = MagicMock()
 
     axes_manager = MagicMock()
     axes_manager.get_UA_axes.return_value = (
@@ -416,12 +416,13 @@ def test_build_ua_vectors_customised_axes_true_calls_get_UA_axes():
 
     force_vecs, torque_vecs = node._build_ua_vectors(
         bead_groups=[bead],
-        residue_atoms=residue_atoms,
+        residue_group=residue_group,
         axes_manager=axes_manager,
         box=np.array([10.0, 10.0, 10.0]),
         force_partitioning=1.0,
         customised_axes=True,
         is_highest=True,
+        res_position=None,
     )
 
     axes_manager.get_UA_axes.assert_called_once()
@@ -451,12 +452,13 @@ def test_build_ua_vectors_vanilla_path_uses_principal_axes_and_vanilla_axes(
 
     force_vecs, torque_vecs = node._build_ua_vectors(
         bead_groups=[bead],
-        residue_atoms=residue_atoms,
+        residue_group=residue_atoms,
         axes_manager=axes_manager,
         box=np.array([10.0, 10.0, 10.0]),
         force_partitioning=1.0,
         customised_axes=False,
         is_highest=True,
+        res_position=None,
     )
 
     axes_manager.get_vanilla_axes.assert_called_once()