diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py
index 4d4cf89c..8bec335d 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 39ff61a3..a2c73aa1 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/docs/science.rst b/docs/science.rst
index 1a98f977..7aa219db 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
diff --git a/tests/regression/baselines/benzaldehyde/combined_forcetorque_false.json b/tests/regression/baselines/benzaldehyde/combined_forcetorque_false.json
index a5d1af0e..27025e62 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 70bb2766..27225723 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 6d50b1a7..c7750ebb 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 16b5e375..ee2c41fc 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 c1ffd000..0f692cbb 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 ff2d09b5..807d8c05 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 a548349d..56de2915 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 d21d437d..dd423267 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 b1a14443..a1a37145 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 e867cee1..2b2d5bf9 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 c35d2211..b561ec3e 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 2cbef740..9ab18f0c 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 2cbef740..9ab18f0c 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 8c68cce7..499fca4c 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 be566d07..911a2667 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 270dbbd5..dd0f2695 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 318c30c6..70be3cef 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 9a45dc98..0f70346d 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 3b31963b..e5f74df8 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 318c30c6..70be3cef 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 920589a7..362cd1c2 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 5c07dcc2..3fcdaef3 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 5f9d18f1..e15b9555 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 2e632568..b92ef4e3 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 c6ae550d..8b368026 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 acccd86b..676010dd 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 2e632568..b92ef4e3 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 2e632568..b92ef4e3 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
     }
   }
 }
diff --git a/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py b/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py
index 9d66e88f..da0cb076 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 e68f4d1d..1200cf13 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,319 @@ 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()
+    assert rot_center.shape == (3,)
+    assert rot_axes.shape == (3, 3)
 
 
-def test_get_residue_axes_from_topology_neighbor_bonds_uses_vanilla_axes(
-    monkeypatch,
-):
+def test_get_custom_residue_moment_of_inertia(monkeypatch):
     ax = AxesCalculator()
-
-    universe = _FakeUniverse(
-        {},
-        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),
     )
-    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,
+    dimensions = np.array([10.0, 10.0, 10.0], dtype=float)
+
+    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(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]))
+    assert moi.shape == (3,)
 
 
-def test_get_UA_axes_from_topology_multiple_heavy_uses_cached_indices_and_box(
-    monkeypatch,
-):
+def test_get_residue_bonded_axes_2neighbours(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 1")
+    residue.__len__.return_value = 1
 
-    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),
+    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])
 
-    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)
+    def _select_atoms(q):
+        if q.endswith("(bonded resindex 0 or resindex 2)"):
+            return edge_atom_set
 
-    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,
+    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_residue_moment_of_inertia",
+        lambda center_of_mass, positions, masses, custom_rot_axes, dimensions: np.array(
+            [1, 1, 1]
+        ),
     )
 
-    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,
+    trans_axes, rot_axes, rot_center, moi = ax.get_residue_axes(
+        u, index=1, relative_index=0
     )
 
+    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_UA_axes_from_topology_single_heavy_uses_residue_principal_axes(
-    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]
-    )
-    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)
 
-    trans_axes, rot_axes, center, moi = ax.get_UA_axes_from_topology(
-        u=universe,
-        residue_atoms=residue_atoms,
-        topology=topology,
-        box=None,
+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])
+    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]),
+        ]
     )
-
-    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])
+    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
 
-def test_get_UA_axes_from_topology_raises_when_cached_bonded_axes_fail(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,))
-
-    monkeypatch.setattr("CodeEntropy.levels.axes.make_whole", lambda _ag: None)
+    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_bonded_axes_from_topology", lambda **kwargs: (None, None)
+        ax,
+        "get_custom_residue_moment_of_inertia",
+        lambda center_of_mass, positions, masses, custom_rot_axes, dimensions: np.array(
+            [1, 1, 1]
+        ),
     )
 
-    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]),
+    trans_axes, rot_axes, rot_center, moi = ax.get_residue_axes(
+        u, index=0, relative_index=0
     )
 
-    assert custom_axes is None
-    assert moi is None
+    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_bonded_axes_from_topology_no_bonded_heavy_uses_vanilla_axes(
-    monkeypatch,
-):
+def test_get_residue_bonded_axes_last_resid(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,),
+    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]),
+        ]
     )
 
-    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)
+    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]
+        ),
+    )
 
-    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=2, relative_index=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()
+    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_bonded_axes_from_topology_one_heavy_no_light_uses_custom_axes(
-    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 = _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=(),
+    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),
     )
 
-    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)
-
-    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_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),
     )
 
-    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 == "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]]
 
-    get_custom_axes = MagicMock(return_value=np.eye(3))
-    monkeypatch.setattr(ax, "get_custom_axes", get_custom_axes)
+    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([1.0, 2.0, 3.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))
-    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 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_two_heavy_uses_heavy_positions_as_b_list(
-    monkeypatch,
-):
+def test_get_ua_axes_bonded_axes_last_resid(monkeypatch):
     ax = AxesCalculator()
-
-    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,
-        }
+    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)),
+        ],
     )
-    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_atom_set = _FakeAtomGroup(
+        [
+            _atom(index=0, mass=12.0, pos=(1, 0, 0)),
+        ],
     )
 
-    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_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_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, 1.0, 1.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=0, res_position=1
     )
 
-    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)
+
+    chain = ax.get_chain(residue=residue, first=heavy_atoms[0], last=heavy_atoms[-1])
+
+    assert len(chain) == 3