diff --git a/.github/workflows/coveralls-main.yaml b/.github/workflows/coveralls-main.yaml
deleted file mode 100644
index 6537127d..00000000
--- a/.github/workflows/coveralls-main.yaml
+++ /dev/null
@@ -1,54 +0,0 @@
-name: Coveralls Main
-
-on:
-  push:
-    branches: [main]
-
-  workflow_dispatch:
-
-permissions:
-  contents: read
-
-concurrency:
-  group: coveralls-main-${{ github.ref }}
-  cancel-in-progress: true
-
-jobs:
-  coveralls:
-    name: Build main coverage baseline
-    runs-on: ubuntu-24.04
-    timeout-minutes: 30
-
-    permissions:
-      contents: read
-      checks: write
-      statuses: write
-
-    steps:
-      - name: Checkout
-        uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8
-
-      - name: Set up Python 3.14
-        uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405
-        with:
-          python-version: "3.14"
-
-      - name: Install testing dependencies
-        run: |
-          python -m pip install --upgrade pip
-          python -m pip install -e .[testing]
-
-      - name: Run coverage
-        run: |
-          python -m pytest tests/unit \
-            --cov CodeEntropy \
-            --cov-report term-missing \
-            --cov-report xml \
-            -q
-
-      - name: Upload coverage to Coveralls
-        uses: coverallsapp/github-action@5cbfd81b66ca5d10c19b062c04de0199c215fb6e
-        with:
-          github-token: ${{ secrets.GITHUB_TOKEN }}
-          file: coverage.xml
-          fail-on-error: false
diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml
index 795348dc..e2269cbe 100644
--- a/.github/workflows/pr.yaml
+++ b/.github/workflows/pr.yaml
@@ -2,17 +2,22 @@ name: Pull Request CI
 
 on:
   pull_request:
+  push:
+    branches: [main]
 
 permissions:
   contents: read
+  checks: write
+  pull-requests: write
 
 concurrency:
-  group: pr-ci-${{ github.ref }}
+  group: pr-${{ github.ref }}
   cancel-in-progress: true
 
 jobs:
   unit:
     name: Unit
+    if: github.event_name == 'pull_request'
     runs-on: ${{ matrix.os }}
     timeout-minutes: 25
     strategy:
@@ -35,11 +40,12 @@ jobs:
           python -m pip install --upgrade pip
           python -m pip install -e .[testing]
 
-      - name: Pytest unit • ${{ matrix.os }}, ${{ matrix.python-version }}
+      - name: Pytest (unit) • ${{ matrix.os }}, ${{ matrix.python-version }}
         run: python -m pytest tests/unit
 
   discover-systems:
     name: Discover regression systems
+    if: github.event_name == 'pull_request'
     runs-on: ubuntu-24.04
     outputs:
       systems: ${{ steps.set-systems.outputs.systems }}
@@ -62,10 +68,11 @@ jobs:
         id: set-systems
         run: |
           SYSTEMS=$(python -m tests.regression.list_systems)
-          echo "systems=$SYSTEMS" >> "$GITHUB_OUTPUT"
+          echo "systems=$SYSTEMS" >> $GITHUB_OUTPUT
 
   regression-quick:
-    name: Regression fast • ${{ matrix.system }}
+    name: Regression (fast) • ${{ matrix.system }}
+    if: github.event_name == 'pull_request'
     needs: discover-systems
     runs-on: ubuntu-24.04
     timeout-minutes: 35
@@ -110,6 +117,7 @@ jobs:
 
   docs:
     name: Docs
+    if: github.event_name == 'pull_request'
     needs: unit
     runs-on: ubuntu-24.04
     timeout-minutes: 25
@@ -142,6 +150,7 @@ jobs:
 
   pre-commit:
     name: Pre-commit
+    if: github.event_name == 'pull_request'
     needs: unit
     runs-on: ubuntu-24.04
     timeout-minutes: 15
@@ -175,12 +184,6 @@ jobs:
     runs-on: ubuntu-24.04
     timeout-minutes: 30
 
-    permissions:
-      contents: read
-      checks: write
-      statuses: write
-      pull-requests: write
-
     steps:
       - name: Checkout
         uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8
diff --git a/CITATION.cff b/CITATION.cff
index b98032b2..e4f657e8 100644
--- a/CITATION.cff
+++ b/CITATION.cff
@@ -76,5 +76,5 @@ keywords:
   - biomolecular simulations
   - protein flexibility
 license: MIT
