From 4a1049302d26d21a817a46554731ed30c6a43e2b Mon Sep 17 00:00:00 2001
From: ioanaapapa <ioanaalexandrapapa@gmail.com>
Date: Fri, 24 Jul 2026 12:18:51 +0100
Subject: [PATCH 1/5] modified residue rotation + UA translation axes and
 corresponding unit tests

---
 CodeEntropy/levels/axes.py                    | 609 +++++++++-------
 CodeEntropy/levels/nodes/covariance.py        | 153 ++--
 .../levels/nodes/test_covariance_node.py      | 121 +---
 tests/unit/CodeEntropy/levels/test_axes.py    | 661 +++++-------------
 4 files changed, 610 insertions(+), 934 deletions(-)

diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py
index 4d4cf89..8bec335 100644
--- a/CodeEntropy/levels/axes.py
+++ b/CodeEntropy/levels/axes.py
@@ -61,13 +61,15 @@ def __init__(self) -> None:
         self._rot_axes = None
         self._number_of_beads = None
 
-    def get_residue_axes(self, data_container, index: int, residue=None):
+    def get_residue_axes(
+        self, data_container, index: int, relative_index: int, residue=None
+    ):
         """Compute residue-level translational and rotational axes.
 
         The translational and rotational axes at the residue level.
 
         - Identify the residue (either provided or selected by `resindex index`).
-        - Determine whether the residue is bonded to neighboring residues
+        - Determine whether the residue is bonded to neighbouring residues
           (previous/next in sequence) using MDAnalysis bonded selections.
         - If there are *no* bonds to other residues:
             * Use a custom principal axes, from a moment-of-inertia (MOI) tensor
@@ -76,13 +78,19 @@ def get_residue_axes(self, data_container, index: int, residue=None):
             * Set translational axes equal to rotational axes (as per the original
               code convention).
         - If bonded to other residues:
-            * Use default axes and MOI (MDAnalysis principal axes / inertia).
+            Find edge heavy atoms (i.e. heavy atoms bonded to neighbour residues)
+            and find the shortest chain between them: the backbone. Edge
+            atoms + backbone COM are used to determine UA translational axes
+            (see get_residue_custom_axes)
 
         Args:
             data_container (MDAnalysis.Universe or AtomGroup):
                 Molecule and trajectory data (the fragment/molecule container).
             index (int):
                 Residue index (resindex) within `data_container`.
+            relative_index (int):
+                Index of first residue within 'data_container'.
+                This is used to obtain index in MDA Universe for atom selections.
             residue (MDAnalysis.AtomGroup, optional):
                 If provided, this residue selection will be used rather than
                 selecting again.
@@ -99,101 +107,95 @@ def get_residue_axes(self, data_container, index: int, residue=None):
                 If the residue selection is empty.
         """
         # TODO refine selection so that it will work for branched polymers
-        index_prev = index - 1
-        index_next = index + 1
+        # match indexing to MDAnalysis indexing
+
+        index_prev = index + relative_index - 1
+        index_next = index + relative_index + 1
 
         if residue is None:
-            residue = data_container.select_atoms(f"resindex {index}")
+            residue = data_container.select_atoms(f"resindex {index + relative_index}")
+            # residue of interest
         if len(residue) == 0:
-            raise ValueError(f"Empty residue selection for resindex={index}")
+            raise ValueError(
+                f"Empty residue selection for resindex={index + relative_index}"
+            )
 
