CCPBioSim · ioanaapapa · Mar 12, 2026 · Mar 17, 2026 · Mar 25, 2026 · Mar 31, 2026
diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py
diff --git a/CodeEntropy/levels/nodes/covariance.py b/CodeEntropy/levels/nodes/covariance.py
@@ -197,7 +197,30 @@ 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):
+            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:
@@ -208,13 +231,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)
@@ -389,23 +413,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 and torque vectors for united-atom beads.
 
         Args:
-            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.
-            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.
@@ -416,13 +443,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)
 

diff --git a/tests/regression/baselines/benzaldehyde/combined_forcetorque_false.json b/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
     }
   }
 }
diff --git a/tests/regression/baselines/benzaldehyde/default.json b/tests/regression/baselines/benzaldehyde/default.json
@@ -2,15 +2,15 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 24.88518233240474,
+        "united_atom:Transvibrational": 24.869562693152986,
         "united_atom:Rovibrational": 27.950376507672583,
         "residue:FTmat-Transvibrational": 71.03412922724692,
         "residue:FTmat-Rovibrational": 59.44169664956799,
         "united_atom:Conformational": 7.522643899702263,
         "residue:Conformational": 0.0,
         "residue:Orientational": 59.43920971558428
       },
-      "total": 250.27323833217878
+      "total": 250.25761869292705
     }
   }
 }
diff --git a/tests/regression/baselines/benzaldehyde/frame_window.json b/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
     }
   }
 }
diff --git a/tests/regression/baselines/benzaldehyde/rad.json b/tests/regression/baselines/benzaldehyde/rad.json
@@ -2,15 +2,15 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 24.88518233240474,
+        "united_atom:Transvibrational": 24.869562693152986,
         "united_atom:Rovibrational": 27.950376507672583,
         "residue:FTmat-Transvibrational": 71.03412922724692,
         "residue:FTmat-Rovibrational": 59.44169664956799,
         "united_atom:Conformational": 7.522643899702263,
         "residue:Conformational": 0.0,
         "residue:Orientational": 25.502722133228936
       },
-      "total": 216.33675074982344
+      "total": 216.32113111057168
     }
   }
 }
diff --git a/tests/regression/baselines/benzaldehyde/selection_subset.json b/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
     }
   }
 }
diff --git a/tests/regression/baselines/benzene/combined_forcetorque_off.json b/tests/regression/baselines/benzene/combined_forcetorque_off.json
@@ -2,15 +2,15 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 0.16311383522946216,
+        "united_atom:Transvibrational": 0.17824017761231503,
         "united_atom:Rovibrational": 40.695812407899396,
         "residue:Transvibrational": 57.59675197041957,
         "residue:Rovibrational": 52.85496348419672,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 0.0
       },
-      "total": 151.31064169774515
+      "total": 151.325768040128
     }
   }
 }
diff --git a/tests/regression/baselines/benzene/default.json b/tests/regression/baselines/benzene/default.json
@@ -2,15 +2,15 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 9.081853431559052,
+        "united_atom:Transvibrational": 9.050489789631982,
         "united_atom:Rovibrational": 27.534380126250227,
         "residue:FTmat-Transvibrational": 72.66211800013413,
         "residue:FTmat-Rovibrational": 59.93761874375924,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 41.58427729275938
       },
-      "total": 210.800247594462
+      "total": 210.76888395253496
     }
   }
 }
diff --git a/tests/regression/baselines/benzene/frame_window.json b/tests/regression/baselines/benzene/frame_window.json
@@ -2,15 +2,15 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 11.46552036084548,
+        "united_atom:Transvibrational": 11.552833417413368,
         "united_atom:Rovibrational": 34.98492056748493,
         "residue:FTmat-Transvibrational": 71.83491878606151,
         "residue:FTmat-Rovibrational": 55.6098400281744,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 0.0
       },
-      "total": 173.89519974256632
+      "total": 173.9825127991342
     }
   }
 }
diff --git a/tests/regression/baselines/benzene/rad.json b/tests/regression/baselines/benzene/rad.json
@@ -2,15 +2,15 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 9.081853431559052,
+        "united_atom:Transvibrational": 9.050489789631982,
         "united_atom:Rovibrational": 27.534380126250227,
         "residue:FTmat-Transvibrational": 72.66211800013413,
         "residue:FTmat-Rovibrational": 59.93761874375924,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 12.386543820026672
       },