-version: 2.3.0
-date-released: '2026-06-26'
+version: 2.2.4
+date-released: '2026-05-15'
diff --git a/CodeEntropy/__init__.py b/CodeEntropy/__init__.py
index bef57d77..b2301f7f 100644
--- a/CodeEntropy/__init__.py
+++ b/CodeEntropy/__init__.py
@@ -8,4 +8,4 @@
 and statistical mechanics.
 """
 
-__version__ = "2.3.0"
+__version__ = "2.2.4"
diff --git a/CodeEntropy/entropy/workflow.py b/CodeEntropy/entropy/workflow.py
index 9d5298e2..f8ddd0a8 100644
--- a/CodeEntropy/entropy/workflow.py
+++ b/CodeEntropy/entropy/workflow.py
@@ -26,8 +26,8 @@
 from CodeEntropy.core.logging import LoggingConfig
 from CodeEntropy.entropy.graph import EntropyGraph
 from CodeEntropy.entropy.water import WaterEntropy
-from CodeEntropy.levels.graph.level_dag import LevelDAG
 from CodeEntropy.levels.hierarchy import HierarchyBuilder
+from CodeEntropy.levels.level_dag import LevelDAG
 from CodeEntropy.trajectory.frames import FrameSelection
 from CodeEntropy.trajectory.source import FrameSource
 
diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py
index 4d4cf89c..f1c8bb75 100644
--- a/CodeEntropy/levels/axes.py
+++ b/CodeEntropy/levels/axes.py
@@ -67,7 +67,7 @@ def get_residue_axes(self, data_container, index: int, residue=None):
         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,7 +76,10 @@ 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):
@@ -99,101 +102,89 @@ 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
+        # match indexing to MDAnalysis indexing
         index_prev = index - 1
         index_next = index + 1
-
         if residue is None:
             residue = data_container.select_atoms(f"resindex {index}")
+            # residue of interest
         if len(residue) == 0:
             raise ValueError(f"Empty residue selection for resindex={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} 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
+            residue = data_container.residues[index]
+            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 +192,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 +212,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 +225,173 @@ 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
-
+        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()
+            else:
+                # residue of interest has at least one neighbour
+                if res_position == -1:
+                    residue = data_container.residues[0]
+                    index_next = residue.resid + 1
+                    # the .resid attribute gives 1-indexing
+                    # substract 1 to match indexing later
+                    second_edge = data_container.select_atoms(
+                        f"resindex {residue.resid - 1} and "
+                        f"bonded resindex {index_next - 1}"
+                    )
+                    edges = [residue.atoms[0], second_edge.atoms[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]
+                    index_prev = residue.resid - 1
+                    index_next = residue.resid + 1
+                    edge_set = data_container.select_atoms(
+                        f"resindex {residue.resid - 1} and "
+                        f"(bonded resindex {index_next - 1} or "
+                        f"resindex {index_prev - 1})"
+                    )
+                    edges = [edge_set[0], edge_set[1]]
+                    backbone = self.get_chain(residue, edge_set[0], edge_set[1])
+
+                else:
+                    # last resid
+                    residue = data_container.residues[1]
+                    index_prev = residue.resid - 1
+                    first_edge = data_container.select_atoms(
+                        f"resindex {residue.resid - 1} and "
+                        f"bonded resindex {index_prev - 1}"
+                    )
+                    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}")
+            rot_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],
             )
-        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 +620,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 +776,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 +807,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 +845,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: MDAnalysis AtomGroup containing chain atoms.
+        """
+        chain = []
+        chain_indices = []
+        # 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
+        chain_indices = [last.index - residue.atoms.indices[0]]
+        # 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_indices.append(current.index - residue.atoms.indices[0])
+        chain_indices = np.flip(chain_indices)
+        # only get in between residues
+        chain_indices = chain_indices[1:-1]
+        # accout for in-residue index
+        chain_AtomGroup = residue.atoms[chain_indices]
+        chain = chain_AtomGroup.atoms.select_atoms("all")
+        return chain
diff --git a/CodeEntropy/levels/graph/conformation_dag.py b/CodeEntropy/levels/conformation_dag.py
similarity index 100%
rename from CodeEntropy/levels/graph/conformation_dag.py
rename to CodeEntropy/levels/conformation_dag.py
diff --git a/CodeEntropy/levels/dihedrals/topology.py b/CodeEntropy/levels/dihedrals/topology.py
index 31e35ffd..680639a1 100644
--- a/CodeEntropy/levels/dihedrals/topology.py
+++ b/CodeEntropy/levels/dihedrals/topology.py
@@ -59,7 +59,9 @@ def _discover_group_dihedral_topology(
         topologies: list[MoleculeDihedralTopology] = []
 
         for molecule_order, molecule_id in enumerate(molecules):
-            mol = self._extract_topology_fragment(data_container, molecule_id)
+            mol = self._universe_operations.extract_fragment(
+                data_container, molecule_id
+            )
             num_residues = len(mol.residues)
             ua_dihedrals_by_residue: dict[int, list[Any]] = {}
             residue_dihedrals: list[Any] = []
@@ -88,26 +90,6 @@ def _discover_group_dihedral_topology(
 
         return topologies
 
-    def _extract_topology_fragment(self, data_container: Any, molecule_id: Any) -> Any:
-        """Return a molecule fragment for topology discovery.
-
-        This uses the lightweight AtomGroup extraction helper when available so
-        static conformational topology discovery does not create a standalone
-        in-memory universe or copy trajectory frames. The fallback preserves
-        compatibility with older ``UniverseOperations`` implementations.
-
-        Args:
-            data_container: Source MDAnalysis universe or universe-like container.
-            molecule_id: Fragment index identifying the molecule to extract.
-
-        Returns:
-            MDAnalysis AtomGroup for the selected molecule
-        """
-        return self._universe_operations.extract_fragment_atomgroup(
-            data_container,
-            int(molecule_id),
-        )
-
     def _select_heavy_residue(self, mol: Any, res_id: int) -> Any:
         """Select heavy atoms in a residue by residue index.
 
@@ -118,15 +100,13 @@ def _select_heavy_residue(self, mol: Any, res_id: int) -> Any:
         Returns:
             AtomGroup containing heavy atoms in the residue selection.
         """
-        residue_atoms = mol.residues[int(res_id)].atoms
-        selection1 = residue_atoms.indices[0]
-        selection2 = residue_atoms.indices[-1]
+        selection1 = mol.residues[res_id].atoms.indices[0]
+        selection2 = mol.residues[res_id].atoms.indices[-1]
 
-        res_container = mol.select_atoms(
-            f"index {selection1}:{selection2}",
-            updating=False,
+        res_container = self._universe_operations.select_atoms(
+            mol, f"index {selection1}:{selection2}"
         )
-        return res_container.select_atoms("prop mass > 1.1", updating=False)
+        return self._universe_operations.select_atoms(res_container, "prop mass > 1.1")
 
     def _get_dihedrals(self, data_container: Any, level: str) -> list[Any]:
         """Return dihedral AtomGroups for a container at a given level.
@@ -141,93 +121,26 @@ def _get_dihedrals(self, data_container: Any, level: str) -> list[Any]:
         atom_groups: list[Any] = []
 
         if level == "united_atom":
-            selected_indices = {int(index) for index in data_container.indices}
-
             for dihedral in data_container.dihedrals:
-                dihedral_atoms = dihedral.atoms
-                dihedral_indices = {int(index) for index in dihedral_atoms.indices}
-
-                if len(dihedral_atoms) == 4 and dihedral_indices.issubset(
-                    selected_indices
-                ):
-                    atom_groups.append(dihedral_atoms)
+                atom_groups.append(dihedral.atoms)
 
         if level == "residue":
             num_residues = len(data_container.residues)
             if num_residues >= 4:
                 for residue in range(4, num_residues + 1):
-                    residue1 = data_container.residues[residue - 4]
-                    residue2 = data_container.residues[residue - 3]
-                    residue3 = data_container.residues[residue - 2]
-                    residue4 = data_container.residues[residue - 1]
-
-                    atom1 = self._atoms_in_source_bonded_to_target(
-                        residue1,
-                        residue2,
+                    atom1 = data_container.select_atoms(
+                        f"resindex {residue - 4} and bonded resindex {residue - 3}"
                     )
-                    atom2 = self._atoms_in_source_bonded_to_target(
-                        residue2,
-                        residue1,
+                    atom2 = data_container.select_atoms(
+                        f"resindex {residue - 3} and bonded resindex {residue - 4}"
                     )
-                    atom3 = self._atoms_in_source_bonded_to_target(
-                        residue3,
-                        residue4,
+                    atom3 = data_container.select_atoms(
+                        f"resindex {residue - 2} and bonded resindex {residue - 1}"
                     )
-                    atom4 = self._atoms_in_source_bonded_to_target(
-                        residue4,
-                        residue3,
+                    atom4 = data_container.select_atoms(
+                        f"resindex {residue - 1} and bonded resindex {residue - 2}"
                     )
-
-                    dihedral_atoms = atom1 + atom2 + atom3 + atom4
-
-                    if len(dihedral_atoms) == 4:
-                        atom_groups.append(dihedral_atoms)
-                    else:
-                        logger.debug(
-                            "Skipping residue-level dihedral for local residues "
-                            "%s-%s-%s-%s because it produced %d atoms.",
-                            residue - 4,
-                            residue - 3,
-                            residue - 2,
-                            residue - 1,
-                            len(dihedral_atoms),
-                        )
+                    atom_groups.append(atom1 + atom2 + atom3 + atom4)
 
         logger.debug("Level: %s, Dihedrals: %s", level, atom_groups)
         return atom_groups
-
-    @staticmethod
-    def _atoms_in_source_bonded_to_target(
-        source_residue: Any,
-        target_residue: Any,
-    ) -> Any:
-        """Return source-residue atoms bonded to atoms in a target residue.
-
-        This helper is used when constructing residue-level dihedral definitions
-        from lightweight molecule AtomGroups. It selects atoms from the source
-        residue that are bonded to any atom in the target residue without using
-        global ``resindex`` selection strings.
-
-        Args:
-            source_residue: Residue whose atoms should be tested for bonds.
-            target_residue: Adjacent residue providing the target bonded atoms.
-
-        Returns:
-            MDAnalysis AtomGroup containing atoms from ``source_residue`` that are
-            bonded to at least one atom in ``target_residue``. If no matching
-            atoms are found, an empty AtomGroup is returned.
-        """
-        source_atoms = source_residue.atoms
-        target_indices = {int(index) for index in target_residue.atoms.indices}
-        selected_indices: list[int] = []
-
-        for atom in source_atoms:
-            bonded_atoms = getattr(atom, "bonded_atoms", None)
-            if bonded_atoms is None:
-                continue
-
-            bonded_indices = {int(index) for index in bonded_atoms.indices}
-            if bonded_indices.intersection(target_indices):
-                selected_indices.append(int(atom.index))
-
-        return source_atoms.universe.atoms[selected_indices]
diff --git a/CodeEntropy/levels/execution/__init__.py b/CodeEntropy/levels/execution/__init__.py
deleted file mode 100644
index e69de29b..00000000
diff --git a/CodeEntropy/levels/execution/scheduler.py b/CodeEntropy/levels/execution/scheduler.py
index 0a5b5063..5f3d8455 100644
--- a/CodeEntropy/levels/execution/scheduler.py
+++ b/CodeEntropy/levels/execution/scheduler.py
@@ -16,7 +16,7 @@
     execute_frame_map_output,
     make_frame_worker_shared_data,
 )
-from CodeEntropy.levels.graph.frame_dag import FrameGraph
+from CodeEntropy.levels.frame_dag import FrameGraph
 from CodeEntropy.levels.neighbors import Neighbors
 from CodeEntropy.results.reporter import _RichProgressSink
 
diff --git a/CodeEntropy/levels/execution/tasks.py b/CodeEntropy/levels/execution/tasks.py
index 872d1ad1..a0997ccf 100644
--- a/CodeEntropy/levels/execution/tasks.py
+++ b/CodeEntropy/levels/execution/tasks.py
@@ -5,7 +5,7 @@
 from dataclasses import dataclass, field
 from typing import Any
 
-from CodeEntropy.levels.graph.frame_dag import FrameGraph
+from CodeEntropy.levels.frame_dag import FrameGraph
 from CodeEntropy.levels.neighbors import Neighbors
 
 FRAME_WORKER_EXCLUDED_SHARED_KEYS = {
diff --git a/CodeEntropy/levels/graph/frame_dag.py b/CodeEntropy/levels/frame_dag.py
similarity index 100%
rename from CodeEntropy/levels/graph/frame_dag.py
rename to CodeEntropy/levels/frame_dag.py
diff --git a/CodeEntropy/levels/graph/__init__.py b/CodeEntropy/levels/graph/__init__.py
deleted file mode 100644
index e69de29b..00000000
diff --git a/CodeEntropy/levels/graph/level_dag.py b/CodeEntropy/levels/level_dag.py
similarity index 95%
rename from CodeEntropy/levels/graph/level_dag.py
rename to CodeEntropy/levels/level_dag.py
index e333c772..cc8129a1 100644
--- a/CodeEntropy/levels/graph/level_dag.py
+++ b/CodeEntropy/levels/level_dag.py
@@ -12,14 +12,13 @@
 import networkx as nx
 
 from CodeEntropy.levels.axes import AxesCalculator
+from CodeEntropy.levels.conformation_dag import ConformationDAG
 from CodeEntropy.levels.execution.policy import ExecutionPolicy
 from CodeEntropy.levels.execution.reducers import NeighborReducer
 from CodeEntropy.levels.execution.scheduler import FrameScheduler
-from CodeEntropy.levels.graph.conformation_dag import ConformationDAG
-from CodeEntropy.levels.graph.frame_dag import FrameGraph
+from CodeEntropy.levels.frame_dag import FrameGraph
 from CodeEntropy.levels.neighbors import Neighbors
 from CodeEntropy.levels.nodes.accumulators import InitCovarianceAccumulatorsNode
-from CodeEntropy.levels.nodes.axes_topology import BuildAxesTopologyNode
 from CodeEntropy.levels.nodes.beads import BuildBeadsNode
 from CodeEntropy.levels.nodes.detect_levels import DetectLevelsNode
 from CodeEntropy.levels.nodes.detect_molecules import DetectMoleculesNode
@@ -50,11 +49,6 @@ def build(self) -> LevelDAG:
         self._add_static("detect_molecules", DetectMoleculesNode())
         self._add_static("detect_levels", DetectLevelsNode(), deps=["detect_molecules"])
         self._add_static("build_beads", BuildBeadsNode(), deps=["detect_levels"])
-        self._add_static(
-            "build_axes_topology",
-            BuildAxesTopologyNode(),
-            deps=["build_beads"],
-        )
         self._add_static(
             "init_covariance_accumulators",
             InitCovarianceAccumulatorsNode(),
diff --git a/CodeEntropy/levels/nodes/axes_topology.py b/CodeEntropy/levels/nodes/axes_topology.py
deleted file mode 100644
index 5197d752..00000000
--- a/CodeEntropy/levels/nodes/axes_topology.py
+++ /dev/null
@@ -1,316 +0,0 @@
-"""Build static axes-topology metadata for frame covariance calculations.
-
-This module caches topology-only atom-index relationships needed by customised
-axes calculations. The cache avoids repeated MDAnalysis selection parsing inside
-the frame-local covariance loop while preserving frame-dependent positions,
-forces, centres, axes, torques, and moments of inertia.
-"""
-
-from __future__ import annotations
-
-import logging
-from dataclasses import dataclass, field
-from typing import Any
-
-import numpy as np
-
-logger = logging.getLogger(__name__)
-
-UAKey = tuple[int, int, int]
-ResidueKey = tuple[int, int]
-
-
-@dataclass(frozen=True)
-class UAAxesTopology:
-    """Static topology required to compute customised united-atom axes.
-
-    Attributes:
-        heavy_atom_index: Reduced-universe atom index for the UA heavy atom.
-        ua_atom_indices: Atom indices for the UA heavy atom and its bonded
-            hydrogens/light atoms.
-        ua_all_atom_indices: Atom indices for the UA heavy atom, bonded heavy
-            atoms, and bonded hydrogens/light atoms.
-        bonded_heavy_indices: Heavy atoms bonded to the UA heavy atom.
-        bonded_light_indices: Hydrogens/light atoms bonded to the UA heavy atom.
-        residue_heavy_indices: Heavy atoms in the parent residue.
-        residue_ua_masses: UA masses for heavy atoms in the parent residue.
-    """
-
-    heavy_atom_index: int
-    ua_atom_indices: np.ndarray
-    ua_all_atom_indices: np.ndarray
-    bonded_heavy_indices: np.ndarray
-    bonded_light_indices: np.ndarray
-    residue_heavy_indices: np.ndarray
-    residue_ua_masses: np.ndarray
-
-
-@dataclass(frozen=True)
-class ResidueAxesTopology:
-    """Static topology required to compute customised residue axes.
-
-    Attributes:
-        residue_heavy_indices: Heavy atom indices in the residue.
-        residue_ua_masses: UA masses for heavy atoms in the residue.
-        has_neighbor_bonds: Whether the residue is bonded to a neighbouring
-            residue according to the original customised residue-axis selection.
-    """
-
-    residue_heavy_indices: np.ndarray
-    residue_ua_masses: np.ndarray
-    has_neighbor_bonds: bool
-
-
-@dataclass(frozen=True)
-class AxesTopology:
-    """Cached axes topology for frame covariance calculations.
-
-    Attributes:
-        ua: Mapping from ``(mol_id, local_residue_id, ua_id)`` to cached
-            united-atom axes topology.
-        residue: Mapping from ``(mol_id, local_residue_id)`` to cached
-            residue axes topology.
-    """
-
-    ua: dict[UAKey, UAAxesTopology] = field(default_factory=dict)
-    residue: dict[ResidueKey, ResidueAxesTopology] = field(default_factory=dict)
-
-
-class BuildAxesTopologyNode:
-    """Build static customised-axes topology before frame covariance execution."""
-
-    def run(self, shared_data: dict[str, Any]) -> dict[str, Any]:
-        """Build cached axes topology and write it into shared data.
-
-        The cache is only populated when ``args.customised_axes`` is true. When
-        customised axes are disabled, an empty cache is still written so later
-        stages can read ``shared_data["axes_topology"]`` safely.
-
-        Args:
-            shared_data: Shared workflow data containing ``args`` and, when
-                customised axes are enabled, ``reduced_universe``, ``levels``,
-                and ``beads``.
-
-        Returns:
-            Dict containing the cached ``axes_topology`` object.
-        """
-        args = shared_data["args"]
-        topology = AxesTopology()
-
-        if not bool(getattr(args, "customised_axes", False)):
-            shared_data["axes_topology"] = topology
-            return {"axes_topology": topology}
-
-        u = shared_data["reduced_universe"]
-        levels = shared_data["levels"]
-        beads = shared_data["beads"]
-
-        ua_topology: dict[UAKey, UAAxesTopology] = {}
-        residue_topology: dict[ResidueKey, ResidueAxesTopology] = {}
-        fragments = u.atoms.fragments
-
-        for mol_id, level_list in enumerate(levels):
-            mol = fragments[mol_id]
-
-            if "residue" in level_list:
-                self._add_residue_topology(
-                    mol=mol,
-                    mol_id=mol_id,
-                    beads=beads,
-                    out=residue_topology,
-                )
-
-            if "united_atom" in level_list:
-                self._add_ua_topology(
-                    u=u,
-                    mol=mol,
-                    mol_id=mol_id,
-                    beads=beads,
-                    out=ua_topology,
-                )
-
-        topology = AxesTopology(ua=ua_topology, residue=residue_topology)
-        shared_data["axes_topology"] = topology
-        return {"axes_topology": topology}
-
-    def _add_residue_topology(
-        self,
-        *,
-        mol: Any,
-        mol_id: int,
-        beads: dict[Any, list[np.ndarray]],
-        out: dict[ResidueKey, ResidueAxesTopology],
-    ) -> None:
-        """Cache static residue axes topology for one molecule.
-
-        Args:
-            mol: Molecule AtomGroup.
-            mol_id: Molecule index.
-            beads: Bead-index mapping produced by ``BuildBeadsNode``.
-            out: Output residue topology mapping mutated in place.
-        """
-        bead_key = (mol_id, "residue")
-        bead_idx_list = beads.get(bead_key, [])
-        if not bead_idx_list:
-            return
-
-        for local_res_i, residue in enumerate(mol.residues):
-            if local_res_i >= len(bead_idx_list):
-                continue
-
-            residue_atoms = residue.atoms
-            residue_heavy = residue_atoms.select_atoms("mass 2 to 999")
-            residue_heavy_indices = residue_heavy.indices.astype(int, copy=True)
-            residue_ua_masses = np.asarray(
-                self._get_ua_masses_from_topology(residue_atoms),
-                dtype=float,
-            )
-            has_neighbor_bonds = self._has_neighbor_bonds(
-                mol=mol,
-                local_res_i=local_res_i,
-            )
-
-            out[(mol_id, local_res_i)] = ResidueAxesTopology(
-                residue_heavy_indices=residue_heavy_indices,
-                residue_ua_masses=residue_ua_masses,
-                has_neighbor_bonds=has_neighbor_bonds,
-            )
-
-    def _add_ua_topology(
-        self,
-        *,
-        u: Any,
-        mol: Any,
-        mol_id: int,
-        beads: dict[Any, list[np.ndarray]],
-        out: dict[UAKey, UAAxesTopology],
-    ) -> None:
-        """Cache static UA axes topology for one molecule.
-
-        Args:
-            u: Reduced universe used to resolve bead atom-index arrays.
-            mol: Molecule AtomGroup.
-            mol_id: Molecule index.
-            beads: Bead-index mapping produced by ``BuildBeadsNode``.
-            out: Output UA topology mapping mutated in place.
-        """
-        for local_res_i, residue in enumerate(mol.residues):
-            bead_key = (mol_id, "united_atom", local_res_i)
-            bead_idx_list = beads.get(bead_key, [])
-
-            if not bead_idx_list:
-                continue
-
-            residue_atoms = residue.atoms
-            residue_heavy = residue_atoms.select_atoms("prop mass > 1.1")
-            residue_heavy_indices = residue_heavy.indices.astype(int, copy=True)
-            residue_ua_masses = np.asarray(
-                self._get_ua_masses_from_topology(residue_atoms),
-                dtype=float,
-            )
-
-            for ua_i, bead_indices in enumerate(bead_idx_list):
-                bead = u.atoms[bead_indices]
-                heavy = bead.select_atoms("prop mass > 1.1")
-
-                if len(heavy) == 0:
-                    logger.warning(
-                        "Skipping UA axes topology with no heavy atom: "
-                        "mol=%s residue=%s ua=%s",
-                        mol_id,
-                        local_res_i,
-                        ua_i,
-                    )
-                    continue
-
-                heavy_atom = heavy[0]
-                bonded_heavy, bonded_light = self._split_bonded_atoms(heavy_atom)
-
-                heavy_index = np.asarray([int(heavy_atom.index)], dtype=int)
-                bonded_heavy_indices = bonded_heavy.indices.astype(int, copy=True)
-                bonded_light_indices = bonded_light.indices.astype(int, copy=True)
-
-                ua_atom_indices = np.concatenate(
-                    [heavy_index, bonded_light_indices],
-                    axis=0,
-                )
-                ua_all_atom_indices = np.concatenate(
-                    [heavy_index, bonded_heavy_indices, bonded_light_indices],
-                    axis=0,
-                )
-
-                out[(mol_id, local_res_i, ua_i)] = UAAxesTopology(
-                    heavy_atom_index=int(heavy_atom.index),
-                    ua_atom_indices=ua_atom_indices,
-                    ua_all_atom_indices=ua_all_atom_indices,
-                    bonded_heavy_indices=bonded_heavy_indices,
-                    bonded_light_indices=bonded_light_indices,
-                    residue_heavy_indices=residue_heavy_indices,
-                    residue_ua_masses=residue_ua_masses,
-                )
-
-    @staticmethod
-    def _has_neighbor_bonds(*, mol: Any, local_res_i: int) -> bool:
-        """Return whether a residue is bonded to neighbouring residues.
-
-        Args:
-            mol: Molecule AtomGroup used for the original bonded-neighbour
-                selection.
-            local_res_i: Residue index local to ``mol``.
-
-        Returns:
-            True when the residue has bonded atoms in the previous or next
-            residue according to the original customised residue-axis query.
-        """
-        index_prev = local_res_i - 1
-        index_next = local_res_i + 1
-        atom_set = mol.select_atoms(
-            f"(resindex {index_prev} or resindex {index_next}) "
-            f"and bonded resid {local_res_i}"
-        )
-        return len(atom_set) > 0
-
-    @staticmethod
-    def _split_bonded_atoms(atom: Any) -> tuple[Any, Any]:
-        """Return bonded heavy and light atoms for one atom.
-
-        Args:
-            atom: MDAnalysis Atom.
-
-        Returns:
-            Tuple containing bonded heavy atoms and bonded hydrogens/light atoms.
-        """
-        bonded_atoms = atom.bonded_atoms
-        bonded_heavy = bonded_atoms.select_atoms("mass 2 to 999")
-        bonded_light = bonded_atoms.select_atoms("mass 1 to 1.1")
-        return bonded_heavy, bonded_light
-
-    @staticmethod
-    def _get_ua_masses_from_topology(atom_group: Any) -> list[float]:
-        """Return UA masses using static bonded atom relationships.
-
-        Args:
-            atom_group: AtomGroup containing atoms from one residue.
-
-        Returns:
-            List of UA masses, one for each heavy atom in ``atom_group``.
-        """
-        ua_masses: list[float] = []
-
-        for atom in atom_group:
-            if atom.mass <= 1.1:
-                continue
-
-            ua_mass = float(atom.mass)
-            bonded_atoms = getattr(atom, "bonded_atoms", None)
-            if bonded_atoms is None:
-                ua_masses.append(ua_mass)
-                continue
-
-            bonded_h_atoms = bonded_atoms.select_atoms("mass 1 to 1.1")
-            for hydrogen in bonded_h_atoms:
-                ua_mass += float(hydrogen.mass)
-
-            ua_masses.append(ua_mass)
-
-        return ua_masses
diff --git a/CodeEntropy/levels/nodes/covariance.py b/CodeEntropy/levels/nodes/covariance.py
index 39ff61a3..38ec3d1c 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,30 @@ def _process_united_atom(
             out_torque: Frame-local torque second-moment accumulator, mutated in place.
             molcount: Per-residue group sample counters, mutated in place.
         """
+
         for local_res_i, res in enumerate(mol.residues):
+            if len(mol.residues) > 1:
+                # there are multiple residues in the molecule
+                # build residue group here
+                if local_res_i == 0:
+                    # first residue
+                    res_position = -1
+                    res_next = mol.residues[1]
+                    residue_group = res + res_next
+                elif local_res_i == len(mol.residues) - 1:
+                    # last residue
+                    res_position = 1
+                    res_prev = mol.residues[-2]
+                    residue_group = res + res_prev
+                else:
+                    res_position = 0
+                    res_prev = mol.residues[local_res_i - 1]
+                    res_next = mol.residues[local_res_i + 1]
+                    residue_group = res_prev + res + res_next
+            else:
+                # only one residue
+                res_position = None
+                residue_group = res
             bead_key = (mol_id, "united_atom", local_res_i)
             bead_idx_list = beads.get(bead_key, [])
             if not bead_idx_list:
@@ -213,17 +231,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 +258,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 +277,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 +297,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 +412,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 +442,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 +486,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 +497,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.
@@ -520,14 +513,10 @@ def _build_residue_vectors(
 
         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,
                 local_res_i=local_res_i,
                 axes_manager=axes_manager,
-                axes_topology=axes_topology,
-                box=box,
                 customised_axes=customised_axes,
             )
 
@@ -556,27 +545,19 @@ def _build_residue_vectors(
     def _get_residue_axes(
         self,
         *,
-        u: Any,
         mol: Any,
-        mol_id: int,
         bead: Any,
         local_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,19 +565,6 @@ 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)
 
         make_whole(mol.atoms)
diff --git a/CodeEntropy/trajectory/mda.py b/CodeEntropy/trajectory/mda.py
index 1cdb6e75..04943a4c 100644
--- a/CodeEntropy/trajectory/mda.py
+++ b/CodeEntropy/trajectory/mda.py
@@ -25,8 +25,7 @@ class UniverseOperations:
     This helper provides methods to:
       - Build reduced universes by selecting subsets of frames or atoms.
       - Extract a single fragment (molecule) into a standalone universe.
-      - Merge coordinates from one trajectory with forces sourced from another
-            trajectory.
+      - Merge coordinates from one trajectory with forces from another trajectory.
     """
 
     def __init__(self) -> None:
@@ -187,25 +186,6 @@ def extract_fragment(
         selection_string = f"index {frag.indices[0]}:{frag.indices[-1]}"
         return self.select_atoms(universe, selection_string)
 
-    def extract_fragment_atomgroup(self, universe: mda.Universe, molecule_id: int):
-        """Return a molecule fragment as an AtomGroup.
-
-        This helper mirrors the atom-index range used by ``extract_fragment`` but
-        avoids building a standalone in-memory universe. It is intended for
-        topology discovery paths where only the static atom selection is needed
-        and trajectory coordinates do not need to be copied.
-
-        Args:
-            universe: Source MDAnalysis universe.
-            molecule_id: Fragment index in ``universe.atoms.fragments``.
-
-        Returns:
-            MDAnalysis AtomGroup containing the atoms in the selected fragment.
-        """
-        frag = universe.atoms.fragments[int(molecule_id)]
-        selection_string = f"index {frag.indices[0]}:{frag.indices[-1]}"
-        return universe.select_atoms(selection_string, updating=False)
-
     def convert_lammps(
         self,
         tprfile: str,
diff --git a/docs/api/CodeEntropy.config.rst b/docs/api/CodeEntropy.config.rst
index d7945e83..c85bcb89 100644
--- a/docs/api/CodeEntropy.config.rst
+++ b/docs/api/CodeEntropy.config.rst
@@ -1,6 +1,11 @@
 CodeEntropy.config package
 ==========================
 
+.. automodule:: CodeEntropy.config
+   :members:
+   :show-inheritance:
+   :undoc-members:
+
 Submodules
 ----------
 
@@ -9,11 +14,3 @@ Submodules
 
    CodeEntropy.config.argparse
    CodeEntropy.config.runtime
-
-Module contents
----------------
-
-.. automodule:: CodeEntropy.config
-   :members:
-   :show-inheritance:
-   :undoc-members:
diff --git a/docs/api/CodeEntropy.core.rst b/docs/api/CodeEntropy.core.rst
index 86ebc99c..695ee31d 100644
--- a/docs/api/CodeEntropy.core.rst
+++ b/docs/api/CodeEntropy.core.rst
@@ -1,6 +1,11 @@
 CodeEntropy.core package
 ========================
 
+.. automodule:: CodeEntropy.core
+   :members:
+   :show-inheritance:
+   :undoc-members:
+
 Submodules
 ----------
 
@@ -9,11 +14,3 @@ Submodules
 
    CodeEntropy.core.dask_clusters
    CodeEntropy.core.logging
-
-Module contents
----------------
-
-.. automodule:: CodeEntropy.core
-   :members:
-   :show-inheritance:
-   :undoc-members:
diff --git a/docs/api/CodeEntropy.entropy.nodes.rst b/docs/api/CodeEntropy.entropy.nodes.rst
index 4ff7516c..2fa9338e 100644
--- a/docs/api/CodeEntropy.entropy.nodes.rst
+++ b/docs/api/CodeEntropy.entropy.nodes.rst
@@ -1,6 +1,11 @@
 CodeEntropy.entropy.nodes package
 =================================
 
+.. automodule:: CodeEntropy.entropy.nodes
+   :members:
+   :show-inheritance:
+   :undoc-members:
+
 Submodules
 ----------
 
@@ -11,11 +16,3 @@ Submodules
    CodeEntropy.entropy.nodes.configurational
    CodeEntropy.entropy.nodes.orientational
    CodeEntropy.entropy.nodes.vibrational
-
-Module contents
----------------
-
-.. automodule:: CodeEntropy.entropy.nodes
-   :members:
-   :show-inheritance:
-   :undoc-members:
diff --git a/docs/api/CodeEntropy.entropy.rst b/docs/api/CodeEntropy.entropy.rst
index f4322465..99ed6b3f 100644
--- a/docs/api/CodeEntropy.entropy.rst
+++ b/docs/api/CodeEntropy.entropy.rst
@@ -1,6 +1,11 @@
 CodeEntropy.entropy package
 ===========================
 
+.. automodule:: CodeEntropy.entropy
+   :members:
+   :show-inheritance:
+   :undoc-members:
+
 Subpackages
 -----------
 
@@ -21,11 +26,3 @@ Submodules
    CodeEntropy.entropy.vibrational
    CodeEntropy.entropy.water
    CodeEntropy.entropy.workflow
-
-Module contents
----------------
-
-.. automodule:: CodeEntropy.entropy
-   :members:
-   :show-inheritance:
-   :undoc-members:
diff --git a/docs/api/CodeEntropy.levels.dihedrals.kernels.rst b/docs/api/CodeEntropy.levels.conformation_dag.rst
similarity index 50%
rename from docs/api/CodeEntropy.levels.dihedrals.kernels.rst
rename to docs/api/CodeEntropy.levels.conformation_dag.rst
index 3e51f72a..ffea54eb 100644
--- a/docs/api/CodeEntropy.levels.dihedrals.kernels.rst
+++ b/docs/api/CodeEntropy.levels.conformation_dag.rst
@@ -1,7 +1,7 @@
-CodeEntropy.levels.dihedrals.kernels module
+CodeEntropy.levels.conformation\_dag module
 ===========================================
 
-.. automodule:: CodeEntropy.levels.dihedrals.kernels
+.. automodule:: CodeEntropy.levels.conformation_dag
    :members:
    :show-inheritance:
    :undoc-members:
diff --git a/docs/api/CodeEntropy.levels.dihedrals.angle_observations.rst b/docs/api/CodeEntropy.levels.dihedrals.angle_observations.rst
deleted file mode 100644
index 16bf9470..00000000
--- a/docs/api/CodeEntropy.levels.dihedrals.angle_observations.rst
+++ /dev/null
@@ -1,7 +0,0 @@
-CodeEntropy.levels.dihedrals.angle\_observations module
-=======================================================
-
-.. automodule:: CodeEntropy.levels.dihedrals.angle_observations
-   :members:
-   :show-inheritance:
-   :undoc-members:
diff --git a/docs/api/CodeEntropy.levels.dihedrals.conformational_state_builder.rst b/docs/api/CodeEntropy.levels.dihedrals.conformational_state_builder.rst
deleted file mode 100644
index 449f21c8..00000000
--- a/docs/api/CodeEntropy.levels.dihedrals.conformational_state_builder.rst
+++ /dev/null
@@ -1,7 +0,0 @@
-CodeEntropy.levels.dihedrals.conformational\_state\_builder module
-==================================================================
-
-.. automodule:: CodeEntropy.levels.dihedrals.conformational_state_builder
-   :members:
-   :show-inheritance:
-   :undoc-members:
diff --git a/docs/api/CodeEntropy.levels.dihedrals.peak_detection.rst b/docs/api/CodeEntropy.levels.dihedrals.peak_detection.rst
deleted file mode 100644
index 11c89969..00000000
--- a/docs/api/CodeEntropy.levels.dihedrals.peak_detection.rst
+++ /dev/null
@@ -1,7 +0,0 @@
-CodeEntropy.levels.dihedrals.peak\_detection module
-===================================================
-
-.. automodule:: CodeEntropy.levels.dihedrals.peak_detection
-   :members:
-   :show-inheritance:
-   :undoc-members:
diff --git a/docs/api/CodeEntropy.levels.dihedrals.rst b/docs/api/CodeEntropy.levels.dihedrals.rst
index 30aa5a67..a69952a9 100644
--- a/docs/api/CodeEntropy.levels.dihedrals.rst
+++ b/docs/api/CodeEntropy.levels.dihedrals.rst
@@ -1,21 +1,5 @@
-CodeEntropy.levels.dihedrals package
-====================================
-
-Submodules
-----------
-
-.. toctree::
-   :maxdepth: 4
-
-   CodeEntropy.levels.dihedrals.angle_observations
-   CodeEntropy.levels.dihedrals.conformational_state_builder
-   CodeEntropy.levels.dihedrals.kernels
-   CodeEntropy.levels.dihedrals.peak_detection
-   CodeEntropy.levels.dihedrals.state_assignment
-   CodeEntropy.levels.dihedrals.topology
-
-Module contents
----------------
+CodeEntropy.levels.dihedrals module
+===================================
 
 .. automodule:: CodeEntropy.levels.dihedrals
    :members:
diff --git a/docs/api/CodeEntropy.levels.dihedrals.state_assignment.rst b/docs/api/CodeEntropy.levels.dihedrals.state_assignment.rst
deleted file mode 100644
index aeed41ce..00000000
--- a/docs/api/CodeEntropy.levels.dihedrals.state_assignment.rst
+++ /dev/null
@@ -1,7 +0,0 @@
-CodeEntropy.levels.dihedrals.state\_assignment module
-=====================================================
-
-.. automodule:: CodeEntropy.levels.dihedrals.state_assignment
-   :members:
-   :show-inheritance:
-   :undoc-members:
diff --git a/docs/api/CodeEntropy.levels.dihedrals.topology.rst b/docs/api/CodeEntropy.levels.dihedrals.topology.rst
deleted file mode 100644
index 30074514..00000000
--- a/docs/api/CodeEntropy.levels.dihedrals.topology.rst
+++ /dev/null
@@ -1,7 +0,0 @@
-CodeEntropy.levels.dihedrals.topology module
-============================================
-
-.. automodule:: CodeEntropy.levels.dihedrals.topology
-   :members:
-   :show-inheritance:
-   :undoc-members:
diff --git a/docs/api/CodeEntropy.levels.execution.chunks.rst b/docs/api/CodeEntropy.levels.execution.chunks.rst
deleted file mode 100644
index 6f7924f8..00000000
--- a/docs/api/CodeEntropy.levels.execution.chunks.rst
+++ /dev/null
@@ -1,7 +0,0 @@
-CodeEntropy.levels.execution.chunks module
-==========================================
-
-.. automodule:: CodeEntropy.levels.execution.chunks
-   :members:
-   :show-inheritance:
-   :undoc-members:
diff --git a/docs/api/CodeEntropy.levels.execution.policy.rst b/docs/api/CodeEntropy.levels.execution.policy.rst
deleted file mode 100644
index 3dd1f5aa..00000000
--- a/docs/api/CodeEntropy.levels.execution.policy.rst
+++ /dev/null
@@ -1,7 +0,0 @@
-CodeEntropy.levels.execution.policy module
-==========================================
-
-.. automodule:: CodeEntropy.levels.execution.policy
-   :members:
-   :show-inheritance:
-   :undoc-members:
diff --git a/docs/api/CodeEntropy.levels.execution.reducers.rst b/docs/api/CodeEntropy.levels.execution.reducers.rst
deleted file mode 100644
index 0beef239..00000000
--- a/docs/api/CodeEntropy.levels.execution.reducers.rst
+++ /dev/null
@@ -1,7 +0,0 @@
-CodeEntropy.levels.execution.reducers module
-============================================
-
-.. automodule:: CodeEntropy.levels.execution.reducers
-   :members:
-   :show-inheritance:
-   :undoc-members:
diff --git a/docs/api/CodeEntropy.levels.execution.rst b/docs/api/CodeEntropy.levels.execution.rst
deleted file mode 100644
index 2932017a..00000000
--- a/docs/api/CodeEntropy.levels.execution.rst
+++ /dev/null
@@ -1,22 +0,0 @@
-CodeEntropy.levels.execution package
-====================================
-
-Submodules
-----------
-
-.. toctree::
-   :maxdepth: 4
-
-   CodeEntropy.levels.execution.chunks
-   CodeEntropy.levels.execution.policy
-   CodeEntropy.levels.execution.reducers
-   CodeEntropy.levels.execution.scheduler
-   CodeEntropy.levels.execution.tasks
-
-Module contents
----------------
-
-.. automodule:: CodeEntropy.levels.execution
-   :members:
-   :show-inheritance:
-   :undoc-members:
diff --git a/docs/api/CodeEntropy.levels.execution.scheduler.rst b/docs/api/CodeEntropy.levels.execution.scheduler.rst
deleted file mode 100644
index b8a95196..00000000
--- a/docs/api/CodeEntropy.levels.execution.scheduler.rst
+++ /dev/null
@@ -1,7 +0,0 @@
-CodeEntropy.levels.execution.scheduler module
-=============================================
-
-.. automodule:: CodeEntropy.levels.execution.scheduler
-   :members:
-   :show-inheritance:
-   :undoc-members:
diff --git a/docs/api/CodeEntropy.levels.execution.tasks.rst b/docs/api/CodeEntropy.levels.execution.tasks.rst
deleted file mode 100644
index 9f2bda92..00000000
--- a/docs/api/CodeEntropy.levels.execution.tasks.rst
+++ /dev/null
@@ -1,7 +0,0 @@
-CodeEntropy.levels.execution.tasks module
-=========================================
-
-.. automodule:: CodeEntropy.levels.execution.tasks
-   :members:
-   :show-inheritance:
-   :undoc-members:
diff --git a/docs/api/CodeEntropy.levels.frame_dag.rst b/docs/api/CodeEntropy.levels.frame_dag.rst
new file mode 100644
index 00000000..9a747a87
--- /dev/null
+++ b/docs/api/CodeEntropy.levels.frame_dag.rst
@@ -0,0 +1,7 @@
+CodeEntropy.levels.frame\_dag module
+====================================
+
+.. automodule:: CodeEntropy.levels.frame_dag
+   :members:
+   :show-inheritance:
+   :undoc-members:
diff --git a/docs/api/CodeEntropy.levels.graph.conformation_dag.rst b/docs/api/CodeEntropy.levels.graph.conformation_dag.rst
deleted file mode 100644
index ee59e581..00000000
--- a/docs/api/CodeEntropy.levels.graph.conformation_dag.rst
+++ /dev/null
@@ -1,7 +0,0 @@
-CodeEntropy.levels.graph.conformation\_dag module
-=================================================
-
-.. automodule:: CodeEntropy.levels.graph.conformation_dag
-   :members:
-   :show-inheritance:
-   :undoc-members:
diff --git a/docs/api/CodeEntropy.levels.graph.frame_dag.rst b/docs/api/CodeEntropy.levels.graph.frame_dag.rst
deleted file mode 100644
index 361e5502..00000000
--- a/docs/api/CodeEntropy.levels.graph.frame_dag.rst
+++ /dev/null
@@ -1,7 +0,0 @@
-CodeEntropy.levels.graph.frame\_dag module
-==========================================
-
-.. automodule:: CodeEntropy.levels.graph.frame_dag
-   :members:
-   :show-inheritance:
-   :undoc-members:
diff --git a/docs/api/CodeEntropy.levels.graph.level_dag.rst b/docs/api/CodeEntropy.levels.graph.level_dag.rst
deleted file mode 100644
index a00e69b5..00000000
--- a/docs/api/CodeEntropy.levels.graph.level_dag.rst
+++ /dev/null
@@ -1,7 +0,0 @@
-CodeEntropy.levels.graph.level\_dag module
-==========================================
-
-.. automodule:: CodeEntropy.levels.graph.level_dag
-   :members:
-   :show-inheritance:
-   :undoc-members:
diff --git a/docs/api/CodeEntropy.levels.graph.rst b/docs/api/CodeEntropy.levels.graph.rst
deleted file mode 100644
index d4f75dbb..00000000
--- a/docs/api/CodeEntropy.levels.graph.rst
+++ /dev/null
@@ -1,20 +0,0 @@
-CodeEntropy.levels.graph package
-================================
-
-Submodules
-----------
-
-.. toctree::
-   :maxdepth: 4
-
-   CodeEntropy.levels.graph.conformation_dag
-   CodeEntropy.levels.graph.frame_dag
-   CodeEntropy.levels.graph.level_dag
-
-Module contents
----------------
-
-.. automodule:: CodeEntropy.levels.graph
-   :members:
-   :show-inheritance:
-   :undoc-members:
diff --git a/docs/api/CodeEntropy.levels.level_dag.rst b/docs/api/CodeEntropy.levels.level_dag.rst
new file mode 100644
index 00000000..78c08786
--- /dev/null
+++ b/docs/api/CodeEntropy.levels.level_dag.rst
@@ -0,0 +1,7 @@
+CodeEntropy.levels.level\_dag module
+====================================
+
+.. automodule:: CodeEntropy.levels.level_dag
+   :members:
+   :show-inheritance:
+   :undoc-members:
diff --git a/docs/api/CodeEntropy.levels.nodes.axes_topology.rst b/docs/api/CodeEntropy.levels.nodes.axes_topology.rst
deleted file mode 100644
index af80bc9e..00000000
--- a/docs/api/CodeEntropy.levels.nodes.axes_topology.rst
+++ /dev/null
@@ -1,7 +0,0 @@
-CodeEntropy.levels.nodes.axes\_topology module
-==============================================
-
-.. automodule:: CodeEntropy.levels.nodes.axes_topology
-   :members:
-   :show-inheritance:
-   :undoc-members:
diff --git a/docs/api/CodeEntropy.levels.nodes.rst b/docs/api/CodeEntropy.levels.nodes.rst
index 048f7fcf..4499f939 100644
--- a/docs/api/CodeEntropy.levels.nodes.rst
+++ b/docs/api/CodeEntropy.levels.nodes.rst
@@ -1,6 +1,11 @@
 CodeEntropy.levels.nodes package
 ================================
 
+.. automodule:: CodeEntropy.levels.nodes
+   :members:
+   :show-inheritance:
+   :undoc-members:
+
 Submodules
 ----------
 
@@ -8,16 +13,7 @@ Submodules
    :maxdepth: 4
 
    CodeEntropy.levels.nodes.accumulators
-   CodeEntropy.levels.nodes.axes_topology
    CodeEntropy.levels.nodes.beads
    CodeEntropy.levels.nodes.covariance
    CodeEntropy.levels.nodes.detect_levels
    CodeEntropy.levels.nodes.detect_molecules
-
-Module contents
----------------
-
-.. automodule:: CodeEntropy.levels.nodes
-   :members:
-   :show-inheritance:
-   :undoc-members:
diff --git a/docs/api/CodeEntropy.levels.rst b/docs/api/CodeEntropy.levels.rst
index 9438f461..cca6398d 100644
--- a/docs/api/CodeEntropy.levels.rst
+++ b/docs/api/CodeEntropy.levels.rst
@@ -1,15 +1,17 @@
 CodeEntropy.levels package
 ==========================
 
+.. automodule:: CodeEntropy.levels
+   :members:
+   :show-inheritance:
+   :undoc-members:
+
 Subpackages
 -----------
 
 .. toctree::
    :maxdepth: 4
 
-   CodeEntropy.levels.dihedrals
-   CodeEntropy.levels.execution
-   CodeEntropy.levels.graph
    CodeEntropy.levels.nodes
 
 Submodules
@@ -19,16 +21,12 @@ Submodules
    :maxdepth: 4
 
    CodeEntropy.levels.axes
+   CodeEntropy.levels.conformation_dag
+   CodeEntropy.levels.dihedrals
    CodeEntropy.levels.forces
+   CodeEntropy.levels.frame_dag
    CodeEntropy.levels.hierarchy
+   CodeEntropy.levels.level_dag
    CodeEntropy.levels.linalg
    CodeEntropy.levels.neighbors
    CodeEntropy.levels.search
-
-Module contents
----------------
-
-.. automodule:: CodeEntropy.levels
-   :members:
-   :show-inheritance:
-   :undoc-members:
diff --git a/docs/api/CodeEntropy.molecules.rst b/docs/api/CodeEntropy.molecules.rst
index 3dd473be..af25747e 100644
--- a/docs/api/CodeEntropy.molecules.rst
+++ b/docs/api/CodeEntropy.molecules.rst
@@ -1,6 +1,11 @@
 CodeEntropy.molecules package
 =============================
 
+.. automodule:: CodeEntropy.molecules
+   :members:
+   :show-inheritance:
+   :undoc-members:
+
 Submodules
 ----------
 
@@ -8,11 +13,3 @@ Submodules
    :maxdepth: 4
 
    CodeEntropy.molecules.grouping
-
-Module contents
----------------
-
-.. automodule:: CodeEntropy.molecules
-   :members:
-   :show-inheritance:
-   :undoc-members:
diff --git a/docs/api/CodeEntropy.results.rst b/docs/api/CodeEntropy.results.rst
index 2606aaf4..40d4eee2 100644
--- a/docs/api/CodeEntropy.results.rst
+++ b/docs/api/CodeEntropy.results.rst
@@ -1,6 +1,11 @@
 CodeEntropy.results package
 ===========================
 
+.. automodule:: CodeEntropy.results
+   :members:
+   :show-inheritance:
+   :undoc-members:
+
 Submodules
 ----------
 
@@ -8,11 +13,3 @@ Submodules
    :maxdepth: 4
 
    CodeEntropy.results.reporter
-
-Module contents
----------------
-
-.. automodule:: CodeEntropy.results
-   :members:
-   :show-inheritance:
-   :undoc-members:
diff --git a/docs/api/CodeEntropy.rst b/docs/api/CodeEntropy.rst
index 2fb3c5bb..ae75c621 100644
--- a/docs/api/CodeEntropy.rst
+++ b/docs/api/CodeEntropy.rst
@@ -1,6 +1,11 @@
 CodeEntropy package
 ===================
 
+.. automodule:: CodeEntropy
+   :members:
+   :show-inheritance:
+   :undoc-members:
+
 Subpackages
 -----------
 
@@ -22,11 +27,3 @@ Submodules
    :maxdepth: 4
 
    CodeEntropy.cli
-
-Module contents
----------------
-
-.. automodule:: CodeEntropy
-   :members:
-   :show-inheritance:
-   :undoc-members:
diff --git a/docs/api/CodeEntropy.trajectory.rst b/docs/api/CodeEntropy.trajectory.rst
index 2ac5a826..14005c6c 100644
--- a/docs/api/CodeEntropy.trajectory.rst
+++ b/docs/api/CodeEntropy.trajectory.rst
@@ -1,6 +1,11 @@
 CodeEntropy.trajectory package
 ==============================
 
+.. automodule:: CodeEntropy.trajectory
+   :members:
+   :show-inheritance:
+   :undoc-members:
+
 Submodules
 ----------
 
@@ -10,11 +15,3 @@ Submodules
    CodeEntropy.trajectory.frames
    CodeEntropy.trajectory.mda
    CodeEntropy.trajectory.source
-
-Module contents
----------------
-
-.. automodule:: CodeEntropy.trajectory
-   :members:
-   :show-inheritance:
-   :undoc-members:
diff --git a/docs/getting_started.rst b/docs/getting_started.rst
index 192323cc..6ddc6a21 100644
--- a/docs/getting_started.rst
+++ b/docs/getting_started.rst
@@ -515,7 +515,7 @@ Example 1: DNA Fragment (Smallest / Fastest)
 
 Data files:
 
-`DNA fragment example (~1MB) <https://file-store.ccpbiosim.org/codeentropy-examples/dna_example.tar>`_
+`DNA fragment example (~1MB) <https://ccpbiosim.ac.uk/file-store/codeentropy-examples/dna_example.tar>`_
 
 Create or edit ``config.yaml`` in your working directory:
 
@@ -551,7 +551,7 @@ Example 2: Lysozyme (Larger / Slower)
 
 Data files:
 
-`Lysozyme example (~1.2GB) <https://file-store.ccpbiosim.org/codeentropy-examples/lysozyme_example.tar>`_
+`Lysozyme example (~1.2GB) <https://ccpbiosim.ac.uk/file-store/codeentropy-examples/lysozyme_example.tar>`_
 
 Create or edit ``config.yaml`` in your working directory:
 
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/default.json b/tests/regression/baselines/benzaldehyde/default.json
index c99a2d15..1ed1a4e7 100644
--- a/tests/regression/baselines/benzaldehyde/default.json
+++ b/tests/regression/baselines/benzaldehyde/default.json
@@ -2,7 +2,7 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 24.88518233240474,
+        "united_atom:Transvibrational": 24.869562693152986,
         "united_atom:Rovibrational": 27.950376507672583,
         "residue:FTmat-Transvibrational": 71.03412922724692,
         "residue:FTmat-Rovibrational": 59.44169664956799,
@@ -10,7 +10,7 @@
         "residue:Conformational": 0.0,
         "residue:Orientational": 59.43920971558428
       },
-      "total": 250.27323833217878
+      "total": 250.25761869292705
     }
   }
 }
diff --git a/tests/regression/baselines/benzaldehyde/frame_window.json b/tests/regression/baselines/benzaldehyde/frame_window.json
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/rad.json b/tests/regression/baselines/benzaldehyde/rad.json
index 8e39e21b..02d56d59 100644
--- a/tests/regression/baselines/benzaldehyde/rad.json
+++ b/tests/regression/baselines/benzaldehyde/rad.json
@@ -2,7 +2,7 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 24.88518233240474,
+        "united_atom:Transvibrational": 24.869562693152986,
         "united_atom:Rovibrational": 27.950376507672583,
         "residue:FTmat-Transvibrational": 71.03412922724692,
         "residue:FTmat-Rovibrational": 59.44169664956799,
@@ -10,7 +10,7 @@
         "residue:Conformational": 0.0,
         "residue:Orientational": 25.502722133228936
       },
-      "total": 216.33675074982344
+      "total": 216.32113111057168
     }
   }
 }
diff --git a/tests/regression/baselines/benzaldehyde/selection_subset.json b/tests/regression/baselines/benzaldehyde/selection_subset.json
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/default.json b/tests/regression/baselines/benzene/default.json
index 86ae25d0..de6eebc0 100644
--- a/tests/regression/baselines/benzene/default.json
+++ b/tests/regression/baselines/benzene/default.json
@@ -2,7 +2,7 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 9.081853431559052,
+        "united_atom:Transvibrational": 9.050489789631982,
         "united_atom:Rovibrational": 27.534380126250227,
         "residue:FTmat-Transvibrational": 72.66211800013413,
         "residue:FTmat-Rovibrational": 59.93761874375924,
@@ -10,7 +10,7 @@
         "residue:Conformational": 0.0,
         "residue:Orientational": 41.58427729275938
       },
-      "total": 210.800247594462
+      "total": 210.76888395253496
     }
   }
 }
diff --git a/tests/regression/baselines/benzene/frame_window.json b/tests/regression/baselines/benzene/frame_window.json
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/rad.json b/tests/regression/baselines/benzene/rad.json
index c4aca55f..2ccf3306 100644
--- a/tests/regression/baselines/benzene/rad.json
+++ b/tests/regression/baselines/benzene/rad.json
@@ -2,7 +2,7 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 9.081853431559052,
+        "united_atom:Transvibrational": 9.050489789631982,
         "united_atom:Rovibrational": 27.534380126250227,
         "residue:FTmat-Transvibrational": 72.66211800013413,
         "residue:FTmat-Rovibrational": 59.93761874375924,
@@ -10,7 +10,7 @@
         "residue:Conformational": 0.0,
         "residue:Orientational": 12.386543820026672
       },
-      "total": 181.6025141217293
+      "total": 181.57115047980224
     }
   }
 }
diff --git a/tests/regression/baselines/benzene/selection_subset.json b/tests/regression/baselines/benzene/selection_subset.json
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/default.json b/tests/regression/baselines/cyclohexane/default.json
index c0233178..e7867e50 100644
--- a/tests/regression/baselines/cyclohexane/default.json
+++ b/tests/regression/baselines/cyclohexane/default.json
@@ -2,7 +2,7 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 11.726471858700231,
+        "united_atom:Transvibrational": 11.65610486557721,
         "united_atom:Rovibrational": 47.71088602161552,
         "residue:FTmat-Transvibrational": 70.3922327806972,
         "residue:FTmat-Rovibrational": 63.44920824648768,
@@ -10,7 +10,7 @@
         "residue:Conformational": 0.0,
         "residue:Orientational": 47.54362774624537
       },
-      "total": 242.11130586044152
+      "total": 242.0409388673185
     }
   }
 }
diff --git a/tests/regression/baselines/cyclohexane/frame_window.json b/tests/regression/baselines/cyclohexane/frame_window.json
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/rad.json b/tests/regression/baselines/cyclohexane/rad.json
index f65d9a4c..8de06e5b 100644
--- a/tests/regression/baselines/cyclohexane/rad.json
+++ b/tests/regression/baselines/cyclohexane/rad.json
@@ -2,7 +2,7 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 11.726471858700231,
+        "united_atom:Transvibrational": 11.65610486557721,
         "united_atom:Rovibrational": 47.71088602161552,
         "residue:FTmat-Transvibrational": 70.3922327806972,
         "residue:FTmat-Rovibrational": 63.44920824648768,
@@ -10,7 +10,7 @@
         "residue:Conformational": 0.0,
         "residue:Orientational": 14.049324060636614
       },
-      "total": 208.61700217483278
+      "total": 208.54663518170975
     }
   }
 }
diff --git a/tests/regression/baselines/cyclohexane/selection_subset.json b/tests/regression/baselines/cyclohexane/selection_subset.json
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..5e724622 100644
--- a/tests/regression/baselines/dna/combined_forcetorque_off.json
+++ b/tests/regression/baselines/dna/combined_forcetorque_off.json
@@ -2,17 +2,17 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 0.0,
+        "united_atom:Transvibrational": 4.762281610623415e-20,
         "united_atom:Rovibrational": 0.002160679012128457,
         "residue:Transvibrational": 0.0,
-        "residue:Rovibrational": 3.376800684085249,
+        "residue:Rovibrational": 6.633599673254765,
         "polymer:Transvibrational": 21.18266215491188,
         "polymer:Rovibrational": 12.837576042626923,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "polymer:Orientational": 4.758905336627712
       },
-      "total": 42.1581048972639
+      "total": 45.41490388643341
     },
     "1": {
       "components": {
diff --git a/tests/regression/baselines/dna/default.json b/tests/regression/baselines/dna/default.json
index 2cbef740..7a2b6c3a 100644
--- a/tests/regression/baselines/dna/default.json
+++ b/tests/regression/baselines/dna/default.json
@@ -2,17 +2,17 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 0.0,
+        "united_atom:Transvibrational": 4.762281610623415e-20,
         "united_atom:Rovibrational": 0.002160679012128457,
         "residue:Transvibrational": 0.0,
-        "residue:Rovibrational": 3.376800684085249,
+        "residue:Rovibrational": 6.633599673254765,
         "polymer: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": {
diff --git a/tests/regression/baselines/dna/frame_window.json b/tests/regression/baselines/dna/frame_window.json
index c35d2211..fd24e0e2 100644
--- a/tests/regression/baselines/dna/frame_window.json
+++ b/tests/regression/baselines/dna/frame_window.json
@@ -5,18 +5,18 @@
         "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,
@@ -26,7 +26,7 @@
         "residue:Conformational": 0.0,
         "polymer:Orientational": 4.758905336627712
       },
-      "total": 97.85965049646045
+      "total": 97.85965049699642
     }
   }
 }
diff --git a/tests/regression/baselines/dna/grouping_each.json b/tests/regression/baselines/dna/grouping_each.json
index 2cbef740..7a2b6c3a 100644
--- a/tests/regression/baselines/dna/grouping_each.json
+++ b/tests/regression/baselines/dna/grouping_each.json
@@ -2,17 +2,17 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 0.0,
+        "united_atom:Transvibrational": 4.762281610623415e-20,
         "united_atom:Rovibrational": 0.002160679012128457,
         "residue:Transvibrational": 0.0,
-        "residue:Rovibrational": 3.376800684085249,
+        "residue:Rovibrational": 6.633599673254765,
         "polymer: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": {
diff --git a/tests/regression/baselines/dna/selection_subset.json b/tests/regression/baselines/dna/selection_subset.json
index 2cbef740..7a2b6c3a 100644
--- a/tests/regression/baselines/dna/selection_subset.json
+++ b/tests/regression/baselines/dna/selection_subset.json
@@ -2,17 +2,17 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 0.0,
+        "united_atom:Transvibrational": 4.762281610623415e-20,
         "united_atom:Rovibrational": 0.002160679012128457,
         "residue:Transvibrational": 0.0,
-        "residue:Rovibrational": 3.376800684085249,
+        "residue:Rovibrational": 6.633599673254765,
         "polymer: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": {
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/default.json b/tests/regression/baselines/ethyl-acetate/default.json
index 077cdbe3..a98d4af8 100644
--- a/tests/regression/baselines/ethyl-acetate/default.json
+++ b/tests/regression/baselines/ethyl-acetate/default.json
@@ -2,7 +2,7 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 19.703248485425714,
+        "united_atom:Transvibrational": 19.67356252280246,
         "united_atom:Rovibrational": 49.35007067210558,
         "residue:FTmat-Transvibrational": 67.91350627765567,
         "residue:FTmat-Rovibrational": 60.042233035077246,
@@ -10,7 +10,7 @@
         "residue:Conformational": 0.0,
         "residue:Orientational": 65.33877737416202
       },
-      "total": 270.52337372549465
+      "total": 270.4936877628714
     }
   }
 }
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/rad.json b/tests/regression/baselines/ethyl-acetate/rad.json
index ee7df624..51bed453 100644
--- a/tests/regression/baselines/ethyl-acetate/rad.json
+++ b/tests/regression/baselines/ethyl-acetate/rad.json
@@ -2,7 +2,7 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 19.703248485425714,
+        "united_atom:Transvibrational": 19.67356252280246,
         "united_atom:Rovibrational": 49.35007067210558,
         "residue:FTmat-Transvibrational": 67.91350627765567,
         "residue:FTmat-Rovibrational": 60.042233035077246,
@@ -10,7 +10,7 @@
         "residue:Conformational": 0.0,
         "residue:Orientational": 30.739736508673403
       },
-      "total": 235.92433286000602
+      "total": 235.89464689738276
     }
   }
 }
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/default.json b/tests/regression/baselines/methane/default.json
index d08ac60f..c6716b00 100644
--- a/tests/regression/baselines/methane/default.json
+++ b/tests/regression/baselines/methane/default.json
@@ -2,12 +2,12 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 40.3239711717637,
-        "united_atom:Rovibrational": 33.60582165153992,
+        "united_atom:Transvibrational": 40.32397117176369,
+        "united_atom:Rovibrational": 33.617646190915714,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 3.596676624528629
       },
-      "total": 77.52646944783224
+      "total": 77.53829398720804
     }
   }
 }
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/rad.json b/tests/regression/baselines/methane/rad.json
index df8fbbd7..2b506858 100644
--- a/tests/regression/baselines/methane/rad.json
+++ b/tests/regression/baselines/methane/rad.json
@@ -2,12 +2,12 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 40.3239711717637,
-        "united_atom:Rovibrational": 33.60582165153992,
+        "united_atom:Transvibrational": 40.32397117176369,
+        "united_atom:Rovibrational": 33.617646190915714,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 6.667209510980459
       },
-      "total": 80.59700233428407
+      "total": 80.60882687365986
     }
   }
 }
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/default.json b/tests/regression/baselines/octonol/default.json
index fdd22382..76b4b259 100644
--- a/tests/regression/baselines/octonol/default.json
+++ b/tests/regression/baselines/octonol/default.json
@@ -2,7 +2,7 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 37.860034170693694,
+        "united_atom:Transvibrational": 37.90014130683313,
         "united_atom:Rovibrational": 84.2608785308744,
         "residue:FTmat-Transvibrational": 66.13571365167344,
         "residue:FTmat-Rovibrational": 57.090651827515686,
@@ -10,7 +10,7 @@
         "residue:Conformational": 0.0,
         "residue:Orientational": 74.78356386019003
       },
-      "total": 354.36024401601327
+      "total": 354.4003511521527
     }
   }
 }
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/rad.json b/tests/regression/baselines/octonol/rad.json
index c2c20839..1d294585 100644
--- a/tests/regression/baselines/octonol/rad.json
+++ b/tests/regression/baselines/octonol/rad.json
@@ -2,7 +2,7 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 37.860034170693694,
+        "united_atom:Transvibrational": 37.90014130683313,
         "united_atom:Rovibrational": 84.2608785308744,
         "residue:FTmat-Transvibrational": 66.13571365167344,
         "residue:FTmat-Rovibrational": 57.090651827515686,
@@ -10,7 +10,7 @@
         "residue:Conformational": 0.0,
         "residue:Orientational": 28.084239413273323
       },
-      "total": 307.66091956909656
+      "total": 307.701026705236
     }
   }
 }
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/default.json b/tests/regression/baselines/water/default.json
index c98bc3eb..7f5baa09 100644
--- a/tests/regression/baselines/water/default.json
+++ b/tests/regression/baselines/water/default.json
@@ -3,11 +3,11 @@
     "0": {
       "components": {
         "united_atom:Transvibrational": 43.725393337304425,
-        "united_atom:Rovibrational": 17.28911070147887,
+        "united_atom:Rovibrational": 17.42448032148262,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 32.193445593362114
       },
-      "total": 93.20794963214541
+      "total": 93.34331925214916
     }
   }
 }
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/rad.json b/tests/regression/baselines/water/rad.json
index cbda37a9..cf8510e3 100644
--- a/tests/regression/baselines/water/rad.json
+++ b/tests/regression/baselines/water/rad.json
@@ -2,12 +2,12 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 79.20298312418278,
-        "united_atom:Rovibrational": 50.90260688502127,
+        "united_atom:Transvibrational": 79.2029831241828,
+        "united_atom:Rovibrational": 51.07316520629525,
         "united_atom:Conformational": 0.0,
         "united_atom:Orientational": 22.760377489372107
       },
-      "total": 152.86596749857614
+      "total": 153.03652581985014
     }
   }
 }
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/regression/helpers.py b/tests/regression/helpers.py
index d5c09a12..9f27b41b 100644
--- a/tests/regression/helpers.py
+++ b/tests/regression/helpers.py
@@ -11,7 +11,7 @@
 
 import yaml
 
-DEFAULT_TESTDATA_BASE_URL = "https://file-store.ccpbiosim.org/codeentropy-testing"
+DEFAULT_TESTDATA_BASE_URL = "https://www.ccpbiosim.ac.uk/file-store/codeentropy-testing"
 
 
 @dataclass(frozen=True)
diff --git a/tests/unit/CodeEntropy/levels/dihedrals/test_topology.py b/tests/unit/CodeEntropy/levels/dihedrals/test_topology.py
index cbc38b19..60affa71 100644
--- a/tests/unit/CodeEntropy/levels/dihedrals/test_topology.py
+++ b/tests/unit/CodeEntropy/levels/dihedrals/test_topology.py
@@ -1,472 +1,146 @@
 from __future__ import annotations
 
-import logging
-from dataclasses import dataclass
-from types import SimpleNamespace
-from typing import Any
-from unittest.mock import Mock, call
+from unittest.mock import MagicMock
 
-import pytest
+import numpy as np
 
-from CodeEntropy.levels.dihedrals.topology import (
-    DihedralTopologyDiscovery,
-    MoleculeDihedralTopology,
-)
+from CodeEntropy.levels.dihedrals.topology import DihedralTopologyDiscovery
 
 
-@dataclass
-class FakeAtom:
-    """Small test double for an MDAnalysis Atom."""
+class _AddableAG:
+    """Minimal addable AtomGroup test double."""
 
-    index: int
-    bonded_atoms: Any = None
+    def __init__(self, name: str) -> None:
+        """Initialize the fake AtomGroup.
 
+        Args:
+            name: Human-readable identifier used in composed names.
+        """
+        self.name = name
 
-class FakeUniverseAtoms:
-    """Indexable fake universe atom container."""
+    def __add__(self, other: _AddableAG) -> _AddableAG:
+        """Return a composed fake AtomGroup.
 
-    def __init__(self, universe: FakeUniverse) -> None:
-        self._universe = universe
+        Args:
+            other: Fake AtomGroup to combine with this object.
 
-    def __getitem__(self, indices: int | list[int] | tuple[int, ...]) -> Any:
-        if isinstance(indices, int):
-            return self._universe.atom_by_index[int(indices)]
+        Returns:
+            New fake AtomGroup containing a composed name.
+        """
+        return _AddableAG(f"({self.name}+{other.name})")
 
-        return FakeAtomGroup(
-            [self._universe.atom_by_index[int(index)] for index in indices],
-            universe=self._universe,
-        )
 
+class _TopologyDiscovery(DihedralTopologyDiscovery):
+    """Concrete topology-discovery helper for unit tests."""
 
-class FakeUniverse:
-    """Small fake universe supporting ``universe.atoms[indices]``."""
+    def __init__(self, universe_operations: MagicMock) -> None:
+        """Initialize the test helper.
 
-    def __init__(self, atoms: list[FakeAtom]) -> None:
-        self.atom_by_index = {int(atom.index): atom for atom in atoms}
-        self.atoms = FakeUniverseAtoms(self)
+        Args:
+            universe_operations: Mock universe-operation adapter.
+        """
+        self._universe_operations = universe_operations
 
 
-class FakeAtomGroup:
-    """Small AtomGroup-like test double for topology discovery tests."""
+def test_select_heavy_residue_builds_expected_selections():
+    uops = MagicMock()
+    helper = _TopologyDiscovery(universe_operations=uops)
 
-    def __init__(
-        self,
-        atoms: list[FakeAtom],
-        *,
-        residues: list[Any] | None = None,
-        dihedrals: list[Any] | None = None,
-        select_map: dict[str, Any] | None = None,
-        universe: FakeUniverse | None = None,
-    ) -> None:
-        self._atoms = list(atoms)
-        self.residues = list(residues or [])
-        self.dihedrals = list(dihedrals or [])
-        self._select_map = dict(select_map or {})
-        self.universe = universe
+    mol = MagicMock()
+    mol.residues = [MagicMock()]
+    mol.residues[0].atoms.indices = np.array([10, 11, 12], dtype=int)
+    uops.select_atoms.side_effect = ["residue_atoms", "heavy_atoms"]
 
-    @property
-    def atoms(self) -> FakeAtomGroup:
-        return self
-
-    @property
-    def indices(self) -> list[int]:
-        return [int(atom.index) for atom in self._atoms]
-
-    def __iter__(self):
-        return iter(self._atoms)
-
-    def __len__(self) -> int:
-        return len(self._atoms)
-
-    def __add__(self, other: FakeAtomGroup) -> FakeAtomGroup:
-        return FakeAtomGroup(
-            self._atoms + other._atoms,
-            universe=self.universe or other.universe,
-        )
-
-    def select_atoms(self, select_string: str, updating: bool = False) -> Any:
-        if select_string not in self._select_map:
-            raise AssertionError(f"Unexpected selection: {select_string!r}")
-        return self._select_map[select_string]
-
-
-@dataclass
-class FakeResidue:
-    """Small residue test double."""
-
-    atoms: FakeAtomGroup
-
-
-@dataclass
-class FakeDihedral:
-    """Small dihedral topology object."""
-
-    atoms: FakeAtomGroup
-
-
-def _make_discovery(
-    universe_operations: Any | None = None,
-) -> DihedralTopologyDiscovery:
-    discovery = DihedralTopologyDiscovery()
-    discovery._universe_operations = universe_operations or Mock()
-    return discovery
-
-
-def test_molecule_dihedral_topology_stores_expected_fields() -> None:
-    topology = MoleculeDihedralTopology(
-        group_id=1,
-        molecule_id=2,
-        molecule_order=3,
-        num_residues=4,
-        ua_dihedrals_by_residue={0: ["ua"]},
-        residue_dihedrals=["res"],
-    )
-
-    assert topology.group_id == 1
-    assert topology.molecule_id == 2
-    assert topology.molecule_order == 3
-    assert topology.num_residues == 4
-    assert topology.ua_dihedrals_by_residue == {0: ["ua"]}
-    assert topology.residue_dihedrals == ["res"]
-
-
-def test_extract_topology_fragment_uses_lightweight_atomgroup_helper() -> None:
-    universe_operations = Mock()
-    universe_operations.extract_fragment_atomgroup.return_value = "fragment_atomgroup"
-    universe_operations.extract_fragment.return_value = "heavy_fragment"
-
-    discovery = _make_discovery(universe_operations)
-
-    result = discovery._extract_topology_fragment("universe", 5)
-
-    assert result == "fragment_atomgroup"
-    universe_operations.extract_fragment_atomgroup.assert_called_once_with(
-        "universe",
-        5,
-    )
-    universe_operations.extract_fragment.assert_not_called()
-
-
-def test_select_heavy_residue_builds_expected_selections() -> None:
-    discovery = _make_discovery()
-
-    residue_atoms = Mock()
-    residue_atoms.indices = [10, 11, 12, 13]
-    residue = SimpleNamespace(atoms=residue_atoms)
-
-    heavy_atoms = object()
-    residue_container = Mock()
-    residue_container.select_atoms.return_value = heavy_atoms
-
-    mol = Mock()
-    mol.residues = [residue]
-    mol.select_atoms.return_value = residue_container
-
-    result = discovery._select_heavy_residue(mol, 0)
-
-    assert result is heavy_atoms
-    mol.select_atoms.assert_called_once_with("index 10:13", updating=False)
-    residue_container.select_atoms.assert_called_once_with(
-        "prop mass > 1.1",
-        updating=False,
-    )
-    discovery._universe_operations.select_atoms.assert_not_called()
-
-
-def test_get_dihedrals_united_atom_collects_atoms_from_dihedral_objects() -> None:
-    discovery = _make_discovery()
-
-    atoms = [FakeAtom(index) for index in range(1, 8)]
-    universe = FakeUniverse(atoms)
-
-    valid_dihedral_atoms = FakeAtomGroup(
-        [atoms[0], atoms[1], atoms[2], atoms[3]],
-        universe=universe,
-    )
-    outside_selection_atoms = FakeAtomGroup(
-        [atoms[0], atoms[1], atoms[2], atoms[6]],
-        universe=universe,
-    )
-    wrong_size_atoms = FakeAtomGroup(
-        [atoms[0], atoms[1], atoms[2]],
-        universe=universe,
-    )
-
-    selected_heavy_atoms = FakeAtomGroup(
-        [atoms[0], atoms[1], atoms[2], atoms[3], atoms[4]],
-        dihedrals=[
-            FakeDihedral(valid_dihedral_atoms),
-            FakeDihedral(outside_selection_atoms),
-            FakeDihedral(wrong_size_atoms),
-        ],
-        universe=universe,
-    )
-
-    result = discovery._get_dihedrals(selected_heavy_atoms, "united_atom")
-
-    assert result == [valid_dihedral_atoms]
-
-
-def test_get_dihedrals_united_atom_returns_empty_when_no_valid_dihedrals() -> None:
-    discovery = _make_discovery()
-
-    atoms = [FakeAtom(index) for index in range(1, 5)]
-    universe = FakeUniverse(atoms)
-
-    invalid_dihedral_atoms = FakeAtomGroup(
-        [atoms[0], atoms[1], atoms[2]],
-        universe=universe,
-    )
-    selected_heavy_atoms = FakeAtomGroup(
-        atoms,
-        dihedrals=[FakeDihedral(invalid_dihedral_atoms)],
-        universe=universe,
-    )
-
-    assert discovery._get_dihedrals(selected_heavy_atoms, "united_atom") == []
-
-
-def test_get_dihedrals_residue_builds_one_dihedral_when_four_residues() -> None:
-    discovery = _make_discovery()
-
-    atom1 = FakeAtom(10)
-    atom2 = FakeAtom(20)
-    atom3 = FakeAtom(30)
-    atom4 = FakeAtom(40)
-
-    universe = FakeUniverse([atom1, atom2, atom3, atom4])
-
-    group1 = FakeAtomGroup([atom1], universe=universe)
-    group2 = FakeAtomGroup([atom2], universe=universe)
-    group3 = FakeAtomGroup([atom3], universe=universe)
-    group4 = FakeAtomGroup([atom4], universe=universe)
-
-    atom1.bonded_atoms = group2
-    atom2.bonded_atoms = group1
-    atom3.bonded_atoms = group4
-    atom4.bonded_atoms = group3
-
-    mol = FakeAtomGroup(
-        [atom1, atom2, atom3, atom4],
-        residues=[
-            FakeResidue(group1),
-            FakeResidue(group2),
-            FakeResidue(group3),
-            FakeResidue(group4),
-        ],
-        universe=universe,
-    )
-
-    result = discovery._get_dihedrals(mol, "residue")
-
-    assert len(result) == 1
-    assert result[0].indices == [10, 20, 30, 40]
-
-
-def test_get_dihedrals_residue_skips_invalid_four_residue_window(
-    caplog: pytest.LogCaptureFixture,
-) -> None:
-    caplog.set_level(logging.DEBUG)
-
-    discovery = _make_discovery()
-
-    atom1 = FakeAtom(10)
-    atom2 = FakeAtom(20)
-    atom3 = FakeAtom(30)
-    atom4 = FakeAtom(40)
-
-    universe = FakeUniverse([atom1, atom2, atom3, atom4])
-
-    group1 = FakeAtomGroup([atom1], universe=universe)
-    group2 = FakeAtomGroup([atom2], universe=universe)
-    group3 = FakeAtomGroup([atom3], universe=universe)
-    group4 = FakeAtomGroup([atom4], universe=universe)
-
-    atom1.bonded_atoms = group2
-    atom2.bonded_atoms = group1
-    atom3.bonded_atoms = FakeAtomGroup([], universe=universe)
-    atom4.bonded_atoms = FakeAtomGroup([], universe=universe)
-
-    mol = FakeAtomGroup(
-        [atom1, atom2, atom3, atom4],
-        residues=[
-            FakeResidue(group1),
-            FakeResidue(group2),
-            FakeResidue(group3),
-            FakeResidue(group4),
-        ],
-        universe=universe,
-    )
-
-    result = discovery._get_dihedrals(mol, "residue")
-
-    assert result == []
-    assert "Skipping residue-level dihedral" in caplog.text
-
-
-def test_get_dihedrals_residue_returns_empty_when_fewer_than_four_residues() -> None:
-    discovery = _make_discovery()
-
-    atoms = [FakeAtom(1), FakeAtom(2), FakeAtom(3)]
-    universe = FakeUniverse(atoms)
-    groups = [FakeAtomGroup([atom], universe=universe) for atom in atoms]
-
-    mol = FakeAtomGroup(
-        atoms,
-        residues=[FakeResidue(group) for group in groups],
-        universe=universe,
-    )
-
-    assert discovery._get_dihedrals(mol, "residue") == []
+    out = helper._select_heavy_residue(mol, res_id=0)
 
+    assert out == "heavy_atoms"
+    assert uops.select_atoms.call_args_list == [
+        ((mol, "index 10:12"),),
+        (("residue_atoms", "prop mass > 1.1"),),
+    ]
 
-def test_atoms_in_source_bonded_to_target_returns_matching_source_atoms() -> None:
-    atom1 = FakeAtom(1)
-    atom2 = FakeAtom(2)
-    atom3 = FakeAtom(3)
 
-    universe = FakeUniverse([atom1, atom2, atom3])
+def test_get_dihedrals_united_atom_collects_atoms_from_dihedral_objects():
+    helper = _TopologyDiscovery(universe_operations=MagicMock())
 
-    source_group = FakeAtomGroup([atom1, atom3], universe=universe)
-    target_group = FakeAtomGroup([atom2], universe=universe)
+    d0 = MagicMock()
+    d0.atoms = "A0"
+    d1 = MagicMock()
+    d1.atoms = "A1"
 
-    atom1.bonded_atoms = target_group
-    atom3.bonded_atoms = FakeAtomGroup([], universe=universe)
+    container = MagicMock()
+    container.dihedrals = [d0, d1]
 
-    source_residue = FakeResidue(source_group)
-    target_residue = FakeResidue(target_group)
+    assert helper._get_dihedrals(container, level="united_atom") == ["A0", "A1"]
 
-    result = DihedralTopologyDiscovery._atoms_in_source_bonded_to_target(
-        source_residue,
-        target_residue,
-    )
 
-    assert result.indices == [1]
+def test_get_dihedrals_residue_returns_empty_when_less_than_four_residues():
+    helper = _TopologyDiscovery(universe_operations=MagicMock())
 
+    mol = MagicMock()
+    mol.residues = [MagicMock(), MagicMock(), MagicMock()]
+    mol.select_atoms = MagicMock()
 
-def test_atoms_in_source_bonded_to_target_returns_empty_when_atom_has_no_bonds() -> (
-    None
-):
-    atom1 = FakeAtom(1)
-    atom2 = FakeAtom(2)
+    assert helper._get_dihedrals(mol, level="residue") == []
+    mol.select_atoms.assert_not_called()
 
-    universe = FakeUniverse([atom1, atom2])
 
-    source_residue = FakeResidue(FakeAtomGroup([atom1], universe=universe))
-    target_residue = FakeResidue(FakeAtomGroup([atom2], universe=universe))
+def test_get_dihedrals_residue_builds_one_dihedral_when_four_residues():
+    helper = _TopologyDiscovery(universe_operations=MagicMock())
 
-    result = DihedralTopologyDiscovery._atoms_in_source_bonded_to_target(
-        source_residue,
-        target_residue,
+    mol = MagicMock()
+    mol.residues = [MagicMock(), MagicMock(), MagicMock(), MagicMock()]
+    mol.select_atoms = MagicMock(
+        side_effect=[
+            _AddableAG("a1"),
+            _AddableAG("a2"),
+            _AddableAG("a3"),
+            _AddableAG("a4"),
+        ]
     )
 
-    assert result.indices == []
-
+    out = helper._get_dihedrals(mol, level="residue")
 
-def test_discover_group_dihedral_topology_builds_one_entry_per_molecule(
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    universe_operations = Mock()
+    assert len(out) == 1
+    assert isinstance(out[0], _AddableAG)
+    assert mol.select_atoms.call_count == 4
 
-    molecule0 = SimpleNamespace(label="molecule0", residues=[object(), object()])
-    molecule1 = SimpleNamespace(label="molecule1", residues=[object()])
 
-    universe_operations.extract_fragment_atomgroup.side_effect = [
-        molecule0,
-        molecule1,
-    ]
-
-    discovery = _make_discovery(universe_operations)
-
-    def fake_select_heavy_residue(mol: Any, res_id: int) -> str:
-        return f"{mol.label}-heavy-{res_id}"
+def test_discover_group_dihedral_topology_builds_one_entry_per_molecule():
+    uops = MagicMock()
+    helper = _TopologyDiscovery(universe_operations=uops)
 
-    def fake_get_dihedrals(data_container: Any, level: str) -> list[str]:
-        if level == "united_atom":
-            return [f"ua-{data_container}"]
-        return [f"res-{data_container.label}"]
+    mol0 = MagicMock()
+    mol0.residues = [MagicMock(), MagicMock()]
+    mol1 = MagicMock()
+    mol1.residues = [MagicMock(), MagicMock()]
+    uops.extract_fragment.side_effect = [mol0, mol1]
 
-    monkeypatch.setattr(
-        discovery,
-        "_select_heavy_residue",
-        fake_select_heavy_residue,
+    helper._select_heavy_residue = MagicMock(
+        side_effect=["heavy0", "heavy1", "heavy2", "heavy3"]
     )
-    monkeypatch.setattr(
-        discovery,
-        "_get_dihedrals",
-        fake_get_dihedrals,
+    helper._get_dihedrals = MagicMock(
+        side_effect=[
+            ["ua0r0"],
+            ["ua0r1"],
+            ["res0"],
+            ["ua1r0"],
+            ["ua1r1"],
+            ["res1"],
+        ]
     )
 
-    topologies = discovery._discover_group_dihedral_topology(
+    topologies = helper._discover_group_dihedral_topology(
         data_container="universe",
-        group_id=9,
-        molecules=[100, 200],
+        group_id=3,
+        molecules=[7, 8],
         level_list=["united_atom", "residue"],
     )
 
-    assert topologies == [
-        MoleculeDihedralTopology(
-            group_id=9,
-            molecule_id=100,
-            molecule_order=0,
-            num_residues=2,
-            ua_dihedrals_by_residue={
-                0: ["ua-molecule0-heavy-0"],
-                1: ["ua-molecule0-heavy-1"],
-            },
-            residue_dihedrals=["res-molecule0"],
-        ),
-        MoleculeDihedralTopology(
-            group_id=9,
-            molecule_id=200,
-            molecule_order=1,
-            num_residues=1,
-            ua_dihedrals_by_residue={
-                0: ["ua-molecule1-heavy-0"],
-            },
-            residue_dihedrals=["res-molecule1"],
-        ),
-    ]
-
-    assert universe_operations.extract_fragment_atomgroup.call_args_list == [
-        call("universe", 100),
-        call("universe", 200),
-    ]
-
-
-def test_discover_group_dihedral_topology_respects_enabled_levels(
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    universe_operations = Mock()
-
-    molecule = SimpleNamespace(label="molecule", residues=[object(), object()])
-    universe_operations.extract_fragment_atomgroup.return_value = molecule
-
-    discovery = _make_discovery(universe_operations)
-
-    select_heavy_residue = Mock(return_value="heavy-residue")
-    get_dihedrals = Mock(return_value=["residue-dihedral"])
-
-    monkeypatch.setattr(discovery, "_select_heavy_residue", select_heavy_residue)
-    monkeypatch.setattr(discovery, "_get_dihedrals", get_dihedrals)
-
-    topologies = discovery._discover_group_dihedral_topology(
-        data_container="universe",
-        group_id=1,
-        molecules=[0],
-        level_list=["residue"],
-    )
-
-    assert topologies == [
-        MoleculeDihedralTopology(
-            group_id=1,
-            molecule_id=0,
-            molecule_order=0,
-            num_residues=2,
-            ua_dihedrals_by_residue={},
-            residue_dihedrals=["residue-dihedral"],
-        )
-    ]
-
-    select_heavy_residue.assert_not_called()
-    get_dihedrals.assert_called_once_with(molecule, "residue")
+    assert [topology.molecule_id for topology in topologies] == [7, 8]
+    assert [topology.molecule_order for topology in topologies] == [0, 1]
+    assert topologies[0].group_id == 3
+    assert topologies[0].ua_dihedrals_by_residue == {0: ["ua0r0"], 1: ["ua0r1"]}
+    assert topologies[0].residue_dihedrals == ["res0"]
+    assert topologies[1].ua_dihedrals_by_residue == {0: ["ua1r0"], 1: ["ua1r1"]}
+    assert topologies[1].residue_dihedrals == ["res1"]
diff --git a/tests/unit/CodeEntropy/levels/nodes/test_axes_topology.py b/tests/unit/CodeEntropy/levels/nodes/test_axes_topology.py
deleted file mode 100644
index 2c319985..00000000
--- a/tests/unit/CodeEntropy/levels/nodes/test_axes_topology.py
+++ /dev/null
@@ -1,415 +0,0 @@
-from __future__ import annotations
-
-from types import SimpleNamespace
-
-import numpy as np
-
-from CodeEntropy.levels.nodes.axes_topology import (
-    AxesTopology,
-    BuildAxesTopologyNode,
-    ResidueAxesTopology,
-    UAAxesTopology,
-)
-
-
-class FakeAtomGroup:
-    """Small AtomGroup-like object for axes-topology tests."""
-
-    def __init__(self, atoms=None, *, name="ag"):
-        self._atoms = list(atoms or [])
-        self.name = name
-        self.indices = np.asarray([atom.index for atom in self._atoms], dtype=int)
-
-    def __iter__(self):
-        return iter(self._atoms)
-
-    def __len__(self):
-        return len(self._atoms)
-
-    def __getitem__(self, index):
-        if isinstance(index, (list, tuple, np.ndarray)):
-            return FakeAtomGroup([self._atoms[int(i)] for i in index])
-        return self._atoms[int(index)]
-
-    def select_atoms(self, selection):
-        """Return atoms matching the small mass selections used by the node."""
-        if selection == "prop mass > 1.1":
-            return FakeAtomGroup([atom for atom in self._atoms if atom.mass > 1.1])
-        if selection == "mass 2 to 999":
-            return FakeAtomGroup(
-                [atom for atom in self._atoms if 2.0 <= atom.mass <= 999.0]
-            )
-        if selection == "mass 1 to 1.1":
-            return FakeAtomGroup(
-                [atom for atom in self._atoms if 1.0 <= atom.mass <= 1.1]
-            )
-        raise AssertionError(f"Unexpected selection: {selection}")
-
-
-class FakeAtom:
-    """Small Atom-like object with mass, index, and bonded atoms."""
-
-    def __init__(self, index, mass, bonded_atoms=None):
-        self.index = index
-        self.mass = mass
-        self.bonded_atoms = bonded_atoms or FakeAtomGroup([])
-
-
-class FakeResidue:
-    """Small residue-like object."""
-
-    def __init__(self, atoms):
-        self.atoms = atoms
-
-
-class FakeMolecule:
-    """Small molecule-like object with residues and selections."""
-
-    def __init__(self, residues, select_map=None):
-        self.residues = residues
-        self._select_map = dict(select_map or {})
-
-    def select_atoms(self, selection):
-        return self._select_map.get(selection, FakeAtomGroup([]))
-
-
-class FakeAtoms:
-    """Container supporting fragments and u.atoms[index_array]."""
-
-    def __init__(self, fragments, atom_map):
-        self.fragments = fragments
-        self._atom_map = dict(atom_map)
-
-    def __getitem__(self, indices):
-        if isinstance(indices, np.ndarray):
-            return FakeAtomGroup([self._atom_map[int(index)] for index in indices])
-        if isinstance(indices, (list, tuple)):
-            return FakeAtomGroup([self._atom_map[int(index)] for index in indices])
-        return self._atom_map[int(indices)]
-
-
-class FakeUniverse:
-    """Small universe-like object."""
-
-    def __init__(self, fragments, atom_map):
-        self.atoms = FakeAtoms(fragments=fragments, atom_map=atom_map)
-
-
-def _args(*, customised_axes):
-    return SimpleNamespace(customised_axes=customised_axes)
-
-
-def _single_molecule_universe():
-    """Build a small molecule containing one residue and one UA bead."""
-    hydrogen = FakeAtom(index=2, mass=1.0)
-    bonded_heavy = FakeAtom(index=3, mass=12.0)
-    heavy = FakeAtom(
-        index=1,
-        mass=12.0,
-        bonded_atoms=FakeAtomGroup([hydrogen, bonded_heavy]),
-    )
-    other_residue_heavy = FakeAtom(index=4, mass=14.0)
-
-    residue_atoms = FakeAtomGroup([heavy, hydrogen, bonded_heavy, other_residue_heavy])
-    residue = FakeResidue(residue_atoms)
-    molecule = FakeMolecule([residue])
-
-    atom_map = {
-        heavy.index: heavy,
-        hydrogen.index: hydrogen,
-        bonded_heavy.index: bonded_heavy,
-        other_residue_heavy.index: other_residue_heavy,
-    }
-    universe = FakeUniverse([molecule], atom_map)
-
-    return universe, molecule, heavy, hydrogen, bonded_heavy, other_residue_heavy
-
-
-def test_ua_axes_topology_dataclass_preserves_arrays():
-    topology = UAAxesTopology(
-        heavy_atom_index=1,
-        ua_atom_indices=np.array([1, 2]),
-        ua_all_atom_indices=np.array([1, 3, 2]),
-        bonded_heavy_indices=np.array([3]),
-        bonded_light_indices=np.array([2]),
-        residue_heavy_indices=np.array([1, 3, 4]),
-        residue_ua_masses=np.array([13.0, 12.0, 14.0]),
-    )
-
-    assert topology.heavy_atom_index == 1
-    np.testing.assert_array_equal(topology.ua_atom_indices, np.array([1, 2]))
-    np.testing.assert_array_equal(topology.ua_all_atom_indices, np.array([1, 3, 2]))
-    np.testing.assert_array_equal(topology.bonded_heavy_indices, np.array([3]))
-    np.testing.assert_array_equal(topology.bonded_light_indices, np.array([2]))
-    np.testing.assert_array_equal(topology.residue_heavy_indices, np.array([1, 3, 4]))
-    np.testing.assert_allclose(topology.residue_ua_masses, np.array([13.0, 12.0, 14.0]))
-
-
-def test_residue_axes_topology_dataclass_preserves_arrays():
-    topology = ResidueAxesTopology(
-        residue_heavy_indices=np.array([1, 3, 4]),
-        residue_ua_masses=np.array([13.0, 12.0, 14.0]),
-        has_neighbor_bonds=True,
-    )
-
-    np.testing.assert_array_equal(topology.residue_heavy_indices, np.array([1, 3, 4]))
-    np.testing.assert_allclose(topology.residue_ua_masses, np.array([13.0, 12.0, 14.0]))
-    assert topology.has_neighbor_bonds is True
-
-
-def test_axes_topology_defaults_to_empty_mappings():
-    topology = AxesTopology()
-
-    assert topology.ua == {}
-    assert topology.residue == {}
-
-
-def test_run_writes_empty_topology_when_customised_axes_disabled():
-    node = BuildAxesTopologyNode()
-    shared_data = {"args": _args(customised_axes=False)}
-
-    result = node.run(shared_data)
-
-    assert isinstance(result["axes_topology"], AxesTopology)
-    assert result["axes_topology"].ua == {}
-    assert result["axes_topology"].residue == {}
-    assert shared_data["axes_topology"] is result["axes_topology"]
-
-
-def test_run_builds_residue_topology_for_residue_levels():
-    node = BuildAxesTopologyNode()
-    universe, molecule, heavy, hydrogen, bonded_heavy, other_residue_heavy = (
-        _single_molecule_universe()
-    )
-    neighbor_query = "(resindex -1 or resindex 1) and bonded resid 0"
-    molecule._select_map[neighbor_query] = FakeAtomGroup([bonded_heavy])
-    shared_data = {
-        "args": _args(customised_axes=True),
-        "reduced_universe": universe,
-        "levels": [["residue"]],
-        "beads": {(0, "residue"): [np.array([1, 2, 3, 4])]},
-    }
-
-    result = node.run(shared_data)
-
-    axes_topology = result["axes_topology"]
-    assert axes_topology.ua == {}
-    assert set(axes_topology.residue) == {(0, 0)}
-
-    residue_topology = axes_topology.residue[(0, 0)]
-    np.testing.assert_array_equal(
-        residue_topology.residue_heavy_indices,
-        np.array([heavy.index, bonded_heavy.index, other_residue_heavy.index]),
-    )
-    np.testing.assert_allclose(
-        residue_topology.residue_ua_masses,
-        np.array([13.0, 12.0, 14.0]),
-    )
-    assert residue_topology.has_neighbor_bonds is True
-
-
-def test_add_residue_topology_skips_when_no_residue_beads():
-    node = BuildAxesTopologyNode()
-    _, molecule, _, _, _, _ = _single_molecule_universe()
-    out = {}
-
-    node._add_residue_topology(
-        mol=molecule,
-        mol_id=0,
-        beads={},
-        out=out,
-    )
-
-    assert out == {}
-
-
-def test_has_neighbor_bonds_uses_original_residue_selection():
-    node = BuildAxesTopologyNode()
-    query = "(resindex 0 or resindex 2) and bonded resid 1"
-    molecule = FakeMolecule([], select_map={query: FakeAtomGroup([FakeAtom(1, 12.0)])})
-
-    assert node._has_neighbor_bonds(mol=molecule, local_res_i=1) is True
-
-
-def test_has_neighbor_bonds_returns_false_for_empty_selection():
-    node = BuildAxesTopologyNode()
-    molecule = FakeMolecule([])
-
-    assert node._has_neighbor_bonds(mol=molecule, local_res_i=1) is False
-
-
-def test_run_builds_ua_topology_for_united_atom_levels():
-    node = BuildAxesTopologyNode()
-    universe, _, heavy, hydrogen, bonded_heavy, other_residue_heavy = (
-        _single_molecule_universe()
-    )
-    shared_data = {
-        "args": _args(customised_axes=True),
-        "reduced_universe": universe,
-        "levels": [["united_atom"]],
-        "beads": {(0, "united_atom", 0): [np.array([1, 2])]},
-    }
-
-    result = node.run(shared_data)
-
-    axes_topology = result["axes_topology"]
-    assert shared_data["axes_topology"] is axes_topology
-    assert set(axes_topology.ua) == {(0, 0, 0)}
-
-    ua_topology = axes_topology.ua[(0, 0, 0)]
-    assert ua_topology.heavy_atom_index == heavy.index
-    np.testing.assert_array_equal(ua_topology.ua_atom_indices, np.array([1, 2]))
-    np.testing.assert_array_equal(ua_topology.ua_all_atom_indices, np.array([1, 3, 2]))
-    np.testing.assert_array_equal(ua_topology.bonded_heavy_indices, np.array([3]))
-    np.testing.assert_array_equal(ua_topology.bonded_light_indices, np.array([2]))
-    np.testing.assert_array_equal(
-        ua_topology.residue_heavy_indices,
-        np.array([heavy.index, bonded_heavy.index, other_residue_heavy.index]),
-    )
-    np.testing.assert_allclose(
-        ua_topology.residue_ua_masses,
-        np.array([13.0, 12.0, 14.0]),
-    )
-
-
-def test_run_ignores_molecules_without_united_atom_level():
-    node = BuildAxesTopologyNode()
-    universe, _, _, _, _, _ = _single_molecule_universe()
-    shared_data = {
-        "args": _args(customised_axes=True),
-        "reduced_universe": universe,
-        "levels": [["residue"]],
-        "beads": {(0, "united_atom", 0): [np.array([1, 2])]},
-    }
-
-    result = node.run(shared_data)
-
-    assert result["axes_topology"].ua == {}
-    assert shared_data["axes_topology"].ua == {}
-
-
-def test_add_ua_topology_skips_residues_without_beads():
-    node = BuildAxesTopologyNode()
-    universe, molecule, _, _, _, _ = _single_molecule_universe()
-    out = {}
-
-    node._add_ua_topology(
-        u=universe,
-        mol=molecule,
-        mol_id=0,
-        beads={},
-        out=out,
-    )
-
-    assert out == {}
-
-
-def test_add_ua_topology_skips_ua_beads_without_heavy_atoms(caplog):
-    node = BuildAxesTopologyNode()
-    hydrogen = FakeAtom(index=2, mass=1.0)
-    residue = FakeResidue(FakeAtomGroup([hydrogen]))
-    molecule = FakeMolecule([residue])
-    universe = FakeUniverse([molecule], {2: hydrogen})
-    out = {}
-
-    node._add_ua_topology(
-        u=universe,
-        mol=molecule,
-        mol_id=5,
-        beads={(5, "united_atom", 0): [np.array([2])]},
-        out=out,
-    )
-
-    assert out == {}
-    assert "Skipping UA axes topology with no heavy atom" in caplog.text
-
-
-def test_add_ua_topology_handles_multiple_residues_and_ua_beads():
-    node = BuildAxesTopologyNode()
-
-    h0 = FakeAtom(index=10, mass=1.0)
-    c0 = FakeAtom(index=11, mass=12.0, bonded_atoms=FakeAtomGroup([h0]))
-    residue0 = FakeResidue(FakeAtomGroup([c0, h0]))
-
-    h1 = FakeAtom(index=20, mass=1.0)
-    c1 = FakeAtom(index=21, mass=12.0, bonded_atoms=FakeAtomGroup([h1]))
-    residue1 = FakeResidue(FakeAtomGroup([c1, h1]))
-
-    molecule = FakeMolecule([residue0, residue1])
-    universe = FakeUniverse(
-        [molecule],
-        {
-            h0.index: h0,
-            c0.index: c0,
-            h1.index: h1,
-            c1.index: c1,
-        },
-    )
-    out = {}
-
-    node._add_ua_topology(
-        u=universe,
-        mol=molecule,
-        mol_id=3,
-        beads={
-            (3, "united_atom", 0): [np.array([11, 10])],
-            (3, "united_atom", 1): [np.array([21, 20])],
-        },
-        out=out,
-    )
-
-    assert set(out) == {(3, 0, 0), (3, 1, 0)}
-    assert out[(3, 0, 0)].heavy_atom_index == 11
-    assert out[(3, 1, 0)].heavy_atom_index == 21
-
-
-def test_split_bonded_atoms_returns_heavy_and_light_atom_groups():
-    hydrogen = FakeAtom(index=2, mass=1.0)
-    heavy_bonded = FakeAtom(index=3, mass=12.0)
-    atom = FakeAtom(
-        index=1,
-        mass=12.0,
-        bonded_atoms=FakeAtomGroup([hydrogen, heavy_bonded]),
-    )
-
-    bonded_heavy, bonded_light = BuildAxesTopologyNode._split_bonded_atoms(atom)
-
-    np.testing.assert_array_equal(bonded_heavy.indices, np.array([3]))
-    np.testing.assert_array_equal(bonded_light.indices, np.array([2]))
-
-
-def test_get_ua_masses_from_topology_adds_bonded_hydrogen_masses():
-    hydrogen = FakeAtom(index=2, mass=1.0)
-    heavy = FakeAtom(index=1, mass=12.0, bonded_atoms=FakeAtomGroup([hydrogen]))
-    other_heavy = FakeAtom(index=3, mass=14.0)
-    atom_group = FakeAtomGroup([heavy, hydrogen, other_heavy])
-
-    masses = BuildAxesTopologyNode._get_ua_masses_from_topology(atom_group)
-
-    assert masses == [13.0, 14.0]
-
-
-def test_get_ua_masses_from_topology_handles_atoms_without_bonded_atoms_attribute():
-    class AtomWithoutBonds:
-        """Small atom-like object without bonded_atoms."""
-
-        def __init__(self, index, mass):
-            self.index = index
-            self.mass = mass
-
-    heavy = AtomWithoutBonds(index=1, mass=12.0)
-    hydrogen = AtomWithoutBonds(index=2, mass=1.0)
-    atom_group = FakeAtomGroup([heavy, hydrogen])
-
-    masses = BuildAxesTopologyNode._get_ua_masses_from_topology(atom_group)
-
-    assert masses == [12.0]
-
-
-def test_get_ua_masses_from_topology_returns_empty_list_for_light_atoms_only():
-    hydrogen = FakeAtom(index=2, mass=1.0)
-    atom_group = FakeAtomGroup([hydrogen])
-
-    masses = BuildAxesTopologyNode._get_ua_masses_from_topology(atom_group)
-
-    assert masses == []
diff --git a/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py b/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py
index 9d66e88f..d989069f 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,22 +436,22 @@ 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
@@ -526,12 +470,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 +484,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 +493,10 @@ 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,
             axes_manager=axes_manager,
-            axes_topology=None,
-            box=None,
             customised_axes=True,
         )
         == expected
@@ -612,14 +521,10 @@ 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,
             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..93740467 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:
@@ -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,25 @@ 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, 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 +166,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 +183,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 +196,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 +213,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 +229,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():
@@ -505,18 +516,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 +538,12 @@ 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, residue=residue)
 
     assert np.allclose(trans, np.eye(3) * 2)
     assert np.allclose(rot, np.eye(3) * 2)
@@ -639,12 +650,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
+
     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 +681,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 +718,37 @@ 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,
-):
-    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,
-    )
-
-    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,
-    )
-
-    np.testing.assert_allclose(trans_axes, np.eye(3) * 2.0)
-    np.testing.assert_allclose(rot_axes, np.eye(3) * 2.0)
-    np.testing.assert_allclose(center, np.array([9.0, 8.0, 7.0]))
-    np.testing.assert_allclose(moi, np.array([3.0, 2.0, 1.0]))
-
-    tensor_kwargs = get_tensor.call_args.kwargs
-    np.testing.assert_allclose(
-        tensor_kwargs["center_of_mass"],
-        np.array([9.0, 8.0, 7.0]),
-    )
-    np.testing.assert_allclose(
-        tensor_kwargs["positions"],
-        np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
-    )
-    np.testing.assert_allclose(tensor_kwargs["masses"], np.array([13.0, 14.0]))
-    np.testing.assert_allclose(tensor_kwargs["dimensions"], box)
-    get_principal.assert_called_once()
-
-
-def test_get_residue_axes_from_topology_neighbor_bonds_uses_vanilla_axes(
-    monkeypatch,
-):
-    ax = AxesCalculator()
-
-    universe = _FakeUniverse(
-        {},
-        dimensions=[11.0, 12.0, 13.0, 90.0, 90.0, 90.0],
-    )
-    mol = MagicMock()
-    mol.atoms.principal_axes.return_value = np.eye(3) * 5.0
-    residue_atoms = MagicMock()
-    residue_atoms.center_of_mass.return_value = np.array([1.0, 2.0, 3.0])
-    topology = _residue_topology(has_neighbor_bonds=True)
-
-    make_whole = MagicMock()
-    get_vanilla = MagicMock(return_value=(np.eye(3) * 6.0, np.array([6.0, 5.0, 4.0])))
-
-    monkeypatch.setattr("CodeEntropy.levels.axes.make_whole", make_whole)
-    monkeypatch.setattr(ax, "get_vanilla_axes", get_vanilla)
-
-    trans_axes, rot_axes, center, moi = ax.get_residue_axes_from_topology(
-        u=universe,
-        mol=mol,
-        residue_atoms=residue_atoms,
-        topology=topology,
-        box=None,
-    )
-
-    make_whole.assert_called_once_with(mol.atoms)
-    mol.atoms.principal_axes.assert_called_once()
-    get_vanilla.assert_called_once_with(residue_atoms)
-    np.testing.assert_allclose(trans_axes, np.eye(3) * 5.0)
-    np.testing.assert_allclose(rot_axes, np.eye(3) * 6.0)
-    np.testing.assert_allclose(center, np.array([1.0, 2.0, 3.0]))
-    np.testing.assert_allclose(moi, np.array([6.0, 5.0, 4.0]))
-
-
-def test_get_UA_axes_from_topology_multiple_heavy_uses_cached_indices_and_box(
-    monkeypatch,
-):
-    ax = AxesCalculator()
-
-    heavy_atom = _FakeAtom(1, 12.0, [1.0, 2.0, 3.0])
-    other_heavy = _FakeAtom(3, 14.0, [4.0, 5.0, 6.0])
-    universe = _FakeUniverse({1: heavy_atom, 3: other_heavy})
-    residue_atoms = MagicMock()
-    residue_atoms.center_of_mass.return_value = np.array([9.0, 8.0, 7.0])
-
-    topology = _ua_topology(
-        heavy_atom_index=1,
-        residue_heavy_indices=(1, 3),
-        residue_ua_masses=(13.0, 14.0),
-    )
-
-    get_tensor = MagicMock(return_value=np.eye(3))
-    get_principal = MagicMock(return_value=(np.eye(3) * 2.0, np.array([3.0, 2.0, 1.0])))
-    get_bonded = MagicMock(return_value=(np.eye(3) * 4.0, np.array([1.0, 1.0, 1.0])))
-
-    monkeypatch.setattr(ax, "get_moment_of_inertia_tensor", get_tensor)
-    monkeypatch.setattr(ax, "get_custom_principal_axes", get_principal)
-    monkeypatch.setattr(ax, "get_bonded_axes_from_topology", get_bonded)
-
-    box = np.array([20.0, 30.0, 40.0])
-    trans_axes, rot_axes, center, moi = ax.get_UA_axes_from_topology(
-        u=universe,
-        residue_atoms=residue_atoms,
-        topology=topology,
-        box=box,
-    )
-
-    np.testing.assert_allclose(trans_axes, np.eye(3) * 2.0)
-    np.testing.assert_allclose(rot_axes, np.eye(3) * 4.0)
-    np.testing.assert_allclose(center, heavy_atom.position)
-    np.testing.assert_allclose(moi, np.array([1.0, 1.0, 1.0]))
-
-    get_tensor.assert_called_once()
-    tensor_kwargs = get_tensor.call_args.kwargs
-    np.testing.assert_allclose(
-        tensor_kwargs["center_of_mass"], np.array([9.0, 8.0, 7.0])
-    )
-    np.testing.assert_allclose(
-        tensor_kwargs["positions"], np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
-    )
-    np.testing.assert_allclose(tensor_kwargs["masses"], np.array([13.0, 14.0]))
-    np.testing.assert_allclose(tensor_kwargs["dimensions"], box)
-
-    get_principal.assert_called_once()
-    np.testing.assert_allclose(get_principal.call_args.args[0], np.eye(3))
-    get_bonded.assert_called_once_with(
-        u=universe,
-        heavy_atom=heavy_atom,
-        topology=topology,
-        dimensions=box,
-    )
-
-
-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,
-    )
-
-    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])
-    )
-
-
-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)
-    monkeypatch.setattr(
-        ax, "get_bonded_axes_from_topology", lambda **kwargs: (None, None)
-    )
-
-    with pytest.raises(ValueError, match="cached UA bead"):
-        ax.get_UA_axes_from_topology(
-            u=universe,
-            residue_atoms=residue_atoms,
-            topology=topology,
-            box=None,
-        )
-
-
-def test_get_bonded_axes_from_topology_non_heavy_returns_none_none():
-    ax = AxesCalculator()
-    light_atom = _FakeAtom(1, 1.0, [0.0, 0.0, 0.0])
-
-    custom_axes, moi = ax.get_bonded_axes_from_topology(
-        u=MagicMock(),
-        heavy_atom=light_atom,
-        topology=_ua_topology(heavy_atom_index=1),
-        dimensions=np.array([10.0, 10.0, 10.0]),
-    )
-
-    assert custom_axes is None
-    assert moi is None
-
-
-def test_get_bonded_axes_from_topology_no_bonded_heavy_uses_vanilla_axes(
-    monkeypatch,
-):
+def test_get_residue_axes_custom_path(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,),
+    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_vanilla = MagicMock(return_value=(np.eye(3) * 7.0, np.array([7.0, 8.0, 9.0])))
-    get_custom_moi = MagicMock()
-    get_flipped = MagicMock(return_value=np.eye(3) * -7.0)
-
-    monkeypatch.setattr(ax, "get_vanilla_axes", get_vanilla)
-    monkeypatch.setattr(ax, "get_custom_moment_of_inertia", get_custom_moi)
-    monkeypatch.setattr(ax, "get_flipped_axes", get_flipped)
-
-    custom_axes, moi = ax.get_bonded_axes_from_topology(
-        u=universe,
-        heavy_atom=heavy_atom,
-        topology=topology,
-        dimensions=np.array([10.0, 10.0, 10.0]),
-    )
-
-    np.testing.assert_allclose(custom_axes, np.eye(3) * -7.0)
-    np.testing.assert_allclose(moi, np.array([7.0, 8.0, 9.0]))
-    get_vanilla.assert_called_once()
-    get_custom_moi.assert_not_called()
-    get_flipped.assert_called_once()
-
-
-def test_get_bonded_axes_from_topology_one_heavy_no_light_uses_custom_axes(
-    monkeypatch,
-):
-    ax = AxesCalculator()
-
-    heavy_atom = _FakeAtom(1, 12.0, [0.0, 0.0, 0.0])
-    bonded_heavy = _FakeAtom(3, 12.0, [1.0, 0.0, 0.0])
-    universe = _FakeUniverse({1: heavy_atom, 3: bonded_heavy})
-    topology = _ua_topology(
-        heavy_atom_index=1,
-        ua_atom_indices=(1,),
-        ua_all_atom_indices=(1, 3),
-        bonded_heavy_indices=(3,),
-        bonded_light_indices=(),
+    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
     )
 
-    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]),
-    )
-
-    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,),
-    )
-
-    get_custom_axes = MagicMock(return_value=np.eye(3))
-    monkeypatch.setattr(ax, "get_custom_axes", get_custom_axes)
-    monkeypatch.setattr(
-        ax,
-        "get_custom_moment_of_inertia",
-        lambda **kwargs: 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]),
-    )
-
-    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
-    )
-
-
-def test_get_bonded_axes_from_topology_two_heavy_uses_heavy_positions_as_b_list(
-    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,
-        }
-    )
-    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=(),
-    )
-
-    get_custom_axes = MagicMock(return_value=np.eye(3) * 4.0)
-    monkeypatch.setattr(ax, "get_custom_axes", get_custom_axes)
-    monkeypatch.setattr(
-        ax,
-        "get_custom_moment_of_inertia",
-        lambda **kwargs: np.array([4.0, 5.0, 6.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]),
-    )
-
-    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 rot_center.shape == (3,)
+    assert rot_axes.shape == (3, 3)
 
 
-def test_get_bonded_axes_from_topology_returns_none_when_custom_axes_none(
-    monkeypatch,
-):
+def test_get_custom_residue_moment_of_inertia(monkeypatch):
     ax = AxesCalculator()
-
-    heavy_atom = _FakeAtom(1, 12.0, [0.0, 0.0, 0.0])
-    bonded_heavy = _FakeAtom(3, 12.0, [1.0, 0.0, 0.0])
-    universe = _FakeUniverse({1: heavy_atom, 3: bonded_heavy})
-    topology = _ua_topology(
-        heavy_atom_index=1,
-        ua_atom_indices=(1,),
-        ua_all_atom_indices=(1, 3),
-        bonded_heavy_indices=(3,),
-        bonded_light_indices=(),
+    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),
     )
+    dimensions = np.array([10.0, 10.0, 10.0], dtype=float)
 
-    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]),
+    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,
     )
 
-    assert custom_axes is None
-    assert moi is None
-    get_custom_moi.assert_not_called()
-    get_flipped.assert_not_called()
+    assert moi.shape == (3,)
diff --git a/tests/unit/CodeEntropy/levels/test_conformation_dag.py b/tests/unit/CodeEntropy/levels/test_conformation_dag.py
index cb634fe8..92568ca9 100644
--- a/tests/unit/CodeEntropy/levels/test_conformation_dag.py
+++ b/tests/unit/CodeEntropy/levels/test_conformation_dag.py
@@ -3,7 +3,7 @@
 from types import SimpleNamespace
 from unittest.mock import MagicMock
 
-from CodeEntropy.levels.graph.conformation_dag import ConformationDAG
+from CodeEntropy.levels.conformation_dag import ConformationDAG
 from CodeEntropy.trajectory.frames import FrameSelection
 
 
diff --git a/tests/unit/CodeEntropy/levels/test_frame_graph.py b/tests/unit/CodeEntropy/levels/test_frame_graph.py
index 59a9e757..e1134772 100644
--- a/tests/unit/CodeEntropy/levels/test_frame_graph.py
+++ b/tests/unit/CodeEntropy/levels/test_frame_graph.py
@@ -2,7 +2,7 @@
 
 import pytest
 
-from CodeEntropy.levels.graph.frame_dag import FrameGraph
+from CodeEntropy.levels.frame_dag import FrameGraph
 
 
 def test_make_frame_ctx_has_required_keys():
diff --git a/tests/unit/CodeEntropy/levels/test_level_dag.py b/tests/unit/CodeEntropy/levels/test_level_dag.py
index 51674c73..1b9f4d4f 100644
--- a/tests/unit/CodeEntropy/levels/test_level_dag.py
+++ b/tests/unit/CodeEntropy/levels/test_level_dag.py
@@ -4,17 +4,16 @@
 
 from unittest.mock import MagicMock, call, patch
 
-from CodeEntropy.levels.graph.level_dag import LevelDAG
+from CodeEntropy.levels.level_dag import LevelDAG
 
 
 def test_build_registers_static_nodes_and_builds_stage_dags():
     with (
-        patch("CodeEntropy.levels.graph.level_dag.DetectMoleculesNode"),
-        patch("CodeEntropy.levels.graph.level_dag.DetectLevelsNode"),
-        patch("CodeEntropy.levels.graph.level_dag.BuildBeadsNode"),
-        patch("CodeEntropy.levels.graph.level_dag.BuildAxesTopologyNode"),
-        patch("CodeEntropy.levels.graph.level_dag.InitCovarianceAccumulatorsNode"),
-        patch("CodeEntropy.levels.graph.level_dag.ConformationDAG"),
+        patch("CodeEntropy.levels.level_dag.DetectMoleculesNode"),
+        patch("CodeEntropy.levels.level_dag.DetectLevelsNode"),
+        patch("CodeEntropy.levels.level_dag.BuildBeadsNode"),
+        patch("CodeEntropy.levels.level_dag.InitCovarianceAccumulatorsNode"),
+        patch("CodeEntropy.levels.level_dag.ConformationDAG"),
     ):
         universe_operations = MagicMock()
         dag = LevelDAG(universe_operations=universe_operations)
@@ -28,7 +27,6 @@ def test_build_registers_static_nodes_and_builds_stage_dags():
         "detect_molecules",
         "detect_levels",
         "build_beads",
-        "build_axes_topology",
         "init_covariance_accumulators",
     }
     assert "find_neighbors" not in dag._static_nodes
@@ -36,7 +34,6 @@ def test_build_registers_static_nodes_and_builds_stage_dags():
 
     assert ("detect_molecules", "detect_levels") in dag._static_graph.edges
     assert ("detect_levels", "build_beads") in dag._static_graph.edges
-    assert ("build_beads", "build_axes_topology") in dag._static_graph.edges
     assert ("detect_levels", "init_covariance_accumulators") in dag._static_graph.edges
     assert (
         "detect_levels",
@@ -59,12 +56,8 @@ def test_execute_sets_default_axes_manager_and_runs_workflow_stages():
     dag._run_frame_stage = MagicMock()
 
     with (
-        patch(
-            "CodeEntropy.levels.graph.level_dag.NeighborReducer.initialise"
-        ) as initialise,
-        patch(
-            "CodeEntropy.levels.graph.level_dag.NeighborReducer.finalise"
-        ) as finalise,
+        patch("CodeEntropy.levels.level_dag.NeighborReducer.initialise") as initialise,
+        patch("CodeEntropy.levels.level_dag.NeighborReducer.finalise") as finalise,
     ):
         out = dag.execute(shared_data, progress=progress)
 
@@ -168,7 +161,7 @@ def test_run_frame_stage_collects_frame_indices_and_delegates_to_scheduler():
 
     shared_data = {"frame_source": frame_source}
 
-    with patch("CodeEntropy.levels.graph.level_dag.FrameScheduler") as Scheduler:
+    with patch("CodeEntropy.levels.level_dag.FrameScheduler") as Scheduler:
         scheduler = Scheduler.return_value
 
         dag._run_frame_stage(shared_data, progress=progress)
@@ -193,7 +186,7 @@ def test_initialise_neighbor_metadata_writes_symmetry_and_linearity():
     groups = {0: [0], 1: [1]}
     shared_data = {"reduced_universe": universe, "groups": groups}
 
-    with patch("CodeEntropy.levels.graph.level_dag.Neighbors") as Neighbors:
+    with patch("CodeEntropy.levels.level_dag.Neighbors") as Neighbors:
         helper = Neighbors.return_value
         helper.get_symmetry.return_value = ({0: 12, 1: 2}, {0: False, 1: True})
 
@@ -208,7 +201,7 @@ def test_initialise_neighbor_metadata_falls_back_to_universe_key():
     universe = object()
     shared_data = {"universe": universe, "groups": {0: [0]}}
 
-    with patch("CodeEntropy.levels.graph.level_dag.Neighbors") as Neighbors:
+    with patch("CodeEntropy.levels.level_dag.Neighbors") as Neighbors:
         helper = Neighbors.return_value
         helper.get_symmetry.return_value = ({0: 1}, {0: False})
 
@@ -243,11 +236,11 @@ def test_level_dag_runs_static_conformation_then_frame(monkeypatch):
     )
 
     monkeypatch.setattr(
-        "CodeEntropy.levels.graph.level_dag.NeighborReducer.initialise",
+        "CodeEntropy.levels.level_dag.NeighborReducer.initialise",
         lambda shared_data: calls.append("neighbor_initialise"),
     )
     monkeypatch.setattr(
-        "CodeEntropy.levels.graph.level_dag.NeighborReducer.finalise",
+        "CodeEntropy.levels.level_dag.NeighborReducer.finalise",
         lambda shared_data: calls.append("neighbor_finalise"),
     )
 
diff --git a/tests/unit/CodeEntropy/levels/test_mda_universe_operations.py b/tests/unit/CodeEntropy/levels/test_mda_universe_operations.py
index d298c486..e8d92329 100644
--- a/tests/unit/CodeEntropy/levels/test_mda_universe_operations.py
+++ b/tests/unit/CodeEntropy/levels/test_mda_universe_operations.py
@@ -1,5 +1,4 @@
-from types import SimpleNamespace
-from unittest.mock import MagicMock, Mock
+from unittest.mock import MagicMock
 
 import numpy as np
 import pytest
@@ -408,19 +407,3 @@ def test_select_frame_indices_raises_when_frame_indices_empty():
         ops.select_frame_indices(u, frame_indices=[])
 
     u.select_atoms.assert_not_called()
-
-
-def test_extract_fragment_atomgroup_returns_lightweight_range_selection() -> None:
-    universe = Mock()
-    universe.atoms.fragments = [
-        SimpleNamespace(indices=[10, 11, 12, 13]),
-    ]
-    universe.select_atoms.return_value = "fragment_atomgroup"
-
-    result = UniverseOperations().extract_fragment_atomgroup(universe, 0)
-
-    assert result == "fragment_atomgroup"
-    universe.select_atoms.assert_called_once_with(
-        "index 10:13",
-        updating=False,
-    )