-        center = residue.atoms.center_of_mass(unwrap=True)
-        atom_set = data_container.select_atoms(
-            f"(resindex {index_prev} or resindex {index_next}) and bonded resid {index}"
+        edge_atom_set = data_container.atoms.select_atoms(
+            f"resindex {index + relative_index} and "
+            f"(bonded resindex {index_prev} or "
+            f"resindex {index_next})"
         )
 
-        if len(atom_set) == 0:
-            # No bonds to other residues.
+        uas = residue.select_atoms("mass 2 to 999")
+        ua_masses = self.get_UA_masses(residue)
+
+        if len(edge_atom_set) == 0:
+            # No UAS are bonded to other residues
             # Use a custom principal axes, from a MOI tensor that uses positions of
             # heavy atoms only, but including masses of heavy atom + bonded H.
-            uas = residue.select_atoms("mass 2 to 999")
-            ua_masses = self.get_UA_masses(residue)
+
             moi_tensor = self.get_moment_of_inertia_tensor(
-                center_of_mass=center,
+                center_of_mass=np.array(residue.center_of_mass()),
                 positions=uas.positions,
                 masses=ua_masses,
                 dimensions=data_container.dimensions[:3],
             )
             rot_axes, moment_of_inertia = self.get_custom_principal_axes(moi_tensor)
             trans_axes = rot_axes  # per original convention
+            rot_center = np.array(residue.center_of_mass())
         else:
-            # If bonded to other residues, use default axes and MOI.
+            # If bonded to other residues, use local axes.
             make_whole(data_container.atoms)
             trans_axes = data_container.atoms.principal_axes()
-            rot_axes, moment_of_inertia = self.get_vanilla_axes(residue)
-            center = residue.center_of_mass(unwrap=True)
-
-        return trans_axes, rot_axes, center, moment_of_inertia
-
-    def get_residue_axes_from_topology(
-        self,
-        *,
-        u,
-        mol,
-        residue_atoms,
-        topology,
-        box: np.ndarray | None,
-    ):
-        """Compute residue axes using cached static topology.
-
-        This is the cached-index equivalent of ``get_residue_axes``. It keeps
-        all frame-dependent numerical work frame-local, but avoids repeated
-        MDAnalysis selections for residue heavy atoms, UA masses, and neighbour
-        bond discovery.
 
-        Args:
-            u: Current-frame universe used to resolve cached atom indices.
-            mol: Current-frame molecule fragment.
-            residue_atoms: AtomGroup for the residue in the current frame.
-            topology: Cached ``ResidueAxesTopology`` for this residue.
-            box: Current periodic box lengths. If omitted, ``u.dimensions`` is used.
-
-        Returns:
-            Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
-                - trans_axes: Translational axes, shape ``(3, 3)``.
-                - rot_axes: Rotational axes, shape ``(3, 3)``.
-                - center: Residue centre, shape ``(3,)``.
-                - moment_of_inertia: Principal moments, shape ``(3,)``.
-        """
-        dimensions = (
-            np.asarray(box, dtype=float)
-            if box is not None
-            else np.asarray(u.dimensions[:3], dtype=float)
-        )
-
-        center = residue_atoms.center_of_mass(unwrap=True)
-
-        if not topology.has_neighbor_bonds:
-            heavy_atoms = u.atoms[topology.residue_heavy_indices]
-            moment_of_inertia_tensor = self.get_moment_of_inertia_tensor(
-                center_of_mass=center,
-                positions=heavy_atoms.positions,
-                masses=topology.residue_ua_masses,
-                dimensions=dimensions,
-            )
-            rot_axes, moment_of_inertia = self.get_custom_principal_axes(
-                moment_of_inertia_tensor
+            if len(edge_atom_set) == 1:
+                if index == 0:
+                    # first residue
+                    # use first heavy atom
+                    edges = [residue.atoms[0], edge_atom_set[0]]
+                    backbone = self.get_chain(
+                        residue, residue.atoms[0], edge_atom_set[0]
+                    )
+                else:
+                    # last residue
+                    last_index = len(uas) - 1
+                    last = None
+                    # look for last heavy atom
+                    # with only one bond to another
+                    while last_index > 0 and last is None:
+                        heavy_atom = uas[last_index]
+                        bonded_atoms = residue.atoms.select_atoms(
+                            f"(mass 2 to 999) and bonded index {heavy_atom.index}"
+                        )
+                        if len(bonded_atoms) == 1:
+                            last = heavy_atom
+                        else:
+                            last_index -= 1
+                    edges = [edge_atom_set[0], last]
+
+                    backbone = self.get_chain(residue, edge_atom_set[0], last)
+            else:
+                # residue has two bonds to other residues
+                edges = [edge_atom_set[0], edge_atom_set[1]]
+                backbone = self.get_chain(residue, edge_atom_set[0], edge_atom_set[1])
+            # get edge atoms of the residue
+            # for terminal residues, this will include the C/N terminus
+            backbone_center = np.zeros(3)
+            for heavy_atom in backbone:
+                backbone_center += heavy_atom.position
+            backbone_center = backbone_center / len(backbone)
+            rot_center, rot_axes = self.get_residue_custom_axes(edges, backbone_center)
+
+            moment_of_inertia = self.get_custom_residue_moment_of_inertia(
+                center_of_mass=rot_center,
+                positions=uas.positions,
+                masses=ua_masses,
+                custom_rot_axes=rot_axes,
+                dimensions=data_container.dimensions[:3],
             )
-            trans_axes = rot_axes
-        else:
-            make_whole(mol.atoms)
-            trans_axes = mol.atoms.principal_axes()
-            rot_axes, moment_of_inertia = self.get_vanilla_axes(residue_atoms)
-            center = residue_atoms.center_of_mass(unwrap=True)
-
-        return trans_axes, rot_axes, center, moment_of_inertia
+        return trans_axes, rot_axes, rot_center, moment_of_inertia
 
-    def get_UA_axes(self, data_container, index: int):
+    def get_UA_axes(self, data_container, index: int, res_position):
         """Compute united-atom-level translational and rotational axes.
 
         The translational and rotational axes at the united-atom level.
@@ -201,7 +203,12 @@ def get_UA_axes(self, data_container, index: int):
         This preserves the original behaviour and its rationale:
 
         - Translational axes:
-            Use the same custom principal-axes approach as residue level:
+            Use the same approach as residue level rotational.
+            Identify residue of interest and neighbours, then select
+            edge heavy atoms (i.e. heavy atoms bonded to neighbour residues)
+            and find the shortest chain between them: the backbone. Edge
+            atoms + backbone COM are used to determine UA translational axes
+            (see get_residue_custom_axes)
             compute a custom MOI tensor using heavy-atom coordinates but UA masses
             (heavy + bonded H masses), then compute the principal axes from it.
 
@@ -216,7 +223,8 @@ def get_UA_axes(self, data_container, index: int):
                 Molecule and trajectory data.
             index (int):
                 Bead index (ordinal among heavy atoms).
-
+            res_position: where the residue of interest is
+                in data_container
         Returns:
             Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
                 - trans_axes: Translational axes (3, 3).
@@ -228,209 +236,180 @@ def get_UA_axes(self, data_container, index: int):
             IndexError:
                 If `index` does not correspond to an existing heavy atom.
             ValueError:
-                If bonded-axis construction fails.
+                If axis construction fails.
         """
 
-        index = int(index)  # bead index
-
+        index = int(index)  # UA bead index
+        heavy_atoms = data_container.select_atoms("mass 2 to 999")
         # use the same customPI trans axes as the residue level
-        heavy_atoms = data_container.select_atoms("prop mass > 1.1")
         if len(heavy_atoms) > 1:
-            UA_masses = self.get_UA_masses(data_container.atoms)
-            center = data_container.atoms.center_of_mass(unwrap=True)
-            moment_of_inertia_tensor = self.get_moment_of_inertia_tensor(
-                center, heavy_atoms.positions, UA_masses, data_container.dimensions[:3]
-            )
-            trans_axes, _moment_of_inertia = self.get_custom_principal_axes(
-                moment_of_inertia_tensor
+            if len(data_container.residues) == 1:
+                # only the one residue => use principal axes
+                residue = data_container
+                trans_center = data_container.atoms.center_of_mass(unwrap=True)
+                trans_axes = data_container.atoms.principal_axes()
+                residue_heavy_atoms = heavy_atoms
+            else:
+                # residue of interest has at least one neighbour
+                if res_position == -1:
+                    residue = data_container.residues[0]
+                    resindex = residue.resindex
+                    resindex_next = resindex + 1
+
+                    second_edge = data_container.select_atoms(
+                        f"resindex {resindex} and bonded resindex {resindex_next}"
+                    )
+
+                    edges = [residue.atoms[0], second_edge[0]]
+                    backbone = self.get_chain(
+                        residue, residue.atoms[0], second_edge.atoms[0]
+                    )
+
+                elif res_position == 0:
+                    # between 2 residues
+                    residue = data_container.residues[1]
+                    resindex = residue.resindex
+                    resindex_next = resindex + 1
+                    resindex_prev = resindex - 1
+
+                    edge_set = data_container.select_atoms(
+                        f"resindex {resindex} and "
+                        f"(bonded resindex {resindex_prev} or "
+                        f"resindex {resindex_next})"
+                    )
+
+                    edges = [edge_set[0], edge_set[1]]
+                    backbone = self.get_chain(residue, edge_set[0], edge_set[1])
+
+                else:
+                    # last resid
+                    # always resindex 1 in data_container
+                    residue = data_container.residues[1]
+                    resindex = residue.resindex
+                    resindex_prev = resindex - 1
+                    first_edge = data_container.select_atoms(
+                        f"resindex {resindex} and bonded resindex {resindex_prev}"
+                    )
+
+                    last_index = len(heavy_atoms) - 1
+                    last = None
+                    # look for last heavy atom
+                    # with only one bond to another
+                    while last_index > 0 and last is None:
+                        heavy_atom = heavy_atoms[last_index]
+                        bonded_atoms = residue.atoms.select_atoms(
+                            f"(mass 2 to 999) and bonded index {heavy_atom.index}"
+                        )
+                        if len(bonded_atoms) == 1:
+                            last = heavy_atom
+                        else:
+                            last_index -= 1
+
+                    edges = [first_edge.atoms[0], last]
+                    backbone = self.get_chain(residue, first_edge.atoms[0], last)
+
+                backbone_center = np.zeros(3)
+                for heavy_atom in backbone:
+                    backbone_center += heavy_atom.position
+                backbone_center = backbone_center / len(backbone)
+
+                trans_center, trans_axes = self.get_residue_custom_axes(
+                    edges, backbone_center
+                )
+                residue_heavy_atoms = residue.atoms.select_atoms("mass 2 to 999")
+
+            # look for heavy atoms in residue of interest
+            heavy_atom_indices = []
+            for atom in residue_heavy_atoms:
+                heavy_atom_indices.append(atom.index)
+            # we find the nth heavy atom
+            # where n is the bead index
+            heavy_atom_index = heavy_atom_indices[index]
+            heavy_atom = residue.atoms.select_atoms(f"index {heavy_atom_index}")[0]
+            rot_center = heavy_atom.position
+            rot_axes, moment_of_inertia = self.get_bonded_axes(
+                system=data_container,
+                atom=heavy_atom,
+                dimensions=data_container.dimensions[:3],
             )
-        else:
-            # use standard PA for UA not bonded to anything else
-            make_whole(data_container.atoms)
-            trans_axes = data_container.atoms.principal_axes()
-
-        # look for heavy atoms in residue of interest
-        heavy_atom_indices = []
-        for atom in heavy_atoms:
-            heavy_atom_indices.append(atom.index)
-        # we find the nth heavy atom
-        # where n is the bead index
-        heavy_atom_index = heavy_atom_indices[index]
-        heavy_atom = data_container.select_atoms(f"index {heavy_atom_index}")
-
-        center = heavy_atom.positions[0]
-        rot_axes, moment_of_inertia = self.get_bonded_axes(
-            system=data_container,
-            atom=heavy_atom[0],
-            dimensions=data_container.dimensions[:3],
-        )
-        if rot_axes is None or moment_of_inertia is None:
-            raise ValueError("Unable to compute bonded axes for UA bead.")
 
-        logger.debug(f"Translational Axes: {trans_axes}")
-        logger.debug(f"Rotational Axes: {rot_axes}")
-        logger.debug(f"Center: {center}")
-        logger.debug(f"Moment of Inertia: {moment_of_inertia}")
-
-        return trans_axes, rot_axes, center, moment_of_inertia
-
-    def get_UA_axes_from_topology(
-        self,
-        *,
-        u,
-        residue_atoms,
-        topology,
-        box: np.ndarray | None,
-    ):
-        """Compute UA axes using cached static topology.
-
-        This is the cached-index equivalent of ``get_UA_axes``. It preserves the
-        frame-dependent numerical calculations, but avoids repeated MDAnalysis
-        selection strings for heavy atoms, bonded atoms, and UA masses.
-
-        Args:
-            u: Current-frame universe.
-            residue_atoms: AtomGroup for the parent residue in the current frame.
-            topology: Cached ``UAAxesTopology`` for this UA bead.
-            box: Current periodic box lengths. If omitted, ``u.dimensions`` is used.
-
-        Returns:
-            Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
-                - trans_axes: Translational axes, shape ``(3, 3)``.
-                - rot_axes: Rotational axes, shape ``(3, 3)``.
-                - center: Rotation centre, shape ``(3,)``.
-                - moment_of_inertia: Principal moments, shape ``(3,)``.
-
-        Raises:
-            ValueError: If cached bonded-axis construction fails.
-        """
-        dimensions = (
-            np.asarray(box, dtype=float)
-            if box is not None
-            else np.asarray(u.dimensions[:3], dtype=float)
-        )
+        else:
+            # 1 heavy atom in the data_container
+            heavy_atom = heavy_atoms[0]
+            # trans and rot centres are centre of mass
+            rot_center = data_container.center_of_mass()
+            rot_axes, moment_of_inertia = self.get_bonded_axes(
+                system=data_container,
+                atom=heavy_atom,
+                dimensions=data_container.dimensions[:3],
+            )
+            trans_center = rot_center
+            # principal axes
+            trans_axes = rot_axes
 
-        heavy_atoms = u.atoms[topology.residue_heavy_indices]
-        heavy_atom = u.atoms[int(topology.heavy_atom_index)]
+        if trans_axes is None:
+            raise ValueError("Unable to compute translation axes for UA bead.")
 
-        if len(heavy_atoms) > 1:
-            center = residue_atoms.center_of_mass(unwrap=True)
-            moment_of_inertia_tensor = self.get_moment_of_inertia_tensor(
-                center_of_mass=center,
-                positions=heavy_atoms.positions,
-                masses=topology.residue_ua_masses,
-                dimensions=dimensions,
-            )
-            trans_axes, _moment_of_inertia = self.get_custom_principal_axes(
-                moment_of_inertia_tensor
-            )
-        else:
-            make_whole(residue_atoms)
-            trans_axes = residue_atoms.principal_axes()
-
-        center = heavy_atom.position
-        rot_axes, moment_of_inertia = self.get_bonded_axes_from_topology(
-            u=u,
-            heavy_atom=heavy_atom,
-            topology=topology,
-            dimensions=dimensions,
-        )
         if rot_axes is None or moment_of_inertia is None:
-            raise ValueError("Unable to compute bonded axes for cached UA bead.")
+            raise ValueError("Unable to compute bonded axes for UA bead.")
 
         logger.debug("Translational Axes: %s", trans_axes)
         logger.debug("Rotational Axes: %s", rot_axes)
-        logger.debug("Center: %s", center)
+        logger.debug("Translational center: %s", trans_center)
+        logger.debug("Rotational center: %s", rot_center)
         logger.debug("Moment of Inertia: %s", moment_of_inertia)
 
-        return trans_axes, rot_axes, center, moment_of_inertia
+        return trans_axes, rot_axes, rot_center, moment_of_inertia
 
-    def get_bonded_axes_from_topology(
-        self,
-        *,
-        u,
-        heavy_atom,
-        topology,
-        dimensions: np.ndarray,
-    ):
-        """Compute UA bonded axes using cached bonded atom indices.
+    def get_residue_custom_axes(self, edges, center):
+        """
+        Compute rotation axes at the residue level, given
+        two edge atoms of the residue (E1+E2),
+        and the centre of geometry of backbone atoms
+        that are not edges (C).
+        x axis is O-E1
+        y axis is O-C (perpendicular to O-E1 in the
+        same plane as E2)
+        z axis is perpendicular to the two other axes
+
+        ::
 
-        This mirrors ``get_bonded_axes`` but receives precomputed bonded atom
-        memberships from ``UAAxesTopology`` instead of rediscovering them with
-        MDAnalysis selection strings inside the frame loop.
+                    C
+                    |
+                    |
+            E1 ---- O --- E2
 
         Args:
-            u: Current-frame universe.
-            heavy_atom: Current-frame heavy atom for the UA bead.
-            topology: Cached ``UAAxesTopology`` for the UA bead.
-            dimensions: Simulation box lengths, shape ``(3,)``.
+            edges: (2,3) positions of two edge atoms
+            center: (3,) coordinates of the inner backbone
+            centre of geometry
 
         Returns:
-            Tuple[np.ndarray | None, np.ndarray | None]:
-                - custom_axes: Custom rotation axes, shape ``(3, 3)``, or ``None``.
-                - custom_moment_of_inertia: Principal moments, shape ``(3,)``, or
-                  ``None``.
+            rot_center: (3,) rotation centre,
+            lies on the E1-E2 vector
+            rot_axes: (3,3) rotation axes of residue
         """
-        if not heavy_atom.mass > 1.1:
-            return None, None
-
-        custom_moment_of_inertia = None
-        custom_axes = None
-
-        heavy_bonded = u.atoms[topology.bonded_heavy_indices]
-        light_bonded = u.atoms[topology.bonded_light_indices]
-        ua = u.atoms[topology.ua_atom_indices]
-        ua_all = u.atoms[topology.ua_all_atom_indices]
-
-        if len(heavy_bonded) == 0:
-            custom_axes, custom_moment_of_inertia = self.get_vanilla_axes(ua_all)
-
-        if len(heavy_bonded) == 1 and len(light_bonded) == 0:
-            custom_axes = self.get_custom_axes(
-                a=heavy_atom.position,
-                b_list=[heavy_bonded[0].position],
-                c=np.zeros(3),
-                dimensions=dimensions,
-            )
-
-        if len(heavy_bonded) == 1 and len(light_bonded) >= 1:
-            custom_axes = self.get_custom_axes(
-                a=heavy_atom.position,
-                b_list=[heavy_bonded[0].position],
-                c=light_bonded[0].position,
-                dimensions=dimensions,
-            )
-
-        if len(heavy_bonded) >= 2:
-            custom_axes = self.get_custom_axes(
-                a=heavy_atom.position,
-                b_list=heavy_bonded.positions,
-                c=heavy_bonded[1].position,
-                dimensions=dimensions,
-            )
-
-        if custom_axes is None:
-            return None, None
-
-        if custom_moment_of_inertia is None:
-            custom_moment_of_inertia = self.get_custom_moment_of_inertia(
-                UA=ua,
-                custom_rotation_axes=custom_axes,
-                center_of_mass=heavy_atom.position,
-                dimensions=dimensions,
-            )
-
-        custom_axes = self.get_flipped_axes(
-            ua,
-            custom_axes,
-            heavy_atom.position,
-            dimensions,
-        )
-
-        return custom_axes, custom_moment_of_inertia
+        # x axis is O-E1
+        E1C_vector = center - edges[0].position
+        # look for projection of E1-O onto E1-E2 (E1-C)
+        E1E2_vector = edges[1].position - edges[0].position
+        E1O_vector = (
+            np.dot(E1E2_vector, E1C_vector) / (np.linalg.norm(E1E2_vector) ** 2)
+        ) * E1E2_vector
+        x_axis = -E1O_vector
+        # O-C = O-E1 + E1-C
+        OC_vector = -E1O_vector + E1C_vector
+        y_axis = OC_vector
+        z_axis = np.cross(x_axis, y_axis)
+        x_axis /= np.linalg.norm(x_axis)
+        y_axis /= np.linalg.norm(y_axis)
+        z_axis /= np.linalg.norm(z_axis)
+        rot_axes = np.array([x_axis, y_axis, z_axis])
+        rot_center = E1O_vector + edges[0].position
+        return rot_center, rot_axes
 
     def get_bonded_axes(self, system, atom, dimensions: np.ndarray):
-        r"""Compute UA rotational axes from bonded topology around a heavy atom.
+        """Compute UA rotational axes from bonded topology around a heavy atom.
 
         For a given heavy atom, use its bonded atoms to get the axes for rotating
         forces around. Few cases for choosing united atom axes, which are dependent
@@ -659,6 +638,41 @@ def get_custom_axes(
         scaled_custom_axes = unscaled_custom_axes / mod[:, np.newaxis]
         return scaled_custom_axes
 
+    def get_custom_residue_moment_of_inertia(
+        self,
+        center_of_mass: np.ndarray,
+        positions: np.ndarray,
+        masses: np.ndarray,
+        custom_rot_axes: np.ndarray,
+        dimensions: np.ndarray,
+    ):
+        """
+        Compute moment of inertia around custom axes for a bead
+        formed of multiple UAs.
+
+        Args:
+            center_of_mass: (3, ) COM for bead
+            positions: (N,3) positions of the UAs in the bead
+            masses: (N,) masses of the UAs in the bead
+            custom_rot_axes: (3,3) array of residue rotation axes
+            dimensions: (3,) simulation_box_dimensions
+
+        Returns:
+            np.ndarray: (3,) moment of inertia array.
+
+        """
+
+        translated_coords = self.get_vector(center_of_mass, positions, dimensions)
+        custom_moment_of_inertia = np.zeros(3, dtype=float)
+
+        for coord, mass in zip(translated_coords, masses, strict=True):
+            axis_component = np.sum(
+                np.cross(custom_rot_axes, coord) ** 2 * mass, axis=1
+            )
+            custom_moment_of_inertia += axis_component
+
+        return custom_moment_of_inertia
+
     def get_custom_moment_of_inertia(
         self,
         UA,
@@ -780,7 +794,6 @@ def get_moment_of_inertia_tensor(
         """
         r = self.get_vector(center_of_mass, positions, dimensions)
         r2 = np.sum(r**2, axis=1)
-
         masses_arr = np.asarray(list(masses), dtype=float)
         moment_of_inertia_tensor = np.eye(3) * np.sum(masses_arr * r2)
         moment_of_inertia_tensor -= np.einsum("i,ij,ik->jk", masses_arr, r, r)
@@ -812,6 +825,7 @@ def get_custom_principal_axes(
                 - principal_axes: (3, 3) principal axes (rows).
                 - moment_of_inertia: (3,) principal moments.
         """
+
         eigenvalues, eigenvectors = np.linalg.eig(moment_of_inertia_tensor)
         order = np.abs(eigenvalues).argsort()[::-1]  # descending order
         transposed = np.transpose(eigenvectors)  # columns -> rows
@@ -849,3 +863,78 @@ def get_UA_masses(self, molecule) -> list[float]:
                     ua_mass += float(h.mass)
                 ua_masses.append(ua_mass)
         return ua_masses
+
+    def get_chain(self, residue, first, last):
+        """
+        For a given MDAnalysis AtomGroup and two given heavy atoms
+        within that AtomGroup, return the
+        shortest path between the two atoms.
+
+        Args:
+            residue: MDAnalysis AtomGroup representing
+            the residue/monomer of interest.
+            first: First heavy atom in the chain
+            last: Last heavy atom in the chain
+
+        Returns:
+            chain: Array containing chain atoms.
+        """
+
+        chain = []
+        # at the beggining we've only visited the first atom
+        visited_dict = {first: True}
+        # keep the previous atom to trace back the path
+        prev = {}
+        # queue of next heavy atoms to visit
+        next_to_visit = [first]
+        # all others heavy atoms in the residue, we have not yet visited
+        remaining_heavy_atoms = residue.atoms.select_atoms(
+            f"(mass 2 to 999) and not index {first.index}"
+        )
+        for atom in remaining_heavy_atoms:
+            visited_dict[atom] = False
+        current = first
+
+        while not visited_dict[last]:
+            # we haven't found a path to the last residue
+            next_to_visit.pop(0)
+            # we're visiting the current atom => we remove it from the queue
+            bonded_atoms = residue.atoms.select_atoms(
+                f"(mass 2 to 999) and bonded index {current.index}"
+            )
+
+            if last in bonded_atoms:
+                # we found a path to the last atom
+                visited_dict[last] = True
+                chain.append(last)
+                prev[last] = current
+
+            else:
+                for bonded_atom in bonded_atoms:
+                    # look for unvisited bonded atoms to the current atom we're visiting
+                    if not visited_dict[bonded_atom]:
+                        # we're going to want to visit the atoms
+                        next_to_visit.append(bonded_atom)
+                        prev[bonded_atom] = current
+                # we visit the next atom in the queue
+                current = next_to_visit[0]
+                visited_dict[current] = True
+
+        # we track the previous atom back to the first atom now
+        current = last
+        chain = [last]
+        # subtract index of first atom in resid
+        # most likely will coincide with first
+        # but this will work even if it doesn't
+        # accout for in-residue index
+        # start from last atom in chain
+        while chain[-1] != first:
+            # we haven't yet returned to the first atom
+            current = prev[current]
+            chain.append(current)
+
+        chain = np.flip(chain)
+        # only get in between residues
+        chain = chain[1:-1]
+        # accout for in-residue index
+        return chain
diff --git a/CodeEntropy/levels/nodes/covariance.py b/CodeEntropy/levels/nodes/covariance.py
index 39ff61a..a2c73aa 100644
--- a/CodeEntropy/levels/nodes/covariance.py
+++ b/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,27 +561,20 @@ 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:
@@ -584,20 +582,9 @@ def _get_residue_axes(
         """
         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)
diff --git a/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py b/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py
index 9d66e88..da0cb07 100644
--- a/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py
+++ b/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py
@@ -1,3 +1,5 @@
+"""Atomic unit tests for frame-local covariance construction."""
+
 from __future__ import annotations
 
 from types import SimpleNamespace
@@ -86,7 +88,6 @@ def test_run_processes_all_levels_and_writes_frame_covariance():
     mol = FakeMolecule()
     universe = FakeUniverse([mol], dimensions=np.array([10.0, 20.0, 30.0, 90.0]))
     axes_manager = object()
-    axes_topology = object()
 
     ctx = {
         "shared": {
@@ -96,7 +97,6 @@ def test_run_processes_all_levels_and_writes_frame_covariance():
             "beads": {},
             "args": _args(combined_forcetorque=True, customised_axes=True),
             "axes_manager": axes_manager,
-            "axes_topology": axes_topology,
         }
     }
 
@@ -115,14 +115,10 @@ def test_run_processes_all_levels_and_writes_frame_covariance():
     assert ua_kwargs["mol_id"] == 0
     assert ua_kwargs["group_id"] == 7
     assert ua_kwargs["axes_manager"] is axes_manager
-    assert ua_kwargs["axes_topology"] is axes_topology
     assert ua_kwargs["force_partitioning"] == 0.5
     assert ua_kwargs["customised_axes"] is True
     assert ua_kwargs["is_highest"] is False
 
-    res_kwargs = node._process_residue.call_args.kwargs
-    assert res_kwargs["axes_topology"] is axes_topology
-
 
 def test_run_omits_forcetorque_when_combined_is_false():
     node = FrameCovarianceNode()
@@ -169,7 +165,6 @@ def test_process_united_atom_updates_outputs_and_molcount():
         group_id=7,
         beads={(0, "united_atom", 0): [np.array([0])]},
         axes_manager="axes",
-        axes_topology=None,
         box=None,
         force_partitioning=0.5,
         customised_axes=False,
@@ -197,7 +192,6 @@ def test_process_united_atom_returns_when_no_beads_or_empty_atom_groups():
         group_id=7,
         beads={},
         axes_manager=None,
-        axes_topology=None,
         box=None,
         force_partitioning=0.5,
         customised_axes=False,
@@ -219,7 +213,6 @@ def test_process_united_atom_returns_when_no_beads_or_empty_atom_groups():
         group_id=7,
         beads={(0, "united_atom", 0): [np.array([0])]},
         axes_manager=None,
-        axes_topology=None,
         box=None,
         force_partitioning=0.5,
         customised_axes=False,
@@ -256,7 +249,6 @@ def test_process_residue_updates_outputs_and_combined_ft():
         group_id=7,
         beads={(0, "residue"): [np.array([0])]},
         axes_manager="axes",
-        axes_topology=None,
         box=None,
         customised_axes=True,
         force_partitioning=0.5,
@@ -286,7 +278,6 @@ def test_process_residue_returns_when_no_beads_or_empty_groups():
         group_id=7,
         beads={},
         axes_manager=None,
-        axes_topology=None,
         box=None,
         customised_axes=False,
         force_partitioning=0.5,
@@ -308,7 +299,6 @@ def test_process_residue_returns_when_no_beads_or_empty_groups():
         group_id=7,
         beads={(0, "residue"): [np.array([0])]},
         axes_manager=None,
-        axes_topology=None,
         box=None,
         customised_axes=False,
         force_partitioning=0.5,
@@ -424,17 +414,14 @@ def test_build_ua_vectors_uses_customised_axes():
     node._ft.get_weighted_torques = MagicMock(return_value=np.array([0.0, 1.0, 0.0]))
 
     force_vecs, torque_vecs = node._build_ua_vectors(
-        u=FakeUniverse([]),
-        mol_id=0,
-        local_res_i=0,
         bead_groups=[FakeAtomGroup("ua")],
-        residue_atoms=FakeAtomGroup("res"),
+        residue_group=FakeAtomGroup("res"),
         axes_manager=axes_manager,
-        axes_topology=None,
         box=None,
         force_partitioning=0.5,
         customised_axes=True,
         is_highest=True,
+        res_position=None,
     )
 
     assert len(force_vecs) == 1
@@ -442,49 +429,6 @@ def test_build_ua_vectors_uses_customised_axes():
     axes_manager.get_UA_axes.assert_called_once()
 
 
-def test_build_ua_vectors_uses_cached_axes_topology_when_available():
-    node = FrameCovarianceNode()
-    axes_manager = MagicMock()
-
-    u = FakeUniverse([])
-    ua_topology = object()
-    axes_topology = SimpleNamespace(ua={(3, 4, 0): ua_topology})
-
-    axes_manager.get_UA_axes_from_topology.return_value = (
-        np.eye(3),
-        2.0 * np.eye(3),
-        np.ones(3),
-        np.array([1.0, 2.0, 3.0]),
-    )
-    node._ft.get_weighted_forces = MagicMock(return_value=np.array([1.0, 0.0, 0.0]))
-    node._ft.get_weighted_torques = MagicMock(return_value=np.array([0.0, 1.0, 0.0]))
-
-    force_vecs, torque_vecs = node._build_ua_vectors(
-        u=u,
-        mol_id=3,
-        local_res_i=4,
-        bead_groups=[FakeAtomGroup("ua")],
-        residue_atoms=FakeAtomGroup("res"),
-        axes_manager=axes_manager,
-        axes_topology=axes_topology,
-        box=None,
-        force_partitioning=0.5,
-        customised_axes=True,
-        is_highest=True,
-    )
-
-    assert len(force_vecs) == 1
-    assert len(torque_vecs) == 1
-
-    called_kwargs = axes_manager.get_UA_axes_from_topology.call_args.kwargs
-    assert called_kwargs["u"] is u
-    assert called_kwargs["topology"] is ua_topology
-    assert called_kwargs["box"] is None
-    assert called_kwargs["residue_atoms"].name == "res"
-
-    axes_manager.get_UA_axes.assert_not_called()
-
-
 def test_build_ua_vectors_uses_vanilla_axes_when_not_customised():
     node = FrameCovarianceNode()
     axes_manager = MagicMock()
@@ -492,31 +436,34 @@ def test_build_ua_vectors_uses_vanilla_axes_when_not_customised():
         np.eye(3),
         np.array([1.0, 2.0, 3.0]),
     )
+    FakeAtomGroup_Atoms = MagicMock(positions=np.zeros((2, 3)))
+    FakeAtomGroup.atoms = FakeAtomGroup_Atoms
+    axes_manager.principal_axes.return_value = np.eye(3)
     node._ft.get_weighted_forces = MagicMock(return_value=np.array([1.0, 0.0, 0.0]))
     node._ft.get_weighted_torques = MagicMock(return_value=np.array([0.0, 1.0, 0.0]))
 
     with patch("CodeEntropy.levels.nodes.covariance.make_whole") as make_whole:
         node._build_ua_vectors(
-            u=FakeUniverse([]),
-            mol_id=0,
-            local_res_i=0,
             bead_groups=[FakeAtomGroup("ua")],
-            residue_atoms=FakeAtomGroup("res"),
+            residue_group=FakeAtomGroup("res"),
             axes_manager=axes_manager,
-            axes_topology=None,
             box=None,
             force_partitioning=0.5,
             customised_axes=False,
             is_highest=False,
+            res_position=None,
         )
 
     assert make_whole.call_count == 2
     axes_manager.get_vanilla_axes.assert_called_once()
 
 
-def test_build_residue_vectors_uses_residue_axes():
+def test_build_residue_vectors_uses_residue_axes(monkeypatch):
     node = FrameCovarianceNode()
     mol = FakeMolecule(n_residues=1)
+    residue = MagicMock()
+    mol.residues = [residue]
+    monkeypatch.setattr(residue, "resindex", lambda resindex: 0)
     axes_manager = MagicMock()
 
     node._get_residue_axes = MagicMock(
@@ -526,12 +473,9 @@ def test_build_residue_vectors_uses_residue_axes():
     node._ft.get_weighted_torques = MagicMock(return_value=np.array([0.0, 1.0, 0.0]))
 
     force_vecs, torque_vecs = node._build_residue_vectors(
-        u=FakeUniverse([mol]),
         mol=mol,
-        mol_id=0,
         bead_groups=[FakeAtomGroup("res")],
         axes_manager=axes_manager,
-        axes_topology=None,
         box=None,
         customised_axes=True,
         force_partitioning=0.5,
@@ -543,34 +487,6 @@ def test_build_residue_vectors_uses_residue_axes():
     node._get_residue_axes.assert_called_once()
 
 
-def test_get_residue_axes_customised_uses_cached_topology_when_available():
-    node = FrameCovarianceNode()
-    mol = FakeMolecule(n_residues=1)
-    axes_manager = MagicMock()
-    expected = (np.eye(3), np.eye(3) * 2.0, np.zeros(3), np.ones(3))
-    residue_topology = object()
-    axes_topology = SimpleNamespace(residue={(3, 0): residue_topology})
-    axes_manager.get_residue_axes_from_topology.return_value = expected
-
-    result = node._get_residue_axes(
-        u=FakeUniverse([mol]),
-        mol=mol,
-        mol_id=3,
-        bead=FakeAtomGroup("res"),
-        local_res_i=0,
-        axes_manager=axes_manager,
-        axes_topology=axes_topology,
-        box=None,
-        customised_axes=True,
-    )
-
-    assert result == expected
-    called_kwargs = axes_manager.get_residue_axes_from_topology.call_args.kwargs
-    assert called_kwargs["topology"] is residue_topology
-    assert called_kwargs["residue_atoms"] is mol.residues[0].atoms
-    axes_manager.get_residue_axes.assert_not_called()
-
-
 def test_get_residue_axes_customised_delegates_to_axes_manager():
     node = FrameCovarianceNode()
     mol = FakeMolecule(n_residues=1)
@@ -580,14 +496,11 @@ def test_get_residue_axes_customised_delegates_to_axes_manager():
 
     assert (
         node._get_residue_axes(
-            u=FakeUniverse([mol]),
             mol=mol,
-            mol_id=0,
             bead=FakeAtomGroup("res"),
             local_res_i=0,
+            relative_res_i=0,
             axes_manager=axes_manager,
-            axes_topology=None,
-            box=None,
             customised_axes=True,
         )
         == expected
@@ -596,6 +509,7 @@ def test_get_residue_axes_customised_delegates_to_axes_manager():
     axes_manager.get_residue_axes.assert_called_once_with(
         mol,
         0,
+        0,
         residue=mol.residues[0].atoms,
     )
 
@@ -612,14 +526,11 @@ def test_get_residue_axes_vanilla_uses_make_whole_and_vanilla_axes():
 
     with patch("CodeEntropy.levels.nodes.covariance.make_whole") as make_whole:
         trans_axes, rot_axes, center, moi = node._get_residue_axes(
-            u=FakeUniverse([mol]),
             mol=mol,
-            mol_id=0,
             bead=bead,
             local_res_i=0,
+            relative_res_i=0,
             axes_manager=axes_manager,
-            axes_topology=None,
-            box=None,
             customised_axes=False,
         )
 
diff --git a/tests/unit/CodeEntropy/levels/test_axes.py b/tests/unit/CodeEntropy/levels/test_axes.py
index e68f4d1..9e1236b 100644
--- a/tests/unit/CodeEntropy/levels/test_axes.py
+++ b/tests/unit/CodeEntropy/levels/test_axes.py
@@ -4,7 +4,6 @@
 import pytest
 
 from CodeEntropy.levels.axes import AxesCalculator
-from CodeEntropy.levels.nodes.axes_topology import ResidueAxesTopology, UAAxesTopology
 
 
 class _FakeAtom:
@@ -76,7 +75,7 @@ def test_get_residue_axes_empty_residue_raises():
     u.select_atoms.return_value = []
 
     with pytest.raises(ValueError):
-        ax.get_residue_axes(u, index=5)
+        ax.get_residue_axes(u, index=5, relative_index=0)
 
 
 def test_get_residue_axes_no_bonds_uses_custom_principal_axes(monkeypatch):
@@ -110,7 +109,7 @@ def _select_atoms(q):
         lambda moi: (np.eye(3), np.array([3.0, 2.0, 1.0])),
     )
 
-    trans, rot, center, moi = ax.get_residue_axes(u, index=7)
+    trans, rot, center, moi = ax.get_residue_axes(u, index=7, relative_index=0)
 
     assert np.allclose(trans, np.eye(3))
     assert np.allclose(rot, np.eye(3))
@@ -118,13 +117,17 @@ def _select_atoms(q):
     assert np.allclose(moi, np.array([3.0, 2.0, 1.0]))
 
 
-def test_get_residue_axes_with_bonds_uses_vanilla_axes(monkeypatch):
+def test_get_residue_axes_uses_vanilla_axes(monkeypatch):
     ax = AxesCalculator()
 
     residue = MagicMock()
     residue.__len__.return_value = 1
     residue.atoms.center_of_mass.return_value = np.array([1.0, 2.0, 3.0])
     residue.center_of_mass.return_value = np.array([1.0, 2.0, 3.0])
+    residue.select_atoms.return_value = MagicMock(
+        positions=[[1.0, 2.0, 3.0], [3.0, 2.0, 1.0]]
+    )
+    uas = MagicMock(positions=np.zeros((2, 3)))
 
     u = MagicMock()
     u.dimensions = np.array([10.0, 10.0, 10.0, 90, 90, 90])
@@ -136,20 +139,27 @@ def _select_atoms(q):
             return [1]  # non-empty
         if q.startswith("resindex "):
             return residue
+        if q == "mass 2 to 999":
+            return uas
         return []
 
+    monkeypatch.setattr(ax, "get_UA_masses", lambda mol: [10.0, 12.0])
     u.select_atoms.side_effect = _select_atoms
+    residue = u.select_atoms("resindex 10")
 
     monkeypatch.setattr("CodeEntropy.levels.axes.make_whole", lambda _ag: None)
-    monkeypatch.setattr(
-        ax, "get_vanilla_axes", lambda mol: (np.eye(3) * 2, np.array([9.0, 8.0, 7.0]))
-    )
+    eigenvalues, eigenvectors = np.linalg.eig([[48, 0, 48], [0, 96, 0], [48, 0, 48]])
+    transposed = np.transpose(eigenvectors)
+    axes = transposed[[2, 0, 1]]
+    axes[2] = -axes[2]
 
-    trans, rot, center, moi = ax.get_residue_axes(u, index=10)
+    trans, rot, center, moi = ax.get_residue_axes(
+        u, index=10, relative_index=0, residue=residue
+    )
 
-    assert np.allclose(trans, np.eye(3))
-    assert np.allclose(rot, np.eye(3) * 2)
-    assert np.allclose(moi, np.array([9.0, 8.0, 7.0]))
+    assert np.allclose(trans, axes)
+    assert np.allclose(rot, axes)
+    assert np.allclose(moi, np.array([96.0, 96.0, 0.0]))
 
 
 def test_get_UA_axes_uses_principal_axes_when_single_heavy(monkeypatch):
@@ -158,13 +168,15 @@ def test_get_UA_axes_uses_principal_axes_when_single_heavy(monkeypatch):
     u = MagicMock()
     u.dimensions = np.array([10.0, 10.0, 10.0, 90, 90, 90])
     u.atoms.principal_axes.return_value = np.eye(3)
+    u.center_of_mass.return_value = np.array([[4.0, 0.0, 0.0]])
 
     # heavy_atoms length <= 1 => principal_axes path
-    heavy_atom = MagicMock(index=5)
+    heavy_atom = MagicMock(index=5, position=np.array([4.0, 0.0, 0.0]))
     heavy_atoms = [heavy_atom]
 
     def _sel(q):
-        if q == "prop mass > 1.1":
+        if q == "mass 2 to 999":
+            # return heavy atoms group
             return heavy_atoms
         if q.startswith("index "):
             # return atom group with positions
@@ -173,6 +185,7 @@ def _sel(q):
             ag.__getitem__.return_value = MagicMock(
                 mass=12.0, position=np.array([4.0, 0.0, 0.0]), index=5
             )
+
             return ag
         return []
 
@@ -185,7 +198,7 @@ def _sel(q):
         lambda system, atom, dimensions: (np.eye(3), np.array([1.0, 2.0, 3.0])),
     )
 
-    trans, rot, center, moi = ax.get_UA_axes(u, index=0)
+    trans, rot, center, moi = ax.get_UA_axes(u, index=0, res_position=None)
 
     assert np.allclose(trans, np.eye(3))
     assert np.allclose(rot, np.eye(3))
@@ -202,7 +215,7 @@ def test_get_UA_axes_raises_when_bonded_axes_fail(monkeypatch):
     heavy_atoms = [heavy_atom]
 
     def _sel(q):
-        if q == "prop mass > 1.1":
+        if q == "mass 2 to 999":
             return heavy_atoms
         if q.startswith("index "):
             ag = MagicMock()
@@ -218,7 +231,7 @@ def _sel(q):
     monkeypatch.setattr(ax, "get_bonded_axes", lambda **kwargs: (None, None))
 
     with pytest.raises(ValueError):
-        ax.get_UA_axes(u, index=0)
+        ax.get_UA_axes(u, index=5, res_position=None)
 
 
 def test_get_custom_axes_degenerate_axis1_raises():
@@ -497,7 +510,7 @@ def _select_atoms(q):
         lambda moi: (np.eye(3), np.array([3.0, 2.0, 1.0])),
     )
 
-    trans, rot, center, moi = ax.get_residue_axes(u, index=7)
+    trans, rot, center, moi = ax.get_residue_axes(u, index=7, relative_index=0)
 
     assert trans.shape == (3, 3)
     assert rot.shape == (3, 3)
@@ -505,18 +518,16 @@ def _select_atoms(q):
     assert np.allclose(moi, np.array([3.0, 2.0, 1.0]))
 
 
-def test_get_residue_axes_with_bonds_vanilla_path(monkeypatch):
+def test_get_residue_axes_vanilla_path(monkeypatch):
     ax = AxesCalculator()
-
     residue = MagicMock()
     residue.__len__.return_value = 1
-    residue.atoms.principal_axes.return_value = np.eye(3) * 2
     residue.atoms.center_of_mass.return_value = np.array([1.0, 2.0, 3.0])
     residue.center_of_mass.return_value = np.array([1.0, 2.0, 3.0])
+    residue.select_atoms.return_value = MagicMock(positions=np.zeros((1, 3)))
 
     u = MagicMock()
     u.dimensions = np.array([10.0, 10.0, 10.0, 90, 90, 90])
-    u.atoms.principal_axes.return_value = np.eye(3) * 2
 
     def _select_atoms(q):
         if q.startswith("(resindex"):
@@ -529,10 +540,14 @@ def _select_atoms(q):
 
     monkeypatch.setattr("CodeEntropy.levels.axes.make_whole", lambda _ag: None)
     monkeypatch.setattr(
-        ax, "get_vanilla_axes", lambda mol: (np.eye(3) * 2, np.array([9.0, 8.0, 7.0]))
+        ax,
+        "get_custom_principal_axes",
+        lambda mol: (np.eye(3) * 2, np.array([9.0, 8.0, 7.0])),
     )
 
-    trans, rot, center, moi = ax.get_residue_axes(u, index=10)
+    trans, rot, center, moi = ax.get_residue_axes(
+        u, index=10, relative_index=0, residue=residue
+    )
 
     assert np.allclose(trans, np.eye(3) * 2)
     assert np.allclose(rot, np.eye(3) * 2)
@@ -621,13 +636,13 @@ def test_get_UA_axes_multiple_heavy_atoms_uses_custom_principal_axes(monkeypatch
 
     heavy_atoms = _FakeAtomGroup(
         [
-            _FakeAtom(0, 12.0, [0, 0, 0]),
-            _FakeAtom(1, 12.0, [1, 0, 0]),
+            _atom(index=0, mass=12.0, pos=[0, 0, 0]),
+            _atom(index=1, mass=12.0, pos=[1, 0, 0]),
         ],
         positions=np.array([[0, 0, 0], [1, 0, 0]], dtype=float),
     )
 
-    system_atom = _FakeAtom(index=0, mass=12.0, position=[0, 0, 0])
+    system_atom = _atom(index=0, mass=12.0, pos=[0, 0, 0])
     heavy_atom_selection = _FakeAtomGroup(
         [system_atom], positions=np.array([[0, 0, 0]], dtype=float)
     )
@@ -639,12 +654,23 @@ def center_of_mass(self, *args, **kwargs):
         def __getitem__(self, idx):
             return system_atom
 
+        def principal_axes(self, *args, **kwargs):
+            return np.eye(3)
+
+        def select_atoms(self, q):
+            if q == "mass 2 to 999":
+                return heavy_atoms
+            if q.startswith("index "):
+                return [heavy_atom_selection[0]]
+
     data_container = MagicMock()
     data_container.atoms = _Atoms()
+    data_container.residues.__len__.return_value = 1
     data_container.dimensions = np.array([10.0, 10.0, 10.0, 90, 90, 90], dtype=float)
+    _FakeAtomGroup.atoms = heavy_atom_selection
 
     def _select_atoms(q):
-        if q == "prop mass > 1.1":
+        if q == "mass 2 to 999":
             return heavy_atoms
         if q.startswith("index "):
             return heavy_atom_selection
@@ -659,20 +685,14 @@ def _select_atoms(q):
     )
     monkeypatch.setattr(ax, "get_UA_masses", lambda _ag: [12.0, 12.0])
 
-    got_tensor = MagicMock(return_value=np.eye(3))
-    monkeypatch.setattr(ax, "get_moment_of_inertia_tensor", got_tensor)
-
-    got_custom_axes = MagicMock(return_value=(np.eye(3), np.array([3.0, 2.0, 1.0])))
-    monkeypatch.setattr(ax, "get_custom_principal_axes", got_custom_axes)
-
-    trans_axes, rot_axes, center, moi = ax.get_UA_axes(data_container, index=0)
+    trans_axes, rot_axes, center, moi = ax.get_UA_axes(
+        data_container, index=0, res_position=None
+    )
 
     assert trans_axes.shape == (3, 3)
     assert rot_axes.shape == (3, 3)
     assert np.allclose(center, np.array([0.0, 0.0, 0.0]))
     assert moi.shape == (3,)
-    got_tensor.assert_called_once()
-    got_custom_axes.assert_called_once()
 
 
 def test_get_bonded_axes_returns_none_none_if_custom_axes_none(monkeypatch):
@@ -702,497 +722,166 @@ def test_get_bonded_axes_returns_none_none_if_custom_axes_none(monkeypatch):
     assert moi is None
 
 
-class _FakeIndexedAtoms:
-    """Container supporting ``u.atoms[index]`` and ``u.atoms[index_array]``."""
-
-    def __init__(self, atom_map):
-        self._atom_map = dict(atom_map)
-
-    def __getitem__(self, index):
-        if isinstance(index, np.ndarray):
-            return _FakeAtomGroup([self._atom_map[int(i)] for i in index])
-        if isinstance(index, (list, tuple)):
-            return _FakeAtomGroup([self._atom_map[int(i)] for i in index])
-        return self._atom_map[int(index)]
-
-
-class _FakeUniverse:
-    """Small universe-like object with indexed atoms and dimensions."""
-
-    def __init__(self, atom_map, dimensions=None):
-        self.atoms = _FakeIndexedAtoms(atom_map)
-        self.dimensions = np.asarray(
-            dimensions
-            if dimensions is not None
-            else [10.0, 10.0, 10.0, 90.0, 90.0, 90.0],
-            dtype=float,
-        )
-
-
-def _ua_topology(
-    *,
-    heavy_atom_index=1,
-    ua_atom_indices=(1,),
-    ua_all_atom_indices=(1,),
-    bonded_heavy_indices=(),
-    bonded_light_indices=(),
-    residue_heavy_indices=(1,),
-    residue_ua_masses=(12.0,),
-):
-    """Build a small cached UA topology fixture."""
-    return UAAxesTopology(
-        heavy_atom_index=int(heavy_atom_index),
-        ua_atom_indices=np.asarray(ua_atom_indices, dtype=int),
-        ua_all_atom_indices=np.asarray(ua_all_atom_indices, dtype=int),
-        bonded_heavy_indices=np.asarray(bonded_heavy_indices, dtype=int),
-        bonded_light_indices=np.asarray(bonded_light_indices, dtype=int),
-        residue_heavy_indices=np.asarray(residue_heavy_indices, dtype=int),
-        residue_ua_masses=np.asarray(residue_ua_masses, dtype=float),
-    )
-
-
-def _residue_topology(
-    *,
-    residue_heavy_indices=(1,),
-    residue_ua_masses=(12.0,),
-    has_neighbor_bonds=False,
-):
-    """Build a small cached residue topology fixture."""
-    return ResidueAxesTopology(
-        residue_heavy_indices=np.asarray(residue_heavy_indices, dtype=int),
-        residue_ua_masses=np.asarray(residue_ua_masses, dtype=float),
-        has_neighbor_bonds=bool(has_neighbor_bonds),
-    )
-
-
-def test_get_residue_axes_from_topology_no_neighbor_bonds_uses_cached_indices(
-    monkeypatch,
-):
+def test_get_residue_axes_custom_path(monkeypatch):
     ax = AxesCalculator()
 
-    heavy_atom = _FakeAtom(1, 12.0, [1.0, 2.0, 3.0])
-    other_heavy = _FakeAtom(3, 14.0, [4.0, 5.0, 6.0])
-    universe = _FakeUniverse({1: heavy_atom, 3: other_heavy})
-    mol = MagicMock()
-    residue_atoms = MagicMock()
-    residue_atoms.center_of_mass.return_value = np.array([9.0, 8.0, 7.0])
-    topology = _residue_topology(
-        residue_heavy_indices=(1, 3),
-        residue_ua_masses=(13.0, 14.0),
-        has_neighbor_bonds=False,
+    edge_atoms = _FakeAtomGroup(
+        [_FakeAtom(8, 12.0, [1, 0, 0]), _FakeAtom(10, 14.0, [0, 0, 0])],
+        positions=np.array([[1.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=float),
     )
 
-    get_tensor = MagicMock(return_value=np.eye(3))
-    get_principal = MagicMock(return_value=(np.eye(3) * 2.0, np.array([3.0, 2.0, 1.0])))
-
-    monkeypatch.setattr(ax, "get_moment_of_inertia_tensor", get_tensor)
-    monkeypatch.setattr(ax, "get_custom_principal_axes", get_principal)
-
-    box = np.array([20.0, 30.0, 40.0])
-    trans_axes, rot_axes, center, moi = ax.get_residue_axes_from_topology(
-        u=universe,
-        mol=mol,
-        residue_atoms=residue_atoms,
-        topology=topology,
-        box=box,
+    backbone_center = np.array([0.0, 1.0, 0.0])
+    rot_center, rot_axes = ax.get_residue_custom_axes(
+        [edge_atoms[0], edge_atoms[1]], backbone_center
     )
 
-    np.testing.assert_allclose(trans_axes, np.eye(3) * 2.0)
-    np.testing.assert_allclose(rot_axes, np.eye(3) * 2.0)
-    np.testing.assert_allclose(center, np.array([9.0, 8.0, 7.0]))
-    np.testing.assert_allclose(moi, np.array([3.0, 2.0, 1.0]))
-
-    tensor_kwargs = get_tensor.call_args.kwargs
-    np.testing.assert_allclose(
-        tensor_kwargs["center_of_mass"],
-        np.array([9.0, 8.0, 7.0]),
-    )
-    np.testing.assert_allclose(
-        tensor_kwargs["positions"],
-        np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
-    )
-    np.testing.assert_allclose(tensor_kwargs["masses"], np.array([13.0, 14.0]))
-    np.testing.assert_allclose(tensor_kwargs["dimensions"], box)
-    get_principal.assert_called_once()
-
-
-def test_get_residue_axes_from_topology_neighbor_bonds_uses_vanilla_axes(
-    monkeypatch,
-):
-    ax = AxesCalculator()
-
-    universe = _FakeUniverse(
-        {},
-        dimensions=[11.0, 12.0, 13.0, 90.0, 90.0, 90.0],
-    )
-    mol = MagicMock()
-    mol.atoms.principal_axes.return_value = np.eye(3) * 5.0
-    residue_atoms = MagicMock()
-    residue_atoms.center_of_mass.return_value = np.array([1.0, 2.0, 3.0])
-    topology = _residue_topology(has_neighbor_bonds=True)
-
-    make_whole = MagicMock()
-    get_vanilla = MagicMock(return_value=(np.eye(3) * 6.0, np.array([6.0, 5.0, 4.0])))
-
-    monkeypatch.setattr("CodeEntropy.levels.axes.make_whole", make_whole)
-    monkeypatch.setattr(ax, "get_vanilla_axes", get_vanilla)
-
-    trans_axes, rot_axes, center, moi = ax.get_residue_axes_from_topology(
-        u=universe,
-        mol=mol,
-        residue_atoms=residue_atoms,
-        topology=topology,
-        box=None,
-    )
-
-    make_whole.assert_called_once_with(mol.atoms)
-    mol.atoms.principal_axes.assert_called_once()
-    get_vanilla.assert_called_once_with(residue_atoms)
-    np.testing.assert_allclose(trans_axes, np.eye(3) * 5.0)
-    np.testing.assert_allclose(rot_axes, np.eye(3) * 6.0)
-    np.testing.assert_allclose(center, np.array([1.0, 2.0, 3.0]))
-    np.testing.assert_allclose(moi, np.array([6.0, 5.0, 4.0]))
-
-
-def test_get_UA_axes_from_topology_multiple_heavy_uses_cached_indices_and_box(
-    monkeypatch,
-):
-    ax = AxesCalculator()
-
-    heavy_atom = _FakeAtom(1, 12.0, [1.0, 2.0, 3.0])
-    other_heavy = _FakeAtom(3, 14.0, [4.0, 5.0, 6.0])
-    universe = _FakeUniverse({1: heavy_atom, 3: other_heavy})
-    residue_atoms = MagicMock()
-    residue_atoms.center_of_mass.return_value = np.array([9.0, 8.0, 7.0])
-
-    topology = _ua_topology(
-        heavy_atom_index=1,
-        residue_heavy_indices=(1, 3),
-        residue_ua_masses=(13.0, 14.0),
-    )
-
-    get_tensor = MagicMock(return_value=np.eye(3))
-    get_principal = MagicMock(return_value=(np.eye(3) * 2.0, np.array([3.0, 2.0, 1.0])))
-    get_bonded = MagicMock(return_value=(np.eye(3) * 4.0, np.array([1.0, 1.0, 1.0])))
-
-    monkeypatch.setattr(ax, "get_moment_of_inertia_tensor", get_tensor)
-    monkeypatch.setattr(ax, "get_custom_principal_axes", get_principal)
-    monkeypatch.setattr(ax, "get_bonded_axes_from_topology", get_bonded)
-
-    box = np.array([20.0, 30.0, 40.0])
-    trans_axes, rot_axes, center, moi = ax.get_UA_axes_from_topology(
-        u=universe,
-        residue_atoms=residue_atoms,
-        topology=topology,
-        box=box,
-    )
-
-    np.testing.assert_allclose(trans_axes, np.eye(3) * 2.0)
-    np.testing.assert_allclose(rot_axes, np.eye(3) * 4.0)
-    np.testing.assert_allclose(center, heavy_atom.position)
-    np.testing.assert_allclose(moi, np.array([1.0, 1.0, 1.0]))
-
-    get_tensor.assert_called_once()
-    tensor_kwargs = get_tensor.call_args.kwargs
-    np.testing.assert_allclose(
-        tensor_kwargs["center_of_mass"], np.array([9.0, 8.0, 7.0])
-    )
-    np.testing.assert_allclose(
-        tensor_kwargs["positions"], np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
-    )
-    np.testing.assert_allclose(tensor_kwargs["masses"], np.array([13.0, 14.0]))
-    np.testing.assert_allclose(tensor_kwargs["dimensions"], box)
-
-    get_principal.assert_called_once()
-    np.testing.assert_allclose(get_principal.call_args.args[0], np.eye(3))
-    get_bonded.assert_called_once_with(
-        u=universe,
-        heavy_atom=heavy_atom,
-        topology=topology,
-        dimensions=box,
-    )
+    assert rot_center.shape == (3,)
+    assert rot_axes.shape == (3, 3)
 
 
-def test_get_UA_axes_from_topology_single_heavy_uses_residue_principal_axes(
-    monkeypatch,
-):
+def test_get_custom_residue_moment_of_inertia(monkeypatch):
     ax = AxesCalculator()
-
-    heavy_atom = _FakeAtom(1, 12.0, [1.0, 0.0, 0.0])
-    universe = _FakeUniverse(
-        {1: heavy_atom}, dimensions=[11.0, 12.0, 13.0, 90.0, 90.0, 90.0]
+    heavy_atoms = _FakeAtomGroup(
+        [_FakeAtom(0, 12.0, [1, 2, 1]), _FakeAtom(1, 12.0, [2, 1, 1])],
+        positions=np.array([[1, 2, 1], [2, 1, 1]], dtype=float),
     )
-    residue_atoms = MagicMock()
-    residue_atoms.principal_axes.return_value = np.eye(3) * 5.0
-
-    topology = _ua_topology(heavy_atom_index=1, residue_heavy_indices=(1,))
-
-    make_whole = MagicMock()
-    get_bonded = MagicMock(return_value=(np.eye(3) * 6.0, np.array([6.0, 5.0, 4.0])))
-
-    monkeypatch.setattr("CodeEntropy.levels.axes.make_whole", make_whole)
-    monkeypatch.setattr(ax, "get_bonded_axes_from_topology", get_bonded)
+    dimensions = np.array([10.0, 10.0, 10.0], dtype=float)
 
-    trans_axes, rot_axes, center, moi = ax.get_UA_axes_from_topology(
-        u=universe,
-        residue_atoms=residue_atoms,
-        topology=topology,
-        box=None,
+    moi = ax.get_custom_residue_moment_of_inertia(
+        center_of_mass=np.array([1, 1, 1]),
+        positions=heavy_atoms.positions,
+        masses=heavy_atoms.masses,
+        custom_rot_axes=np.eye(3),
+        dimensions=dimensions,
     )
 
-    make_whole.assert_called_once_with(residue_atoms)
-    residue_atoms.principal_axes.assert_called_once()
-    np.testing.assert_allclose(trans_axes, np.eye(3) * 5.0)
-    np.testing.assert_allclose(rot_axes, np.eye(3) * 6.0)
-    np.testing.assert_allclose(center, heavy_atom.position)
-    np.testing.assert_allclose(moi, np.array([6.0, 5.0, 4.0]))
-
-    called_kwargs = get_bonded.call_args.kwargs
-    np.testing.assert_allclose(
-        called_kwargs["dimensions"], np.array([11.0, 12.0, 13.0])
-    )
+    assert moi.shape == (3,)
 
 
-def test_get_UA_axes_from_topology_raises_when_cached_bonded_axes_fail(monkeypatch):
+def test_get_residue_bonded_axes_2neighbours(monkeypatch):
     ax = AxesCalculator()
-
-    heavy_atom = _FakeAtom(1, 12.0, [1.0, 0.0, 0.0])
-    universe = _FakeUniverse({1: heavy_atom})
-    residue_atoms = MagicMock()
-    residue_atoms.principal_axes.return_value = np.eye(3)
-    topology = _ua_topology(heavy_atom_index=1, residue_heavy_indices=(1,))
-
+    u = MagicMock()
+    u.dimensions = np.array([10.0, 10.0, 10.0, 90, 90, 90])
     monkeypatch.setattr("CodeEntropy.levels.axes.make_whole", lambda _ag: None)
-    monkeypatch.setattr(
-        ax, "get_bonded_axes_from_topology", lambda **kwargs: (None, None)
-    )
-
-    with pytest.raises(ValueError, match="cached UA bead"):
-        ax.get_UA_axes_from_topology(
-            u=universe,
-            residue_atoms=residue_atoms,
-            topology=topology,
-            box=None,
-        )
-
-
-def test_get_bonded_axes_from_topology_non_heavy_returns_none_none():
-    ax = AxesCalculator()
-    light_atom = _FakeAtom(1, 1.0, [0.0, 0.0, 0.0])
-
-    custom_axes, moi = ax.get_bonded_axes_from_topology(
-        u=MagicMock(),
-        heavy_atom=light_atom,
-        topology=_ua_topology(heavy_atom_index=1),
-        dimensions=np.array([10.0, 10.0, 10.0]),
-    )
-
-    assert custom_axes is None
-    assert moi is None
-
-
-def test_get_bonded_axes_from_topology_no_bonded_heavy_uses_vanilla_axes(
-    monkeypatch,
-):
-    ax = AxesCalculator()
-
-    heavy_atom = _FakeAtom(1, 12.0, [0.0, 0.0, 0.0])
-    hydrogen = _FakeAtom(2, 1.0, [1.0, 0.0, 0.0])
-    universe = _FakeUniverse({1: heavy_atom, 2: hydrogen})
-    topology = _ua_topology(
-        heavy_atom_index=1,
-        ua_atom_indices=(1, 2),
-        ua_all_atom_indices=(1, 2),
-        bonded_heavy_indices=(),
-        bonded_light_indices=(2,),
-    )
-
-    get_vanilla = MagicMock(return_value=(np.eye(3) * 7.0, np.array([7.0, 8.0, 9.0])))
-    get_custom_moi = MagicMock()
-    get_flipped = MagicMock(return_value=np.eye(3) * -7.0)
-
-    monkeypatch.setattr(ax, "get_vanilla_axes", get_vanilla)
-    monkeypatch.setattr(ax, "get_custom_moment_of_inertia", get_custom_moi)
-    monkeypatch.setattr(ax, "get_flipped_axes", get_flipped)
-
-    custom_axes, moi = ax.get_bonded_axes_from_topology(
-        u=universe,
-        heavy_atom=heavy_atom,
-        topology=topology,
-        dimensions=np.array([10.0, 10.0, 10.0]),
-    )
-
-    np.testing.assert_allclose(custom_axes, np.eye(3) * -7.0)
-    np.testing.assert_allclose(moi, np.array([7.0, 8.0, 9.0]))
-    get_vanilla.assert_called_once()
-    get_custom_moi.assert_not_called()
-    get_flipped.assert_called_once()
-
-
-def test_get_bonded_axes_from_topology_one_heavy_no_light_uses_custom_axes(
-    monkeypatch,
-):
-    ax = AxesCalculator()
-
-    heavy_atom = _FakeAtom(1, 12.0, [0.0, 0.0, 0.0])
-    bonded_heavy = _FakeAtom(3, 12.0, [1.0, 0.0, 0.0])
-    universe = _FakeUniverse({1: heavy_atom, 3: bonded_heavy})
-    topology = _ua_topology(
-        heavy_atom_index=1,
-        ua_atom_indices=(1,),
-        ua_all_atom_indices=(1, 3),
-        bonded_heavy_indices=(3,),
-        bonded_light_indices=(),
-    )
-
-    get_custom_axes = MagicMock(return_value=np.eye(3) * 2.0)
-    get_custom_moi = MagicMock(return_value=np.array([2.0, 3.0, 4.0]))
-    get_flipped = MagicMock(return_value=np.eye(3) * 3.0)
-
-    monkeypatch.setattr(ax, "get_custom_axes", get_custom_axes)
-    monkeypatch.setattr(ax, "get_custom_moment_of_inertia", get_custom_moi)
-    monkeypatch.setattr(ax, "get_flipped_axes", get_flipped)
+    residue = u.select_atoms("resindex 1")
+    residue.__len__.return_value = 1
 
-    custom_axes, moi = ax.get_bonded_axes_from_topology(
-        u=universe,
-        heavy_atom=heavy_atom,
-        topology=topology,
-        dimensions=np.array([10.0, 10.0, 10.0]),
+    edge_atom_set = _FakeAtomGroup(
+        [
+            _atom(index=8, mass=12.0, pos=[1, 0, 0]),
+            _atom(index=10, mass=14.0, pos=[0, 0, 0]),
+        ],
+        positions=np.array([[1.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=float),
     )
+    backbone_atom = _atom(index=9, mass=12.0, pos=[0, 1, 0])
 
-    np.testing.assert_allclose(custom_axes, np.eye(3) * 3.0)
-    np.testing.assert_allclose(moi, np.array([2.0, 3.0, 4.0]))
-
-    kwargs = get_custom_axes.call_args.kwargs
-    np.testing.assert_allclose(kwargs["a"], heavy_atom.position)
-    np.testing.assert_allclose(kwargs["b_list"][0], bonded_heavy.position)
-    np.testing.assert_allclose(kwargs["c"], np.zeros(3))
-    get_custom_moi.assert_called_once()
-    get_flipped.assert_called_once()
-
-
-def test_get_bonded_axes_from_topology_one_heavy_with_light_uses_light_as_c(
-    monkeypatch,
-):
-    ax = AxesCalculator()
-
-    heavy_atom = _FakeAtom(1, 12.0, [0.0, 0.0, 0.0])
-    bonded_heavy = _FakeAtom(3, 12.0, [1.0, 0.0, 0.0])
-    bonded_light = _FakeAtom(2, 1.0, [0.0, 1.0, 0.0])
-    universe = _FakeUniverse({1: heavy_atom, 2: bonded_light, 3: bonded_heavy})
-    topology = _ua_topology(
-        heavy_atom_index=1,
-        ua_atom_indices=(1, 2),
-        ua_all_atom_indices=(1, 3, 2),
-        bonded_heavy_indices=(3,),
-        bonded_light_indices=(2,),
-    )
+    def _select_atoms(q):
+        if q.endswith("(bonded resindex 0 or resindex 2)"):
+            return edge_atom_set
 
-    get_custom_axes = MagicMock(return_value=np.eye(3))
-    monkeypatch.setattr(ax, "get_custom_axes", get_custom_axes)
+    u.atoms.select_atoms.side_effect = _select_atoms
+    u.atoms.principal_axes.return_value = np.eye(3)
+    monkeypatch.setattr(ax, "get_chain", backbone_atom)
     monkeypatch.setattr(
         ax,
-        "get_custom_moment_of_inertia",
-        lambda **kwargs: np.array([1.0, 2.0, 3.0]),
+        "get_custom_residue_moment_of_inertia",
+        lambda center_of_mass, positions, masses, custom_rot_axes, dimensions: np.array(
+            [1, 1, 1]
+        ),
     )
-    monkeypatch.setattr(ax, "get_flipped_axes", lambda ua, axes, com, dims: axes)
 
-    custom_axes, moi = ax.get_bonded_axes_from_topology(
-        u=universe,
-        heavy_atom=heavy_atom,
-        topology=topology,
-        dimensions=np.array([10.0, 10.0, 10.0]),
+    trans_axes, rot_axes, rot_center, moi = ax.get_residue_axes(
+        u, index=1, relative_index=0
     )
 
-    np.testing.assert_allclose(custom_axes, np.eye(3))
-    np.testing.assert_allclose(moi, np.array([1.0, 2.0, 3.0]))
-    np.testing.assert_allclose(
-        get_custom_axes.call_args.kwargs["c"], bonded_light.position
-    )
+    assert len(edge_atom_set) == 2
+    assert np.allclose(trans_axes, np.eye(3))
+    assert rot_axes.shape == (3, 3)
+    assert rot_center.shape == (3,)
+    assert np.allclose(moi, np.array([1, 1, 1]))
 
 
-def test_get_bonded_axes_from_topology_two_heavy_uses_heavy_positions_as_b_list(
-    monkeypatch,
-):
+def test_get_ua_axes_bonded_axes_2neighbours(monkeypatch):
     ax = AxesCalculator()
+    residue_group = MagicMock()
+    residue_group.__len__ = 3
+    residue = residue_group.residues[1]
 
-    heavy_atom = _FakeAtom(1, 12.0, [0.0, 0.0, 0.0])
-    bonded_heavy_0 = _FakeAtom(3, 12.0, [1.0, 0.0, 0.0])
-    bonded_heavy_1 = _FakeAtom(4, 12.0, [0.0, 1.0, 0.0])
-    universe = _FakeUniverse(
-        {
-            1: heavy_atom,
-            3: bonded_heavy_0,
-            4: bonded_heavy_1,
-        }
+    heavy_atoms = _FakeAtomGroup(
+        [
+            _atom(index=3, mass=12.0, pos=(1, 0, 0)),
+            _atom(index=5, mass=12.0, pos=(0, 1, 0)),
+            _atom(index=7, mass=12.0, pos=(0, 0, 1)),
+        ],
+        positions=np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=float),
     )
-    topology = _ua_topology(
-        heavy_atom_index=1,
-        ua_atom_indices=(1,),
-        ua_all_atom_indices=(1, 3, 4),
-        bonded_heavy_indices=(3, 4),
-        bonded_light_indices=(),
+
+    edge_atoms = _FakeAtomGroup(
+        [
+            _atom(index=3, mass=12.0, pos=(1, 0, 0)),
+            _atom(index=7, mass=12.0, pos=(0, 0, 1)),
+        ],
+        positions=np.array([[1, 0, 0], [0, 0, 1]], dtype=float),
     )
 
-    get_custom_axes = MagicMock(return_value=np.eye(3) * 4.0)
-    monkeypatch.setattr(ax, "get_custom_axes", get_custom_axes)
+    def _select_atoms(q):
+        if q == "mass 2 to 999":
+            # return heavy atoms group
+            return heavy_atoms
+        if q.startswith("resindex "):
+            return edge_atoms
+        if q.startswith("index "):
+            return [heavy_atoms[1]]
+
+    residue_group.select_atoms.side_effect = _select_atoms
+    residue.atoms.select_atoms.side_effect = _select_atoms
+    monkeypatch.setattr(ax, "get_chain", lambda residue, first, last: heavy_atoms[1])
     monkeypatch.setattr(
         ax,
-        "get_custom_moment_of_inertia",
-        lambda **kwargs: np.array([4.0, 5.0, 6.0]),
+        "get_bonded_axes",
+        lambda system, atom, dimensions: (np.eye(3), np.array([1.0, 2.0, 3.0])),
     )
-    monkeypatch.setattr(ax, "get_flipped_axes", lambda ua, axes, com, dims: axes)
 
-    custom_axes, moi = ax.get_bonded_axes_from_topology(
-        u=universe,
-        heavy_atom=heavy_atom,
-        topology=topology,
-        dimensions=np.array([10.0, 10.0, 10.0]),
+    trans_axes, rot_axes, rot_center, moi = ax.get_UA_axes(
+        data_container=residue_group, index=1, res_position=0
     )
 
-    np.testing.assert_allclose(custom_axes, np.eye(3) * 4.0)
-    np.testing.assert_allclose(moi, np.array([4.0, 5.0, 6.0]))
-    np.testing.assert_allclose(
-        get_custom_axes.call_args.kwargs["b_list"],
-        np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]),
-    )
-    np.testing.assert_allclose(
-        get_custom_axes.call_args.kwargs["c"],
-        bonded_heavy_1.position,
-    )
+    assert trans_axes.shape == (3, 3)
+    assert rot_axes.shape == (3, 3)
+    assert moi.shape == (3,)
+    assert rot_center.shape == (3,)
 
 
-def test_get_bonded_axes_from_topology_returns_none_when_custom_axes_none(
-    monkeypatch,
-):
+def test_get_chain(monkeypatch):
     ax = AxesCalculator()
+    residue = MagicMock()
+    residue.__len__ = 5
+    heavy_atoms = [
+        _atom(index=0, mass=12.0, pos=(1, 0, 0)),
+        _atom(index=1, mass=12.0, pos=(0, 1, 0)),
+        _atom(index=2, mass=12.0, pos=(0, 0, 1)),
+        _atom(index=3, mass=12.0, pos=(0, 1, 1)),
+        _atom(index=4, mass=12.0, pos=(1, 0, 1)),
+    ]
 
-    heavy_atom = _FakeAtom(1, 12.0, [0.0, 0.0, 0.0])
-    bonded_heavy = _FakeAtom(3, 12.0, [1.0, 0.0, 0.0])
-    universe = _FakeUniverse({1: heavy_atom, 3: bonded_heavy})
-    topology = _ua_topology(
-        heavy_atom_index=1,
-        ua_atom_indices=(1,),
-        ua_all_atom_indices=(1, 3),
-        bonded_heavy_indices=(3,),
-        bonded_light_indices=(),
-    )
-
-    get_custom_moi = MagicMock()
-    get_flipped = MagicMock()
-
-    monkeypatch.setattr(ax, "get_custom_axes", lambda **kwargs: None)
-    monkeypatch.setattr(ax, "get_custom_moment_of_inertia", get_custom_moi)
-    monkeypatch.setattr(ax, "get_flipped_axes", get_flipped)
-
-    custom_axes, moi = ax.get_bonded_axes_from_topology(
-        u=universe,
-        heavy_atom=heavy_atom,
-        topology=topology,
-        dimensions=np.array([10.0, 10.0, 10.0]),
-    )
-
-    assert custom_axes is None
-    assert moi is None
-    get_custom_moi.assert_not_called()
-    get_flipped.assert_not_called()
+    def _select_atoms(q):
+        if q.endswith("bonded index 0"):
+            return [heavy_atoms[1]]
+        elif q.endswith("bonded index 1"):
+            return [heavy_atoms[0], heavy_atoms[2]]
+        elif q.endswith("bonded index 2"):
+            return [heavy_atoms[1], heavy_atoms[3]]
+        elif q.endswith("bonded index 3"):
+            return [heavy_atoms[2], heavy_atoms[4]]
+        elif q.endswith("bonded index 4"):
+            return [heavy_atoms[3]]
+        elif q.endswith("not index 0"):
+            return heavy_atoms[1:]
+
+    residue.atoms.select_atoms.side_effect = _select_atoms
+    residue.atoms.indices = np.arange(5)
+    print(f"The indices in our residue: {residue.atoms.indices}")
+
+    chain = ax.get_chain(residue=residue, first=heavy_atoms[0], last=heavy_atoms[-1])
+
+    assert len(chain) == 3

From db4d8aeafed31c3fb210ce661903fd6757cecb7a Mon Sep 17 00:00:00 2001
From: ioanaapapa <ioanaalexandrapapa@gmail.com>
Date: Fri, 24 Jul 2026 14:00:09 +0100
Subject: [PATCH 2/5] added unit tests for terminal residues

---
 tests/unit/CodeEntropy/levels/test_axes.py | 98 ++++++++++++++++++++++
 1 file changed, 98 insertions(+)

diff --git a/tests/unit/CodeEntropy/levels/test_axes.py b/tests/unit/CodeEntropy/levels/test_axes.py
index 9e1236b..e118eaa 100644
--- a/tests/unit/CodeEntropy/levels/test_axes.py
+++ b/tests/unit/CodeEntropy/levels/test_axes.py
@@ -801,6 +801,53 @@ def _select_atoms(q):
     assert np.allclose(moi, np.array([1, 1, 1]))
 
 
+def test_get_residue_bonded_axes_terminal_resid(monkeypatch):
+    ax = AxesCalculator()
+    u = MagicMock()
+    u.dimensions = np.array([10.0, 10.0, 10.0, 90, 90, 90])
+    monkeypatch.setattr("CodeEntropy.levels.axes.make_whole", lambda _ag: None)
+    residue = u.select_atoms("resindex 0")
+    residue.__len__.return_value = 3
+    residue.atoms = _FakeAtomGroup(
+        [
+            _atom(index=0, mass=12.0, pos=[1, 0, 0]),
+            _atom(index=1, mass=12.0, pos=[0, 1, 0]),
+            _atom(index=2, mass=12.0, pos=[0, 0, 0]),
+        ]
+    )
+    edge_atom_set = _FakeAtomGroup(
+        [
+            _atom(index=2, mass=12.0, pos=[0, 0, 0]),
+        ]
+    )
+
+    def _select_atoms(q):
+        if q.endswith("(bonded resindex -1 or resindex 1)"):
+            return edge_atom_set
+
+    backbone_atom = residue.atoms[1]
+    u.atoms.principal_axes.return_value = np.eye(3)
+    u.atoms.select_atoms.side_effect = _select_atoms
+    monkeypatch.setattr(ax, "get_chain", backbone_atom)
+    monkeypatch.setattr(
+        ax,
+        "get_custom_residue_moment_of_inertia",
+        lambda center_of_mass, positions, masses, custom_rot_axes, dimensions: np.array(
+            [1, 1, 1]
+        ),
+    )
+
+    trans_axes, rot_axes, rot_center, moi = ax.get_residue_axes(
+        u, index=0, relative_index=0
+    )
+
+    assert len(edge_atom_set) == 1
+    assert np.allclose(trans_axes, np.eye(3))
+    assert rot_axes.shape == (3, 3)
+    assert rot_center.shape == (3,)
+    assert np.allclose(moi, np.array([1, 1, 1]))
+
+
 def test_get_ua_axes_bonded_axes_2neighbours(monkeypatch):
     ax = AxesCalculator()
     residue_group = MagicMock()
@@ -852,6 +899,57 @@ def _select_atoms(q):
     assert rot_center.shape == (3,)
 
 
+def test_get_ua_axes_bonded_axes_terminal_resid(monkeypatch):
+    ax = AxesCalculator()
+    residue_group = MagicMock()
+    residue_group.__len__ = 2
+    residue = residue_group.residues[1]
+    heavy_atoms = _FakeAtomGroup(
+        [
+            _atom(index=0, mass=12.0, pos=(1, 0, 0)),
+            _atom(index=1, mass=12.0, pos=(0, 1, 0)),
+            _atom(index=2, mass=12.0, pos=(0, 0, 1)),
+        ],
+    )
+
+    edge_atom_set = _FakeAtomGroup(
+        [
+            _atom(index=0, mass=12.0, pos=(1, 0, 0)),
+        ],
+    )
+
+    def _select_atoms(q):
+        if q == "mass 2 to 999":
+            # return heavy atoms group
+            return heavy_atoms
+        if q.startswith("resindex "):
+            return edge_atom_set
+        if q.startswith("index "):
+            return [heavy_atoms[0]]
+        if q == ("(mass 2 to 999) and bonded index 2"):
+            return [heavy_atoms[1]]
+        if q == ("(mass 2 to 999) and bonded index 1"):
+            return [heavy_atoms[0], heavy_atoms[2]]
+
+    residue_group.select_atoms.side_effect = _select_atoms
+    residue.atoms.select_atoms.side_effect = _select_atoms
+    edge_atom_set.atoms = [edge_atom_set[0]]
+    monkeypatch.setattr(ax, "get_chain", lambda residue, first, last: heavy_atoms[1])
+    monkeypatch.setattr(
+        ax,
+        "get_bonded_axes",
+        lambda system, atom, dimensions: (np.eye(3), np.array([1.0, 1.0, 1.0])),
+    )
+    trans_axes, rot_axes, rot_center, moi = ax.get_UA_axes(
+        data_container=residue_group, index=0, res_position=1
+    )
+
+    assert trans_axes.shape == (3, 3)
+    assert rot_axes.shape == (3, 3)
+    assert moi.shape == (3,)
+    assert rot_center.shape == (3,)
+
+
 def test_get_chain(monkeypatch):
     ax = AxesCalculator()
     residue = MagicMock()

From e215e6a1527547f0cd91b4873954f33b735c27ad Mon Sep 17 00:00:00 2001
From: ioanaapapa <ioanaalexandrapapa@gmail.com>
Date: Fri, 24 Jul 2026 14:10:02 +0100
Subject: [PATCH 3/5] updated regression baselines

---
 .../combined_forcetorque_false.json           |  4 +-
 .../baselines/benzaldehyde/frame_window.json  |  4 +-
 .../benzaldehyde/selection_subset.json        |  4 +-
 .../benzene/combined_forcetorque_off.json     |  4 +-
 .../baselines/benzene/frame_window.json       |  4 +-
 .../baselines/benzene/selection_subset.json   |  4 +-
 .../cyclohexane/combined_forcetorque_off.json |  4 +-
 .../baselines/cyclohexane/frame_window.json   |  4 +-
 .../cyclohexane/selection_subset.json         |  4 +-
 .../dna/combined_forcetorque_off.json         | 10 ++--
 .../baselines/dna/frame_window.json           | 10 ++--
 .../baselines/dna/grouping_each.json          | 10 ++--
 .../baselines/dna/selection_subset.json       | 10 ++--
 .../combined_forcetorque_off.json             |  4 +-
 .../baselines/ethyl-acetate/frame_window.json |  4 +-
 .../ethyl-acetate/selection_subset.json       |  4 +-
 .../methane/combined_forcetorque_off.json     |  6 +--
 .../baselines/methane/frame_window.json       |  4 +-
 .../baselines/methane/grouping_each.json      | 44 +++++++++---------
 .../baselines/methane/selection_subset.json   |  6 +--
 .../octonol/combined_forcetorque_off.json     |  4 +-
 .../baselines/octonol/frame_window.json       |  4 +-
 .../baselines/octonol/selection_subset.json   |  4 +-
 .../water/combined_forcetorque_off.json       |  6 +--
 .../baselines/water/frame_window.json         |  4 +-
 .../baselines/water/grouping_each.json        | 46 +++++++++----------
 .../baselines/water/selection_subset.json     |  6 +--
 .../regression/baselines/water/water_off.json |  6 +--
 28 files changed, 114 insertions(+), 114 deletions(-)

diff --git a/tests/regression/baselines/benzaldehyde/combined_forcetorque_false.json b/tests/regression/baselines/benzaldehyde/combined_forcetorque_false.json
index a5d1af0..27025e6 100644
--- a/tests/regression/baselines/benzaldehyde/combined_forcetorque_false.json
+++ b/tests/regression/baselines/benzaldehyde/combined_forcetorque_false.json
@@ -2,7 +2,7 @@
   "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,
@@ -10,7 +10,7 @@
         "residue:Conformational": 0.0,
         "residue:Orientational": 20.481571492615355
       },
-      "total": 208.1878625488175
+      "total": 208.2064989406355
     }
   }
 }
diff --git a/tests/regression/baselines/benzaldehyde/frame_window.json b/tests/regression/baselines/benzaldehyde/frame_window.json
index 70bb276..2722572 100644
--- a/tests/regression/baselines/benzaldehyde/frame_window.json
+++ b/tests/regression/baselines/benzaldehyde/frame_window.json
@@ -2,7 +2,7 @@
   "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,
@@ -10,7 +10,7 @@
         "residue:Conformational": 0.0,
         "residue:Orientational": 20.481571492615355
       },
-      "total": 230.29580627384846
+      "total": 231.74961624553984
     }
   }
 }
diff --git a/tests/regression/baselines/benzaldehyde/selection_subset.json b/tests/regression/baselines/benzaldehyde/selection_subset.json
index 6d50b1a..c7750eb 100644
--- a/tests/regression/baselines/benzaldehyde/selection_subset.json
+++ b/tests/regression/baselines/benzaldehyde/selection_subset.json
@@ -2,7 +2,7 @@
   "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,
@@ -10,7 +10,7 @@
         "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
index 16b5e37..ee2c41f 100644
--- a/tests/regression/baselines/benzene/combined_forcetorque_off.json
+++ b/tests/regression/baselines/benzene/combined_forcetorque_off.json
@@ -2,7 +2,7 @@
   "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,
@@ -10,7 +10,7 @@
         "residue:Conformational": 0.0,
         "residue:Orientational": 0.0
       },
-      "total": 151.31064169774515
+      "total": 151.325768040128
     }
   }
 }