-      "total": 181.6025141217293
+      "total": 181.57115047980224
     }
   }
 }
diff --git a/tests/regression/baselines/benzene/selection_subset.json b/tests/regression/baselines/benzene/selection_subset.json
@@ -2,15 +2,15 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 0.16311383522946216,
+        "united_atom:Transvibrational": 0.17824017761231503,
         "united_atom:Rovibrational": 40.695812407899396,
         "residue:FTmat-Transvibrational": 71.15771063929333,
         "residue:FTmat-Rovibrational": 47.26253953880574,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 0.0
       },
-      "total": 159.27917642122793
+      "total": 159.2943027636108
     }
   }
 }
diff --git a/tests/regression/baselines/cyclohexane/combined_forcetorque_off.json b/tests/regression/baselines/cyclohexane/combined_forcetorque_off.json
@@ -2,15 +2,15 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 0.5949308536734588,
+        "united_atom:Transvibrational": 0.6825851535582973,
         "united_atom:Rovibrational": 24.234154676578637,
         "residue:Transvibrational": 69.9872567772153,
         "residue:Rovibrational": 68.68521019685565,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 0.0
       },
-      "total": 163.50155250432303
+      "total": 163.5892068042079
     }
   }
 }
diff --git a/tests/regression/baselines/cyclohexane/default.json b/tests/regression/baselines/cyclohexane/default.json
@@ -2,15 +2,15 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 11.726471858700231,
+        "united_atom:Transvibrational": 11.65610486557721,
         "united_atom:Rovibrational": 47.71088602161552,
         "residue:FTmat-Transvibrational": 70.3922327806972,
         "residue:FTmat-Rovibrational": 63.44920824648768,
         "united_atom:Conformational": 1.288879206695518,
         "residue:Conformational": 0.0,
         "residue:Orientational": 47.54362774624537
       },
-      "total": 242.11130586044152
+      "total": 242.0409388673185
     }
   }
 }
diff --git a/tests/regression/baselines/cyclohexane/frame_window.json b/tests/regression/baselines/cyclohexane/frame_window.json
@@ -2,15 +2,15 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 17.398382397359153,
+        "united_atom:Transvibrational": 16.605273867798914,
         "united_atom:Rovibrational": 73.96792995794405,
         "residue:FTmat-Transvibrational": 76.4267996157354,
         "residue:FTmat-Rovibrational": 63.30469126284744,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 0.0
       },
-      "total": 231.09780323388605
+      "total": 230.3046947043258
     }
   }
 }
diff --git a/tests/regression/baselines/cyclohexane/rad.json b/tests/regression/baselines/cyclohexane/rad.json
@@ -2,15 +2,15 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 11.726471858700231,
+        "united_atom:Transvibrational": 11.65610486557721,
         "united_atom:Rovibrational": 47.71088602161552,
         "residue:FTmat-Transvibrational": 70.3922327806972,
         "residue:FTmat-Rovibrational": 63.44920824648768,
         "united_atom:Conformational": 1.288879206695518,
         "residue:Conformational": 0.0,
         "residue:Orientational": 14.049324060636614
       },
-      "total": 208.61700217483278
+      "total": 208.54663518170975
     }
   }
 }
diff --git a/tests/regression/baselines/cyclohexane/selection_subset.json b/tests/regression/baselines/cyclohexane/selection_subset.json
@@ -2,15 +2,15 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 0.5949308536734588,
+        "united_atom:Transvibrational": 0.6825851535582973,
         "united_atom:Rovibrational": 24.234154676578637,
         "residue:FTmat-Transvibrational": 84.37184730911717,
         "residue:FTmat-Rovibrational": 59.52377096811085,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 0.0
       },
-      "total": 168.72470380748013
+      "total": 168.81235810736496
     }
   }
 }
diff --git a/tests/regression/baselines/dna/combined_forcetorque_off.json b/tests/regression/baselines/dna/combined_forcetorque_off.json
@@ -2,17 +2,17 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 0.0,
+        "united_atom:Transvibrational": 4.762281610623415e-20,
         "united_atom:Rovibrational": 0.002160679012128457,
         "residue:Transvibrational": 0.0,
-        "residue:Rovibrational": 3.376800684085249,
+        "residue:Rovibrational": 6.633599673254765,
         "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.41490388643341
     },
     "1": {
       "components": {