diff --git a/tests/regression/baselines/benzene/frame_window.json b/tests/regression/baselines/benzene/frame_window.json
index c1ffd00..0f692cb 100644
--- a/tests/regression/baselines/benzene/frame_window.json
+++ b/tests/regression/baselines/benzene/frame_window.json
@@ -2,7 +2,7 @@
   "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,
@@ -10,7 +10,7 @@
         "residue:Conformational": 0.0,
         "residue:Orientational": 0.0
       },
-      "total": 173.89519974256632
+      "total": 173.9825127991342
     }
   }
 }
diff --git a/tests/regression/baselines/benzene/selection_subset.json b/tests/regression/baselines/benzene/selection_subset.json
index ff2d09b..807d8c0 100644
--- a/tests/regression/baselines/benzene/selection_subset.json
+++ b/tests/regression/baselines/benzene/selection_subset.json
@@ -2,7 +2,7 @@
   "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,
@@ -10,7 +10,7 @@
         "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
index a548349..56de291 100644
--- a/tests/regression/baselines/cyclohexane/combined_forcetorque_off.json
+++ b/tests/regression/baselines/cyclohexane/combined_forcetorque_off.json
@@ -2,7 +2,7 @@
   "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,
@@ -10,7 +10,7 @@
         "residue:Conformational": 0.0,
         "residue:Orientational": 0.0
       },
-      "total": 163.50155250432303
+      "total": 163.5892068042079
     }
   }
 }
diff --git a/tests/regression/baselines/cyclohexane/frame_window.json b/tests/regression/baselines/cyclohexane/frame_window.json
index d21d437..dd42326 100644
--- a/tests/regression/baselines/cyclohexane/frame_window.json
+++ b/tests/regression/baselines/cyclohexane/frame_window.json
@@ -2,7 +2,7 @@
   "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,
@@ -10,7 +10,7 @@
         "residue:Conformational": 0.0,
         "residue:Orientational": 0.0
       },
-      "total": 231.09780323388605
+      "total": 230.3046947043258
     }
   }
 }
diff --git a/tests/regression/baselines/cyclohexane/selection_subset.json b/tests/regression/baselines/cyclohexane/selection_subset.json
index b1a1444..a1a3714 100644
--- a/tests/regression/baselines/cyclohexane/selection_subset.json
+++ b/tests/regression/baselines/cyclohexane/selection_subset.json
@@ -2,7 +2,7 @@
   "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,
@@ -10,7 +10,7 @@
         "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
index e867cee..2b2d5bf 100644
--- a/tests/regression/baselines/dna/combined_forcetorque_off.json
+++ b/tests/regression/baselines/dna/combined_forcetorque_off.json
@@ -2,31 +2,31 @@
   "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": {
         "united_atom:Transvibrational": 0.0,
         "united_atom:Rovibrational": 0.01846427765949586,
         "residue:Transvibrational": 0.0,
-        "residue:Rovibrational": 2.3863201082544565,
+        "residue:Rovibrational": 5.702835600675442,
         "polymer:Transvibrational": 16.607667396609116,
         "polymer:Rovibrational": 12.304363914795593,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "polymer:Orientational": 4.758905336627712
       },
-      "total": 36.07572103394637
+      "total": 39.39223652636736
     }
   }
 }
diff --git a/tests/regression/baselines/dna/frame_window.json b/tests/regression/baselines/dna/frame_window.json
index c35d221..b561ec3 100644
--- a/tests/regression/baselines/dna/frame_window.json
+++ b/tests/regression/baselines/dna/frame_window.json
@@ -5,28 +5,28 @@
         "united_atom:Transvibrational": 0.0,
         "united_atom:Rovibrational": 1.5821720528374943,
         "residue:Transvibrational": 0.0,
-        "residue:Rovibrational": 27.397449238560412,
+        "residue:Rovibrational": 35.89784662172834,
         "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.44373670937406
     },
     "1": {
       "components": {
-        "united_atom:Transvibrational": 0.0,
+        "united_atom:Transvibrational": 5.359713875687138e-10,
         "united_atom:Rovibrational": 2.5277936366208014,
         "residue:Transvibrational": 0.0,
-        "residue:Rovibrational": 24.80670067454149,
+        "residue:Rovibrational": 33.62040396388513,
         "polymer:FTmat-Transvibrational": 60.47397935339153,
         "polymer:FTmat-Rovibrational": 0.0,
         "united_atom:Conformational": 5.292271495278918,
         "residue:Conformational": 0.0,
         "polymer:Orientational": 4.758905336627712
       },
-      "total": 97.85965049646045
+      "total": 106.67335378634006
     }
   }
 }
diff --git a/tests/regression/baselines/dna/grouping_each.json b/tests/regression/baselines/dna/grouping_each.json
index 2cbef74..9ab18f0 100644
--- a/tests/regression/baselines/dna/grouping_each.json
+++ b/tests/regression/baselines/dna/grouping_each.json
@@ -2,31 +2,31 @@
   "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: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.735770036087217
     },
     "1": {
       "components": {
         "united_atom:Transvibrational": 0.0,
         "united_atom:Rovibrational": 0.01846427765949586,
         "residue:Transvibrational": 0.0,
-        "residue:Rovibrational": 2.3863201082544565,
+        "residue:Rovibrational": 5.702835600675442,
         "polymer:FTmat-Transvibrational": 11.11037253388596,
         "polymer:FTmat-Rovibrational": 0.0,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "polymer:Orientational": 4.758905336627712
       },
-      "total": 18.274062256427623
+      "total": 21.59057774884861
     }
   }
 }
diff --git a/tests/regression/baselines/dna/selection_subset.json b/tests/regression/baselines/dna/selection_subset.json
index 2cbef74..9ab18f0 100644
--- a/tests/regression/baselines/dna/selection_subset.json
+++ b/tests/regression/baselines/dna/selection_subset.json
@@ -2,31 +2,31 @@
   "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: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.735770036087217
     },
     "1": {
       "components": {
         "united_atom:Transvibrational": 0.0,
         "united_atom:Rovibrational": 0.01846427765949586,
         "residue:Transvibrational": 0.0,
-        "residue:Rovibrational": 2.3863201082544565,
+        "residue:Rovibrational": 5.702835600675442,
         "polymer:FTmat-Transvibrational": 11.11037253388596,
         "polymer:FTmat-Rovibrational": 0.0,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "polymer:Orientational": 4.758905336627712
       },
-      "total": 18.274062256427623
+      "total": 21.59057774884861
     }
   }
 }
diff --git a/tests/regression/baselines/ethyl-acetate/combined_forcetorque_off.json b/tests/regression/baselines/ethyl-acetate/combined_forcetorque_off.json
index 8c68cce..499fca4 100644
--- a/tests/regression/baselines/ethyl-acetate/combined_forcetorque_off.json
+++ b/tests/regression/baselines/ethyl-acetate/combined_forcetorque_off.json
@@ -2,7 +2,7 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 1.196367074472673,
+        "united_atom:Transvibrational": 0.8081103094129065,
         "united_atom:Rovibrational": 82.36455154278131,
         "residue:Transvibrational": 64.68750154134017,
         "residue:Rovibrational": 58.903938937880085,
@@ -10,7 +10,7 @@
         "residue:Conformational": 0.0,
         "residue:Orientational": 0.0
       },
-      "total": 215.05334534012698
+      "total": 214.6650885750672
     }
   }
 }
diff --git a/tests/regression/baselines/ethyl-acetate/frame_window.json b/tests/regression/baselines/ethyl-acetate/frame_window.json
index be566d0..911a266 100644
--- a/tests/regression/baselines/ethyl-acetate/frame_window.json
+++ b/tests/regression/baselines/ethyl-acetate/frame_window.json
@@ -2,7 +2,7 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 28.312711247966337,
+        "united_atom:Transvibrational": 29.222699118677973,
         "united_atom:Rovibrational": 51.99036142818903,
         "residue:FTmat-Transvibrational": 67.80560626717748,
         "residue:FTmat-Rovibrational": 55.1631201009248,
@@ -10,7 +10,7 @@
         "residue:Conformational": 0.0,
         "residue:Orientational": 3.8984476469128193
       },
-      "total": 215.80571814986257
+      "total": 216.7157060205742
     }
   }
 }
diff --git a/tests/regression/baselines/ethyl-acetate/selection_subset.json b/tests/regression/baselines/ethyl-acetate/selection_subset.json
index 270dbbd..dd0f269 100644
--- a/tests/regression/baselines/ethyl-acetate/selection_subset.json
+++ b/tests/regression/baselines/ethyl-acetate/selection_subset.json
@@ -2,7 +2,7 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 1.196367074472673,
+        "united_atom:Transvibrational": 0.8081103094129065,
         "united_atom:Rovibrational": 82.36455154278131,
         "residue:FTmat-Transvibrational": 74.5939232239247,
         "residue:FTmat-Rovibrational": 55.42964192531005,
@@ -10,7 +10,7 @@
         "residue:Conformational": 0.0,
         "residue:Orientational": 0.0
       },
-      "total": 221.48547001014148
+      "total": 221.09721324508172
     }
   }
 }
diff --git a/tests/regression/baselines/methane/combined_forcetorque_off.json b/tests/regression/baselines/methane/combined_forcetorque_off.json
index 318c30c..70be3ce 100644
--- a/tests/regression/baselines/methane/combined_forcetorque_off.json
+++ b/tests/regression/baselines/methane/combined_forcetorque_off.json
@@ -2,12 +2,12 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 42.51619525142699,
-        "united_atom:Rovibrational": 34.651932322640015,
+        "united_atom:Transvibrational": 42.516195251426986,
+        "united_atom:Rovibrational": 34.69255616696802,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 77.168127574067
+      "total": 77.208751418395
     }
   }
 }
diff --git a/tests/regression/baselines/methane/frame_window.json b/tests/regression/baselines/methane/frame_window.json
index 9a45dc9..0f70346 100644
--- a/tests/regression/baselines/methane/frame_window.json
+++ b/tests/regression/baselines/methane/frame_window.json
@@ -3,11 +3,11 @@
     "0": {
       "components": {
         "united_atom:Transvibrational": 40.70376001258969,
-        "united_atom:Rovibrational": 33.778792396823484,
+        "united_atom:Rovibrational": 33.83818524319629,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 74.48255240941317
+      "total": 74.54194525578598
     }
   }
 }
diff --git a/tests/regression/baselines/methane/grouping_each.json b/tests/regression/baselines/methane/grouping_each.json
index 3b31963..e5f74df 100644
--- a/tests/regression/baselines/methane/grouping_each.json
+++ b/tests/regression/baselines/methane/grouping_each.json
@@ -3,92 +3,92 @@
     "0": {
       "components": {
         "united_atom:Transvibrational": 24.315699188266926,
-        "united_atom:Rovibrational": 21.0483901242056,
+        "united_atom:Rovibrational": 21.06129596642336,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 45.36408931247253
+      "total": 45.376995154690285
     },
     "1": {
       "components": {
         "united_atom:Transvibrational": 23.76608267046309,
-        "united_atom:Rovibrational": 16.769636839671296,
+        "united_atom:Rovibrational": 16.800440306419098,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 40.535719510134385
+      "total": 40.56652297688218
     },
     "2": {
       "components": {
         "united_atom:Transvibrational": 24.833896219788347,
-        "united_atom:Rovibrational": 16.996392060962357,
+        "united_atom:Rovibrational": 17.0032604785275,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 41.83028828075071
+      "total": 41.83715669831585
     },
     "3": {
       "components": {
         "united_atom:Transvibrational": 6.6723857079283215,
-        "united_atom:Rovibrational": 4.986376282798274,
+        "united_atom:Rovibrational": 4.958640584432368,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 11.658761990726596
+      "total": 11.63102629236069
     },
     "4": {
       "components": {
-        "united_atom:Transvibrational": 6.880461178899526,
-        "united_atom:Rovibrational": 4.441585994100298,
+        "united_atom:Transvibrational": 6.880461178899524,
+        "united_atom:Rovibrational": 4.47005523879151,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 11.322047172999824
+      "total": 11.350516417691033
     },
     "5": {
       "components": {
         "united_atom:Transvibrational": 15.473092273577388,
-        "united_atom:Rovibrational": 7.777083478334702,
+        "united_atom:Rovibrational": 7.7776511992146515,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 23.25017575191209
+      "total": 23.25074347279204
     },
     "6": {
       "components": {
         "united_atom:Transvibrational": 25.694168010993888,
-        "united_atom:Rovibrational": 6.731392483940261,
+        "united_atom:Rovibrational": 6.727950082765358,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 32.425560494934146
+      "total": 32.422118093759245
     },
     "7": {
       "components": {
-        "united_atom:Transvibrational": 10.533257828925997,
-        "united_atom:Rovibrational": 8.97563913301586,
+        "united_atom:Transvibrational": 10.533257828925999,
+        "united_atom:Rovibrational": 9.059251947808503,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 19.508896961941858
+      "total": 19.592509776734502
     },
     "8": {
       "components": {
         "united_atom:Transvibrational": 7.332978272264879,
-        "united_atom:Rovibrational": 3.0418502362906086,
+        "united_atom:Rovibrational": 3.084945359076361,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 10.374828508555488
+      "total": 10.417923631341239
     },
     "9": {
       "components": {
         "united_atom:Transvibrational": 4.023932761002867,
-        "united_atom:Rovibrational": 8.289324383636089,
+        "united_atom:Rovibrational": 8.271395156034108,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 12.313257144638957
+      "total": 12.295327917036975
     }
   }
 }
diff --git a/tests/regression/baselines/methane/selection_subset.json b/tests/regression/baselines/methane/selection_subset.json
index 318c30c..70be3ce 100644
--- a/tests/regression/baselines/methane/selection_subset.json
+++ b/tests/regression/baselines/methane/selection_subset.json
@@ -2,12 +2,12 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 42.51619525142699,
-        "united_atom:Rovibrational": 34.651932322640015,
+        "united_atom:Transvibrational": 42.516195251426986,
+        "united_atom:Rovibrational": 34.69255616696802,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 77.168127574067
+      "total": 77.208751418395
     }
   }
 }
diff --git a/tests/regression/baselines/octonol/combined_forcetorque_off.json b/tests/regression/baselines/octonol/combined_forcetorque_off.json
index 920589a..362cd1c 100644
--- a/tests/regression/baselines/octonol/combined_forcetorque_off.json
+++ b/tests/regression/baselines/octonol/combined_forcetorque_off.json
@@ -2,7 +2,7 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 0.34382967782781193,
+        "united_atom:Transvibrational": 0.2907129992470817,
         "united_atom:Rovibrational": 19.154228264604278,
         "residue:Transvibrational": 74.72816183790984,
         "residue:Rovibrational": 60.84390144550801,
@@ -10,7 +10,7 @@
         "residue:Conformational": 0.0,
         "residue:Orientational": 25.79114991860739
       },
-      "total": 197.70076468713353
+      "total": 197.64764800855278
     }
   }
 }
diff --git a/tests/regression/baselines/octonol/frame_window.json b/tests/regression/baselines/octonol/frame_window.json
index 5c07dcc..3fcdaef 100644
--- a/tests/regression/baselines/octonol/frame_window.json
+++ b/tests/regression/baselines/octonol/frame_window.json
@@ -2,7 +2,7 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 65.86600728016141,
+        "united_atom:Transvibrational": 68.20033834493445,
         "united_atom:Rovibrational": 162.26463085986236,
         "residue:FTmat-Transvibrational": 74.84326467422125,
         "residue:FTmat-Rovibrational": 60.95710412746361,
@@ -10,7 +10,7 @@
         "residue:Conformational": 0.0,
         "residue:Orientational": 25.79114991860739
       },
-      "total": 406.37330124036134
+      "total": 408.7076323051344
     }
   }
 }
diff --git a/tests/regression/baselines/octonol/selection_subset.json b/tests/regression/baselines/octonol/selection_subset.json
index 5f9d18f..e15b955 100644
--- a/tests/regression/baselines/octonol/selection_subset.json
+++ b/tests/regression/baselines/octonol/selection_subset.json
@@ -2,7 +2,7 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 0.34382967782781193,
+        "united_atom:Transvibrational": 0.2907129992470817,
         "united_atom:Rovibrational": 19.154228264604278,
         "residue:FTmat-Transvibrational": 83.17889385114952,
         "residue:FTmat-Rovibrational": 57.480841206687685,
@@ -10,7 +10,7 @@
         "residue:Conformational": 0.0,
         "residue:Orientational": 25.79114991860739
       },
-      "total": 202.78843646155286
+      "total": 202.73531978297217
     }
   }
 }
diff --git a/tests/regression/baselines/water/combined_forcetorque_off.json b/tests/regression/baselines/water/combined_forcetorque_off.json
index 2e63256..b92ef4e 100644
--- a/tests/regression/baselines/water/combined_forcetorque_off.json
+++ b/tests/regression/baselines/water/combined_forcetorque_off.json
@@ -2,12 +2,12 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 51.22145722965398,
-        "united_atom:Rovibrational": 17.232634718413053,
+        "united_atom:Transvibrational": 51.22145722965397,
+        "united_atom:Rovibrational": 17.307840626636313,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 68.45409194806703
+      "total": 68.52929785629028
     }
   }
 }
diff --git a/tests/regression/baselines/water/frame_window.json b/tests/regression/baselines/water/frame_window.json
index c6ae550..8b36802 100644
--- a/tests/regression/baselines/water/frame_window.json
+++ b/tests/regression/baselines/water/frame_window.json
@@ -3,11 +3,11 @@
     "0": {
       "components": {
         "united_atom:Transvibrational": 46.684099714230925,
-        "united_atom:Rovibrational": 17.694014405444744,
+        "united_atom:Rovibrational": 17.66359738549709,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 64.37811411967567
+      "total": 64.34769709972801
     }
   }
 }
diff --git a/tests/regression/baselines/water/grouping_each.json b/tests/regression/baselines/water/grouping_each.json
index acccd86..676010d 100644
--- a/tests/regression/baselines/water/grouping_each.json
+++ b/tests/regression/baselines/water/grouping_each.json
@@ -3,92 +3,92 @@
     "0": {
       "components": {
         "united_atom:Transvibrational": 14.945907953805552,
-        "united_atom:Rovibrational": 1.8947243953648254,
+        "united_atom:Rovibrational": 1.94076792014105,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 16.840632349170377
+      "total": 16.8866758739466
     },
     "1": {
       "components": {
-        "united_atom:Transvibrational": 8.111522208287779,
-        "united_atom:Rovibrational": 0.9607747968334215,
+        "united_atom:Transvibrational": 8.11152220828778,
+        "united_atom:Rovibrational": 0.9798912337680555,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 9.0722970051212
+      "total": 9.091413442055837
     },
     "2": {
       "components": {
-        "united_atom:Transvibrational": 15.79497347968759,
-        "united_atom:Rovibrational": 4.606072530590255,
+        "united_atom:Transvibrational": 15.794973479687599,
+        "united_atom:Rovibrational": 4.4692860930919585,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 20.401046010277845
+      "total": 20.264259572779558
     },
     "3": {
       "components": {
         "united_atom:Transvibrational": 9.462487265119448,
-        "united_atom:Rovibrational": 1.6509513986984767,
+        "united_atom:Rovibrational": 1.6054615846390614,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 11.113438663817924
+      "total": 11.067948849758508
     },
     "4": {
       "components": {
         "united_atom:Transvibrational": 16.254161245740693,
-        "united_atom:Rovibrational": 1.367597374110767,
+        "united_atom:Rovibrational": 1.448700454979353,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 17.62175861985146
+      "total": 17.702861700720046
     },
     "5": {
       "components": {
         "united_atom:Transvibrational": 11.57779876203329,
-        "united_atom:Rovibrational": 7.351163568004953,
+        "united_atom:Rovibrational": 7.23849459157533,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 18.928962330038242
+      "total": 18.81629335360862
     },
     "6": {
       "components": {
         "united_atom:Transvibrational": 10.172061262201801,
-        "united_atom:Rovibrational": 3.671537296888502,
+        "united_atom:Rovibrational": 4.086982297438225,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 13.843598559090303
+      "total": 14.259043559640027
     },
     "7": {
       "components": {
-        "united_atom:Transvibrational": 13.49951045186194,
-        "united_atom:Rovibrational": 1.7029491215274863,
+        "united_atom:Transvibrational": 13.499510451861939,
+        "united_atom:Rovibrational": 1.6178142656646386,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 15.202459573389428
+      "total": 15.117324717526577
     },
     "8": {
       "components": {
         "united_atom:Transvibrational": 16.597759171924583,
-        "united_atom:Rovibrational": 1.6190289365617598,
+        "united_atom:Rovibrational": 1.6676178150275727,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 18.21678810848634
+      "total": 18.265376986952155
     },
     "9": {
       "components": {
         "united_atom:Transvibrational": 13.731375187064538,
-        "united_atom:Rovibrational": 5.650620990305761,
+        "united_atom:Rovibrational": 5.161816204002109,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 19.381996177370297
+      "total": 18.893191391066647
     }
   }
 }
diff --git a/tests/regression/baselines/water/selection_subset.json b/tests/regression/baselines/water/selection_subset.json
index 2e63256..b92ef4e 100644
--- a/tests/regression/baselines/water/selection_subset.json
+++ b/tests/regression/baselines/water/selection_subset.json
@@ -2,12 +2,12 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 51.22145722965398,
-        "united_atom:Rovibrational": 17.232634718413053,
+        "united_atom:Transvibrational": 51.22145722965397,
+        "united_atom:Rovibrational": 17.307840626636313,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 68.45409194806703
+      "total": 68.52929785629028
     }
   }
 }
diff --git a/tests/regression/baselines/water/water_off.json b/tests/regression/baselines/water/water_off.json
index 2e63256..b92ef4e 100644
--- a/tests/regression/baselines/water/water_off.json
+++ b/tests/regression/baselines/water/water_off.json
@@ -2,12 +2,12 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 51.22145722965398,
-        "united_atom:Rovibrational": 17.232634718413053,
+        "united_atom:Transvibrational": 51.22145722965397,
+        "united_atom:Rovibrational": 17.307840626636313,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 0.0
       },
-      "total": 68.45409194806703
+      "total": 68.52929785629028
     }
   }
 }

From c75b1e9205a4a006c6df571b1c310a737e85dea1 Mon Sep 17 00:00:00 2001
From: ioanaapapa <ioanaalexandrapapa@gmail.com>
Date: Fri, 24 Jul 2026 15:04:09 +0100
Subject: [PATCH 4/5] test get_residue_axes last resid test

---
 tests/unit/CodeEntropy/levels/test_axes.py | 61 ++++++++++++++++++++--
 1 file changed, 58 insertions(+), 3 deletions(-)

diff --git a/tests/unit/CodeEntropy/levels/test_axes.py b/tests/unit/CodeEntropy/levels/test_axes.py
index e118eaa..1200cf1 100644
--- a/tests/unit/CodeEntropy/levels/test_axes.py
+++ b/tests/unit/CodeEntropy/levels/test_axes.py
@@ -801,7 +801,7 @@ def _select_atoms(q):
     assert np.allclose(moi, np.array([1, 1, 1]))
 
 
-def test_get_residue_bonded_axes_terminal_resid(monkeypatch):
+def test_get_residue_bonded_axes_first_resid(monkeypatch):
     ax = AxesCalculator()
     u = MagicMock()
     u.dimensions = np.array([10.0, 10.0, 10.0, 90, 90, 90])
@@ -848,6 +848,62 @@ def _select_atoms(q):
     assert np.allclose(moi, np.array([1, 1, 1]))
 
 
+def test_get_residue_bonded_axes_last_resid(monkeypatch):
+    ax = AxesCalculator()
+    u = MagicMock()
+    u.dimensions = np.array([10.0, 10.0, 10.0, 90, 90, 90])
+    monkeypatch.setattr("CodeEntropy.levels.axes.make_whole", lambda _ag: None)
+    residue = u.select_atoms("resindex 2")
+    residue.__len__.return_value = 3
+    heavy_atoms = _FakeAtomGroup(
+        [
+            _atom(index=4, mass=12.0, pos=[1, 0, 0]),
+            _atom(index=5, mass=12.0, pos=[0, 1, 0]),
+            _atom(index=6, mass=12.0, pos=[0, 0, 0]),
+        ]
+    )
+    edge_atom_set = _FakeAtomGroup(
+        [
+            _atom(index=4, mass=12.0, pos=[0, 0, 0]),
+        ]
+    )
+
+    def _select_atoms(q):
+        if q == "mass 2 to 999":
+            # return heavy atoms group
+            return heavy_atoms
+        if q.endswith("(bonded resindex 1 or resindex 3)"):
+            return edge_atom_set
+        if q == ("(mass 2 to 999) and bonded index 6"):
+            return [heavy_atoms[1]]
+        if q == ("(mass 2 to 999) and bonded index 5"):
+            return [heavy_atoms[0], heavy_atoms[2]]
+
+    backbone_atom = heavy_atoms[1]
+    u.atoms.principal_axes.return_value = np.eye(3)
+    u.atoms.select_atoms.side_effect = _select_atoms
+    residue.select_atoms.side_effect = _select_atoms
+    residue.atoms.select_atoms.side_effect = _select_atoms
+    monkeypatch.setattr(ax, "get_chain", backbone_atom)
+    monkeypatch.setattr(
+        ax,
+        "get_custom_residue_moment_of_inertia",
+        lambda center_of_mass, positions, masses, custom_rot_axes, dimensions: np.array(
+            [1, 1, 1]
+        ),
+    )
+
+    trans_axes, rot_axes, rot_center, moi = ax.get_residue_axes(
+        u, index=2, relative_index=0
+    )
+
+    assert len(edge_atom_set) == 1
+    assert np.allclose(trans_axes, np.eye(3))
+    assert rot_axes.shape == (3, 3)
+    assert rot_center.shape == (3,)
+    assert np.allclose(moi, np.array([1, 1, 1]))
+
+
 def test_get_ua_axes_bonded_axes_2neighbours(monkeypatch):
     ax = AxesCalculator()
     residue_group = MagicMock()
@@ -899,7 +955,7 @@ def _select_atoms(q):
     assert rot_center.shape == (3,)
 
 
-def test_get_ua_axes_bonded_axes_terminal_resid(monkeypatch):
+def test_get_ua_axes_bonded_axes_last_resid(monkeypatch):
     ax = AxesCalculator()
     residue_group = MagicMock()
     residue_group.__len__ = 2
@@ -978,7 +1034,6 @@ def _select_atoms(q):
 
     residue.atoms.select_atoms.side_effect = _select_atoms
     residue.atoms.indices = np.arange(5)
-    print(f"The indices in our residue: {residue.atoms.indices}")
 
     chain = ax.get_chain(residue=residue, first=heavy_atoms[0], last=heavy_atoms[-1])
 

From 356d5abb61446b4d51c1c8ea6f3889c571cf7893 Mon Sep 17 00:00:00 2001
From: ioanaapapa <ioanaalexandrapapa@gmail.com>
Date: Fri, 24 Jul 2026 15:10:49 +0100
Subject: [PATCH 5/5] updated documentation

---
 docs/science.rst | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/science.rst b/docs/science.rst
index 1a98f97..7aa219d 100644
--- a/docs/science.rst
+++ b/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