From 88127096113408380bce2af4441ac93f1dd24b6b Mon Sep 17 00:00:00 2001
From: Ioana Papa <iapapa1@sheffield.ac.uk>
Date: Mon, 17 Aug 2026 12:26:50 +0100
Subject: [PATCH 01/17] tried a coordinate system for terminal residue using
 average of bonded atoms + average of all other heavy atoms in resid

---
 CodeEntropy/levels/axes.py | 155 ++++++++++++++++++++-----------------
 1 file changed, 84 insertions(+), 71 deletions(-)

diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py
index 7d93fd9..b2608ad 100644
--- a/CodeEntropy/levels/axes.py
+++ b/CodeEntropy/levels/axes.py
@@ -77,7 +77,10 @@ def get_residue_axes(
               heavy atom + bonded hydrogens.
             * Set translational axes equal to rotational axes (as per the original
               code convention).
-        - If bonded to other residues:
+        - If bonded to only one other residue:
+            * Translational axes are principal axes of data_container.
+            * Find edge heavy atom (i.e. heavy atoms bonded to neighbour residue).
+        - If bonded to at least two other residues:
             * Translational axes are principal axes of data_container.
             * Find edge heavy atoms (i.e. heavy atoms bonded to neighbour residues)
               and find the shortest chain between them: the backbone. Edge
@@ -145,30 +148,36 @@ def get_residue_axes(
             trans_axes = data_container.atoms.principal_axes()
 
             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 heavy atom
-                    last_index = len(uas) - 1
-                    last = None
-                    if last_index > 0 and last is None:
-                        heavy_atom = uas[last_index]
-                        last = heavy_atom
-                    edges = [edge_atom_set[0], last]
-
-                    backbone = self.get_chain(residue, edge_atom_set[0], last)
+                edge_atom = edge_atom_set[0]
+                bonded_atoms = uas.select_atoms(f"bonded index {edge_atom.index}")
+                # find the average position of heavy atoms bonded to edge atom
+                average_bonded_atom = np.zeros(3)
+                for atom in bonded_atoms:
+                    average_bonded_atom += atom.position
+                average_bonded_atom /= len(bonded_atoms)
+                # find the average position of all other heavy atoms in residue
+                other_atoms = []
+                for atom in uas:
+                    if atom != edge_atom and atom not in bonded_atoms:
+                        other_atoms.append(atom)
+                average_other_atoms = np.zeros(3)
+                for atom in other_atoms:
+                    average_other_atoms += atom.position
+                average_other_atoms /= len(other_atoms)
+                rot_center, rot_axes = self.get_residue_custom_axes(
+                    [edge_atom.position, average_other_atoms], average_bonded_atom
+                )
+
             else:
-                edges = [edge_atom_set[0], edge_atom_set[1]]
+                edges = [edge_atom_set[0].position, edge_atom_set[1].position]
                 backbone = self.get_chain(residue, edge_atom_set[0], edge_atom_set[1])
-            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)
+                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,
@@ -248,7 +257,8 @@ def get_UA_axes(self, data_container, index: int, res_position):
             Identify residue of interest and neighbours, then select
             edge heavy atoms (i.e. heavy atoms bonded to neighbour residues).
             If there are no bonds to neighbouring residues, use residue
-            .principal axes Otherwise, find the shortest chain between edge
+            principal axes. Otherwise, for residues with at least two neighbours,
+            find the shortest chain between edge
             residues: the backbone. Edge atoms + backbone COM are used to
             determine UA translational axes (see get_residue_custom_axes)
 
@@ -293,67 +303,70 @@ def get_UA_axes(self, data_container, index: int, res_position):
                 residue_heavy_atoms = heavy_atoms
             else:
                 # residue of interest has at least one neighbour
-                if res_position == -1:
-                    residue = data_container.residues[0]
-                    resindex = residue.resindex
-                    resindex_next = resindex + 1
-
-                    second_edge = data_container.select_atoms(
-                        f"resindex {resindex} and bonded resindex {resindex_next}"
+                if res_position == -1 or res_position == 1:
+                    # look at a terminal residue
+                    if res_position == -1:
+                        # first residue
+                        residue = data_container.residues[0]
+                        resindex = residue.resindex
+                        resindex_next = resindex + 1
+                        edge_atom = data_container.select_atoms(
+                            f"resindex {resindex} and bonded resindex {resindex_next}"
+                        )
+                    else:
+                        # last residue
+                        residue = data_container.residues[1]
+                        resindex = residue.resindex
+                        resindex_prev = resindex - 1
+                        edge_atom = data_container.select_atoms(
+                            f"resindex {resindex} and bonded resindex {resindex_prev}"
+                        )
+                    residue_heavy_atoms = residue.atoms.select_atoms("mass 2 to 999")
+                    bonded_atoms = residue_heavy_atoms.select_atoms(
+                        f"bonded index {edge_atom[0].index}"
                     )
-
-                    edges = [residue.atoms[0], second_edge[0]]
-                    backbone = self.get_chain(
-                        residue, residue.atoms[0], second_edge.atoms[0]
+                    # find the average position of heavy atoms bonded to edge atom
+                    average_bonded_atom = np.zeros(3)
+                    for atom in bonded_atoms:
+                        average_bonded_atom += atom.position
+                    average_bonded_atom /= len(bonded_atoms)
+                    # find the average position of all other heavy atoms in residue
+                    other_atoms = []
+                    for atom in residue_heavy_atoms:
+                        if atom != edge_atom and atom not in bonded_atoms:
+                            other_atoms.append(atom)
+                    average_other_atoms = np.zeros(3)
+                    for atom in other_atoms:
+                        average_other_atoms += atom.position
+                    average_other_atoms /= len(other_atoms)
+                    trans_center, trans_axes = self.get_residue_custom_axes(
+                        [edge_atom.positions[0], average_other_atoms],
+                        average_bonded_atom,
                     )
-
-                elif res_position == 0:
+                else:
                     # between 2 residues
                     residue = data_container.residues[1]
                     resindex = residue.resindex
                     resindex_next = resindex + 1
                     resindex_prev = resindex - 1
-
+                    residue_heavy_atoms = residue.atoms.select_atoms("mass 2 to 999")
                     edge_set = data_container.select_atoms(
                         f"resindex {resindex} and "
                         f"(bonded resindex {resindex_prev} or "
                         f"resindex {resindex_next})"
                     )
 
-                    edges = [edge_set[0], edge_set[1]]
+                    edges = edge_set.positions
                     backbone = self.get_chain(residue, edge_set[0], edge_set[1])
+                    backbone_center = np.zeros(3)
+                    for heavy_atom in backbone:
+                        backbone_center += heavy_atom.position
+                    backbone_center = backbone_center / len(backbone)
 
-                else:
-                    # last resid
-                    # always resindex 1 in data_container
-                    residue = data_container.residues[1]
-                    resindex = residue.resindex
-                    resindex_prev = resindex - 1
-                    first_edge = data_container.select_atoms(
-                        f"resindex {resindex} and bonded resindex {resindex_prev}"
+                    trans_center, trans_axes = self.get_residue_custom_axes(
+                        edges, backbone_center
                     )
 
-                    last_index = len(heavy_atoms) - 1
-                    last = None
-                    # look for last heavy atom
-                    # with only one bond to another
-                    if last_index > 0 and last is None:
-                        heavy_atom = heavy_atoms[last_index]
-                        last = heavy_atom
-
-                    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:
@@ -580,9 +593,9 @@ def get_residue_custom_axes(self, edges, center):
             lies on the E1-E2 vector
             rot_axes: (3,3) rotation axes of residue
         """
-        first_edge_centre_of_geometry_vector = center - edges[0].position
+        first_edge_centre_of_geometry_vector = center - edges[0]
         # look for projection of E1-O onto E1-E2 (E1-C)
-        first_edge_second_edge_vector = edges[1].position - edges[0].position
+        first_edge_second_edge_vector = edges[1] - edges[0]
         first_edge_origin_vector = (
             np.dot(first_edge_second_edge_vector, first_edge_centre_of_geometry_vector)
             / (np.linalg.norm(first_edge_second_edge_vector) ** 2)
@@ -598,7 +611,7 @@ def get_residue_custom_axes(self, edges, center):
         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 = first_edge_origin_vector + edges[0].position
+        rot_center = first_edge_origin_vector + edges[0]
         return rot_center, rot_axes
 
     def get_bonded_axes(self, system, atom, dimensions: np.ndarray):

From 578b596cfbf553ceebd6266a877f6069d88b5d28 Mon Sep 17 00:00:00 2001
From: ioanaapapa <ioanaalexandrapapa@gmail.com>
Date: Tue, 25 Aug 2026 15:19:47 +0100
Subject: [PATCH 02/17] updated docstring

---
 CodeEntropy/levels/axes.py | 34 +++++++++++++++++++++++-----------
 1 file changed, 23 insertions(+), 11 deletions(-)

diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py
index b2608ad..4ae31b3 100644
--- a/CodeEntropy/levels/axes.py
+++ b/CodeEntropy/levels/axes.py
@@ -73,20 +73,26 @@ def get_residue_axes(
           (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
-              that uses positions of heavy atoms only, but including masses of
+              that uses positions of heavy atoms only, but includes masses of
               heavy atom + bonded hydrogens.
             * Set translational axes equal to rotational axes (as per the original
               code convention).
         - If bonded to only one other residue:
             * Translational axes are principal axes of data_container.
             * Find edge heavy atom (i.e. heavy atoms bonded to neighbour residue).
+              Find all heavy atoms bonded to edge heavy atom and compute their average
+              position.
+              Find all other heavy atoms in residue and compute their average position.
+              The three points are now used to obtain determine residue rotational axes.
+              (see get_residue_custom_axes)
         - If bonded to at least two other residues:
             * Translational axes are principal axes of data_container.
             * 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 residue rotational axes.
-              (see get_residue_custom_axes).Compute a custom MOI, using heavy atom
-              positions and heavy atom + hydrogen masses.
+              (see get_residue_custom_axes).
+        Compute a custom MOI, using heavy atom positions and
+          heavy atom + hydrogen masses.
 
         Args:
             data_container (MDAnalysis.Universe or AtomGroup):
@@ -256,19 +262,25 @@ def get_UA_axes(self, data_container, index: int, res_position):
             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).
-            If there are no bonds to neighbouring residues, use residue
-            principal axes. Otherwise, for residues with at least two neighbours,
-            find the shortest chain between edge
-            residues: the backbone. Edge atoms + backbone COM are used to
-            determine UA translational axes (see get_residue_custom_axes)
-
+            - If there are *no* bonds to other residues, use a custom principal axes
+              from a moment-of-inertia (MOI) tensor that uses positions of heavy atoms
+              only, but includes masses of heavy atom + bonded hydrogens.
+            - If bonded to only one other residue, find edge heavy atom
+              (i.e. heavy atom bonded to neighbour residue). Find all heavy atoms
+              bonded to edge heavy atom and compute their average position.
+              Find all other heavy atoms in residue and compute their average position.
+              The three points are now used to obtain determine residue rotational axes.
+              (see get_residue_custom_axes)
+            - If bonded to at least two other residues, 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 residue rotational axes. (see get_residue_custom_axes).
         - Rotational axes:
             Identify heavy atoms in the residue/molecule of interest and choose
             the `index`-th heavy atom (where index corresponds to the bead index).
             Use bonded topology around that heavy atom to determine UA rotational
             axes (see :meth:`get_bonded_axes`).
-            Compute a custom MOI tensor using heavy-atom coordinates but UA masses
-            (heavy + bonded H masses), then compute the principal axes from it.
+            Compute a custom MOI tensor.
 
         Args:
             data_container (MDAnalysis.Universe or AtomGroup):

From 391f32ad7cd93892359921350348d4d713be981a Mon Sep 17 00:00:00 2001
From: ioanaapapa <ioanaalexandrapapa@gmail.com>
Date: Tue, 25 Aug 2026 15:22:05 +0100
Subject: [PATCH 03/17] updated docstring

---
 CodeEntropy/levels/axes.py | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py
index 4ae31b3..b6715c0 100644
--- a/CodeEntropy/levels/axes.py
+++ b/CodeEntropy/levels/axes.py
@@ -263,18 +263,18 @@ def get_UA_axes(self, data_container, index: int, res_position):
             Identify residue of interest and neighbours, then select
             edge heavy atoms (i.e. heavy atoms bonded to neighbour residues).
             - If there are *no* bonds to other residues, use a custom principal axes
-              from a moment-of-inertia (MOI) tensor that uses positions of heavy atoms
-              only, but includes masses of heavy atom + bonded hydrogens.
+            from a moment-of-inertia (MOI) tensor that uses positions of heavy atoms
+            only, but includes masses of heavy atom + bonded hydrogens.
             - If bonded to only one other residue, find edge heavy atom
-              (i.e. heavy atom bonded to neighbour residue). Find all heavy atoms
-              bonded to edge heavy atom and compute their average position.
-              Find all other heavy atoms in residue and compute their average position.
-              The three points are now used to obtain determine residue rotational axes.
-              (see get_residue_custom_axes)
+            (i.e. heavy atom bonded to neighbour residue). Find all heavy atoms
+            bonded to edge heavy atom and compute their average position.
+            Find all other heavy atoms in residue and compute their average position.
+             The three points are now used to obtain determine residue rotational axes.
+            (see get_residue_custom_axes)
             - If bonded to at least two other residues, 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 residue rotational axes. (see get_residue_custom_axes).
+            chain between them: the backbone. Edge atoms + backbone COM are used
+            to determine residue rotational axes. (see get_residue_custom_axes).
         - Rotational axes:
             Identify heavy atoms in the residue/molecule of interest and choose
             the `index`-th heavy atom (where index corresponds to the bead index).

From 89342428c658aee538bfe3def8166a4c40ee1c97 Mon Sep 17 00:00:00 2001
From: ioanaapapa <ioanaalexandrapapa@gmail.com>
Date: Tue, 25 Aug 2026 15:24:52 +0100
Subject: [PATCH 04/17] updated docstring

---
 CodeEntropy/levels/axes.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py
index b6715c0..34ce091 100644
--- a/CodeEntropy/levels/axes.py
+++ b/CodeEntropy/levels/axes.py
@@ -265,16 +265,19 @@ def get_UA_axes(self, data_container, index: int, res_position):
             - If there are *no* bonds to other residues, use a custom principal axes
             from a moment-of-inertia (MOI) tensor that uses positions of heavy atoms
             only, but includes masses of heavy atom + bonded hydrogens.
+
             - If bonded to only one other residue, find edge heavy atom
             (i.e. heavy atom bonded to neighbour residue). Find all heavy atoms
             bonded to edge heavy atom and compute their average position.
             Find all other heavy atoms in residue and compute their average position.
-             The three points are now used to obtain determine residue rotational axes.
+            The three points are now used to obtain determine residue rotational axes.
             (see get_residue_custom_axes)
+
             - If bonded to at least two other residues, 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 residue rotational axes. (see get_residue_custom_axes).
+
         - Rotational axes:
             Identify heavy atoms in the residue/molecule of interest and choose
             the `index`-th heavy atom (where index corresponds to the bead index).

From 10b9edab34c9ec023d4326b8bbf48548153673bb Mon Sep 17 00:00:00 2001
From: ioanaapapa <ioanaalexandrapapa@gmail.com>
Date: Tue, 25 Aug 2026 15:25:31 +0100
Subject: [PATCH 05/17] updated docstring

---
 CodeEntropy/levels/axes.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py
index 34ce091..24536cc 100644
--- a/CodeEntropy/levels/axes.py
+++ b/CodeEntropy/levels/axes.py
@@ -265,14 +265,12 @@ def get_UA_axes(self, data_container, index: int, res_position):
             - If there are *no* bonds to other residues, use a custom principal axes
             from a moment-of-inertia (MOI) tensor that uses positions of heavy atoms
             only, but includes masses of heavy atom + bonded hydrogens.
-
             - If bonded to only one other residue, find edge heavy atom
             (i.e. heavy atom bonded to neighbour residue). Find all heavy atoms
             bonded to edge heavy atom and compute their average position.
             Find all other heavy atoms in residue and compute their average position.
             The three points are now used to obtain determine residue rotational axes.
             (see get_residue_custom_axes)
-
             - If bonded to at least two other residues, 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

From 4ab4534645d041e2aaa8b24179e1e6d85b77536d Mon Sep 17 00:00:00 2001
From: ioanaapapa <ioanaalexandrapapa@gmail.com>
Date: Tue, 25 Aug 2026 15:26:37 +0100
Subject: [PATCH 06/17] updated docstring

---
 CodeEntropy/levels/axes.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py
index 24536cc..c058975 100644
--- a/CodeEntropy/levels/axes.py
+++ b/CodeEntropy/levels/axes.py
@@ -280,8 +280,7 @@ def get_UA_axes(self, data_container, index: int, res_position):
             Identify heavy atoms in the residue/molecule of interest and choose
             the `index`-th heavy atom (where index corresponds to the bead index).
             Use bonded topology around that heavy atom to determine UA rotational
-            axes (see :meth:`get_bonded_axes`).
-            Compute a custom MOI tensor.
+            axes (see :meth:`get_bonded_axes`). Compute a custom MOI tensor.
 
         Args:
             data_container (MDAnalysis.Universe or AtomGroup):

From 97f72a7eb7972c29b513081cb90693fb7dce045a Mon Sep 17 00:00:00 2001
From: ioanaapapa <ioanaalexandrapapa@gmail.com>
Date: Tue, 25 Aug 2026 15:30:52 +0100
Subject: [PATCH 07/17] updated docstring

---
 CodeEntropy/levels/axes.py | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py
index c058975..4c59bd1 100644
--- a/CodeEntropy/levels/axes.py
+++ b/CodeEntropy/levels/axes.py
@@ -80,18 +80,19 @@ def get_residue_axes(
         - If bonded to only one other residue:
             * Translational axes are principal axes of data_container.
             * Find edge heavy atom (i.e. heavy atoms bonded to neighbour residue).
-              Find all heavy atoms bonded to edge heavy atom and compute their average
-              position.
-              Find all other heavy atoms in residue and compute their average position.
-              The three points are now used to obtain determine residue rotational axes.
-              (see get_residue_custom_axes)
+            Find all heavy atoms bonded to edge heavy atom and compute their average
+            position. Find all other heavy atoms in residue and compute their average
+            position.The three points are now used to obtain determine residue
+            rotational axes. (see get_residue_custom_axes)
+            Compute a custom MOI, using heavy atom positions and
+          heavy atom + hydrogen masses.
         - If bonded to at least two other residues:
             * Translational axes are principal axes of data_container.
             * 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 residue rotational axes.
-              (see get_residue_custom_axes).
-        Compute a custom MOI, using heavy atom positions and
+            and find the shortest chain between them: the backbone. Edge
+            atoms + backbone COM are used to determine residue rotational axes.
+            (see get_residue_custom_axes).
+            Compute a custom MOI, using heavy atom positions and
           heavy atom + hydrogen masses.
 
         Args:

From 66425475b9fb7a2579b9b8c17cc462dbdf348163 Mon Sep 17 00:00:00 2001
From: ioanaapapa <ioanaalexandrapapa@gmail.com>
Date: Tue, 25 Aug 2026 15:46:18 +0100
Subject: [PATCH 08/17] updated docstring

---
 CodeEntropy/levels/axes.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py
index 4c59bd1..2dda35b 100644
--- a/CodeEntropy/levels/axes.py
+++ b/CodeEntropy/levels/axes.py
@@ -77,23 +77,23 @@ def get_residue_axes(
               heavy atom + bonded hydrogens.
             * Set translational axes equal to rotational axes (as per the original
               code convention).
+
         - If bonded to only one other residue:
             * Translational axes are principal axes of data_container.
             * Find edge heavy atom (i.e. heavy atoms bonded to neighbour residue).
             Find all heavy atoms bonded to edge heavy atom and compute their average
             position. Find all other heavy atoms in residue and compute their average
             position.The three points are now used to obtain determine residue
-            rotational axes. (see get_residue_custom_axes)
-            Compute a custom MOI, using heavy atom positions and
-          heavy atom + hydrogen masses.
+            rotational axes. (see get_residue_custom_axes) Compute a custom MOI,
+            using heavy atom positions and heavy atom + hydrogen masses.
+
         - If bonded to at least two other residues:
             * Translational axes are principal axes of data_container.
             * 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 residue rotational axes.
-            (see get_residue_custom_axes).
-            Compute a custom MOI, using heavy atom positions and
-          heavy atom + hydrogen masses.
+            (see get_residue_custom_axes). Compute a custom MOI, using heavy
+            atom positions and heavy atom + hydrogen masses.
 
         Args:
             data_container (MDAnalysis.Universe or AtomGroup):

From ec7399545ad2121f43f893a235ac40b031226d88 Mon Sep 17 00:00:00 2001
From: ioanaapapa <ioanaalexandrapapa@gmail.com>
Date: Tue, 25 Aug 2026 15:52:19 +0100
Subject: [PATCH 09/17] updated docstring

---
 CodeEntropy/levels/axes.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py
index 2dda35b..1c3556a 100644
--- a/CodeEntropy/levels/axes.py
+++ b/CodeEntropy/levels/axes.py
@@ -80,8 +80,8 @@ def get_residue_axes(
 
         - If bonded to only one other residue:
             * Translational axes are principal axes of data_container.
-            * Find edge heavy atom (i.e. heavy atoms bonded to neighbour residue).
-            Find all heavy atoms bonded to edge heavy atom and compute their average
+            * Find edge heavy atom (i.e. heavy atoms bonded to next residue). Find
+            all heavy atoms bonded to edge heavy atom and compute their average
             position. Find all other heavy atoms in residue and compute their average
             position.The three points are now used to obtain determine residue
             rotational axes. (see get_residue_custom_axes) Compute a custom MOI,
@@ -90,8 +90,8 @@ def get_residue_axes(
         - If bonded to at least two other residues:
             * Translational axes are principal axes of data_container.
             * 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 residue rotational axes.
+            and find the shortest chain between them: the backbone. Edge atoms
+            + backbone COM are used to determine residue rotational axes.
             (see get_residue_custom_axes). Compute a custom MOI, using heavy
             atom positions and heavy atom + hydrogen masses.
 

From 7a929163bb764f621d8fe749d0a641ae6e5a8b64 Mon Sep 17 00:00:00 2001
From: ioanaapapa <ioanaalexandrapapa@gmail.com>
Date: Tue, 25 Aug 2026 15:53:53 +0100
Subject: [PATCH 10/17] updated docstring

---
 CodeEntropy/levels/axes.py | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py
index 1c3556a..1b27635 100644
--- a/CodeEntropy/levels/axes.py
+++ b/CodeEntropy/levels/axes.py
@@ -80,20 +80,20 @@ def get_residue_axes(
 
         - If bonded to only one other residue:
             * Translational axes are principal axes of data_container.
-            * Find edge heavy atom (i.e. heavy atoms bonded to next residue). Find
-            all heavy atoms bonded to edge heavy atom and compute their average
-            position. Find all other heavy atoms in residue and compute their average
-            position.The three points are now used to obtain determine residue
-            rotational axes. (see get_residue_custom_axes) Compute a custom MOI,
-            using heavy atom positions and heavy atom + hydrogen masses.
+            * Find edge heavy atom (i.e. heavy atoms bonded to neighbour residue). Find
+              all heavy atoms bonded to edge heavy atom and compute their average
+              position. Find all other heavy atoms in residue and compute their average
+              position.The three points are now used to obtain determine residue
+              rotational axes. (see get_residue_custom_axes) Compute a custom MOI,
+              using heavy atom positions and heavy atom + hydrogen masses.
 
         - If bonded to at least two other residues:
             * Translational axes are principal axes of data_container.
             * 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 residue rotational axes.
-            (see get_residue_custom_axes). Compute a custom MOI, using heavy
-            atom positions and heavy atom + hydrogen masses.
+              and find the shortest chain between them: the backbone. Edge atoms
+              + backbone COM are used to determine residue rotational axes.
+              (see get_residue_custom_axes). Compute a custom MOI, using heavy
+              atom positions and heavy atom + hydrogen masses.
 
         Args:
             data_container (MDAnalysis.Universe or AtomGroup):

From 11b68d589ce2e29f94995f454ba33afd4629cc90 Mon Sep 17 00:00:00 2001
From: ioanaapapa <ioanaalexandrapapa@gmail.com>
Date: Tue, 25 Aug 2026 16:10:13 +0100
Subject: [PATCH 11/17] updated documentation to reflect terminal residue
 changes

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

diff --git a/docs/science.rst b/docs/science.rst
index 7aa219d..b42af0a 100644
--- a/docs/science.rst
+++ b/docs/science.rst
@@ -70,9 +70,11 @@ The axes for this transformation are calculated for each bead in each time step.
 
 For the polymer level, the translational and rotational axes are defined as the principal axes of the molecule.
 
-For the residue level, there are two situations.
+For the residue level, there are three situations.
 When the residue is not bonded to any other residues, the translational and rotational axes are the principal axes of the molecule.
-When the residue is part of a larger polymer, the translational axes are the principal axes of the polymer, and the rotational axes are defined from the two heavy atoms bonded to neighbour residues(E1,E2) and the average position of all other backbone atoms in the residue (C). The backbone of a residue is defined as the shortest path between the two edge atoms of the residue, i.e. the two heavy atoms bonded to neighbour residues.The centre of rotation is located at the point where the perpendicular from C meets the E1-E2 vector.
+When the residue is part of a larger polymer and is not a terminus of that polymer, the translational axes are the principal axes of the polymer, and the rotational axes are defined from the two heavy atoms bonded to neighbour residues (E1,E2) and the average position of all other backbone atoms in the residue (C). The backbone of a residue is defined as the shortest path between the two edge atoms of the residue, i.e.the two heavy atoms bonded to neighbour residues.The centre of rotation (O) is located at the point where the perpendicular from C meets the E1-E2 vector.
+When the residue is part of a larger polymer and is a terminus of that polymer, the translational axes are the principal axes of the polymer, and the rotational axes are defined from the heavy atom bonded to a
+neighbour residue (E1), the average position of all heavy atoms bonded to E1 (C) and the average position of all other heavy atoms in the residue (E2). The centre of rotation (O) is defined the same as above forthe non-terminal residue case.
 
 For the united atom level, the translational axes are defined as the residue rotational axes and the rotational axes are defined from the average position of the bonds to neighbouring heavy atoms.
 If there are no bonds to other heavy atoms, the principal axes of the molecule are used.

From 5ff1611411854e57f751c604cd5473f4dcce397a Mon Sep 17 00:00:00 2001
From: ioanaapapa <ioanaalexandrapapa@gmail.com>
Date: Wed, 26 Aug 2026 14:12:05 +0100
Subject: [PATCH 12/17] update unit tests + fix atom selection bug

---
 CodeEntropy/levels/axes.py                 |  66 ++++++-----
 tests/unit/CodeEntropy/levels/test_axes.py | 122 +++++----------------
 2 files changed, 68 insertions(+), 120 deletions(-)

diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py
index 1b27635..1b040ba 100644
--- a/CodeEntropy/levels/axes.py
+++ b/CodeEntropy/levels/axes.py
@@ -156,21 +156,26 @@ def get_residue_axes(
 
             if len(edge_atom_set) == 1:
                 edge_atom = edge_atom_set[0]
-                bonded_atoms = uas.select_atoms(f"bonded index {edge_atom.index}")
+                bonded_atoms = residue.select_atoms(
+                    f"(mass 2 to 999) and bonded index {edge_atom.index}"
+                )
                 # find the average position of heavy atoms bonded to edge atom
-                average_bonded_atom = np.zeros(3)
-                for atom in bonded_atoms:
-                    average_bonded_atom += atom.position
-                average_bonded_atom /= len(bonded_atoms)
+                if len(bonded_atoms) > 0:
+                    average_bonded_atom = np.zeros(3)
+                    for bonded_atom in bonded_atoms:
+                        average_bonded_atom += bonded_atom.position
+                    average_bonded_atom /= len(bonded_atoms)
                 # find the average position of all other heavy atoms in residue
                 other_atoms = []
                 for atom in uas:
                     if atom != edge_atom and atom not in bonded_atoms:
                         other_atoms.append(atom)
-                average_other_atoms = np.zeros(3)
-                for atom in other_atoms:
-                    average_other_atoms += atom.position
-                average_other_atoms /= len(other_atoms)
+                if len(other_atoms) > 0:
+                    average_other_atoms = np.zeros(3)
+                    for atom in other_atoms:
+                        average_other_atoms += atom.position
+                    average_other_atoms /= len(average_other_atoms)
+
                 rot_center, rot_axes = self.get_residue_custom_axes(
                     [edge_atom.position, average_other_atoms], average_bonded_atom
                 )
@@ -179,9 +184,10 @@ def get_residue_axes(
                 edges = [edge_atom_set[0].position, edge_atom_set[1].position]
                 backbone = self.get_chain(residue, edge_atom_set[0], edge_atom_set[1])
                 backbone_center = np.zeros(3)
-                for heavy_atom in backbone:
-                    backbone_center += heavy_atom.position
-                backbone_center = backbone_center / len(backbone)
+                if len(backbone) > 0:
+                    for heavy_atom in backbone:
+                        backbone_center += heavy_atom.position
+                    backbone_center /= len(backbone)
                 rot_center, rot_axes = self.get_residue_custom_axes(
                     edges, backbone_center
                 )
@@ -323,7 +329,7 @@ def get_UA_axes(self, data_container, index: int, res_position):
                         residue = data_container.residues[0]
                         resindex = residue.resindex
                         resindex_next = resindex + 1
-                        edge_atom = data_container.select_atoms(
+                        edge_atom_set = data_container.select_atoms(
                             f"resindex {resindex} and bonded resindex {resindex_next}"
                         )
                     else:
@@ -331,29 +337,32 @@ def get_UA_axes(self, data_container, index: int, res_position):
                         residue = data_container.residues[1]
                         resindex = residue.resindex
                         resindex_prev = resindex - 1
-                        edge_atom = data_container.select_atoms(
+                        edge_atom_set = data_container.select_atoms(
                             f"resindex {resindex} and bonded resindex {resindex_prev}"
                         )
+                    edge_atom = edge_atom_set[0]
                     residue_heavy_atoms = residue.atoms.select_atoms("mass 2 to 999")
-                    bonded_atoms = residue_heavy_atoms.select_atoms(
-                        f"bonded index {edge_atom[0].index}"
+                    bonded_atoms = residue.atoms.select_atoms(
+                        f"(mass 2 to 999) and bonded index {edge_atom.index}"
                     )
                     # find the average position of heavy atoms bonded to edge atom
-                    average_bonded_atom = np.zeros(3)
-                    for atom in bonded_atoms:
-                        average_bonded_atom += atom.position
-                    average_bonded_atom /= len(bonded_atoms)
+                    if len(bonded_atoms) > 0:
+                        average_bonded_atom = np.zeros(3)
+                        for atom in bonded_atoms:
+                            average_bonded_atom += atom.position
+                        average_bonded_atom /= len(bonded_atoms)
                     # find the average position of all other heavy atoms in residue
                     other_atoms = []
                     for atom in residue_heavy_atoms:
                         if atom != edge_atom and atom not in bonded_atoms:
                             other_atoms.append(atom)
                     average_other_atoms = np.zeros(3)
-                    for atom in other_atoms:
-                        average_other_atoms += atom.position
-                    average_other_atoms /= len(other_atoms)
+                    if len(other_atoms) > 0:
+                        for atom in other_atoms:
+                            average_other_atoms += atom.position
+                        average_other_atoms /= len(other_atoms)
                     trans_center, trans_axes = self.get_residue_custom_axes(
-                        [edge_atom.positions[0], average_other_atoms],
+                        [edge_atom.position, average_other_atoms],
                         average_bonded_atom,
                     )
                 else:
@@ -371,10 +380,11 @@ def get_UA_axes(self, data_container, index: int, res_position):
 
                     edges = edge_set.positions
                     backbone = self.get_chain(residue, edge_set[0], edge_set[1])
-                    backbone_center = np.zeros(3)
-                    for heavy_atom in backbone:
-                        backbone_center += heavy_atom.position
-                    backbone_center = backbone_center / len(backbone)
+                    if len(backbone) > 0:
+                        backbone_center = np.zeros(3)
+                        for heavy_atom in backbone:
+                            backbone_center += heavy_atom.position
+                        backbone_center /= len(backbone)
 
                     trans_center, trans_axes = self.get_residue_custom_axes(
                         edges, backbone_center
diff --git a/tests/unit/CodeEntropy/levels/test_axes.py b/tests/unit/CodeEntropy/levels/test_axes.py
index c926f95..f9f0d71 100644
--- a/tests/unit/CodeEntropy/levels/test_axes.py
+++ b/tests/unit/CodeEntropy/levels/test_axes.py
@@ -1220,7 +1220,7 @@ def test_get_residue_axes_custom_path(monkeypatch):
 
     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
+        [edge_atoms[0].position, edge_atoms[1].position], backbone_center
     )
 
     assert rot_center.shape == (3,)
@@ -1337,90 +1337,33 @@ def _select_atoms(q):
     assert np.allclose(moi, np.array([1, 1, 1]))
 
 
-def test_get_residue_bonded_axes_first_resid(monkeypatch):
+def test_get_residue_bonded_axes_terminal_resid(monkeypatch):
     ax = AxesCalculator()
     u = MagicMock()
     u.dimensions = np.array([10.0, 10.0, 10.0, 90, 90, 90])
     monkeypatch.setattr("CodeEntropy.levels.axes.make_whole", lambda _ag: None)
     residue = u.select_atoms("resindex 0")
     residue.__len__.return_value = 3
-    residue.atoms = _FakeAtomGroup(
+    uas = _FakeAtomGroup(
         [
             _atom(index=0, mass=12.0, pos=[1, 0, 0]),
             _atom(index=1, mass=12.0, pos=[0, 1, 0]),
             _atom(index=2, mass=12.0, pos=[0, 0, 0]),
         ]
     )
-    edge_atom_set = _FakeAtomGroup(
-        [
-            _atom(index=2, mass=12.0, pos=[0, 0, 0]),
-        ]
-    )
-
-    def _select_atoms(q):
-        if q.endswith("(bonded resindex -1 or resindex 1)"):
-            return edge_atom_set
-
-    backbone_atom = residue.atoms[1]
-    u.atoms.principal_axes.return_value = np.eye(3)
-    u.atoms.select_atoms.side_effect = _select_atoms
-    monkeypatch.setattr(ax, "get_chain", backbone_atom)
-    monkeypatch.setattr(
-        ax,
-        "get_custom_residue_moment_of_inertia",
-        lambda center_of_mass, positions, masses, custom_rot_axes, dimensions: np.array(
-            [1, 1, 1]
-        ),
-    )
-
-    trans_axes, rot_axes, rot_center, moi = ax.get_residue_axes(
-        u, index=0, relative_index=0
-    )
-
-    assert len(edge_atom_set) == 1
-    assert np.allclose(trans_axes, np.eye(3))
-    assert rot_axes.shape == (3, 3)
-    assert rot_center.shape == (3,)
-    assert np.allclose(moi, np.array([1, 1, 1]))
-
-
-def test_get_residue_bonded_axes_last_resid(monkeypatch):
-    ax = AxesCalculator()
-    u = MagicMock()
-    u.dimensions = np.array([10.0, 10.0, 10.0, 90, 90, 90])
-    monkeypatch.setattr("CodeEntropy.levels.axes.make_whole", lambda _ag: None)
-    residue = u.select_atoms("resindex 2")
-    residue.__len__.return_value = 3
-    heavy_atoms = _FakeAtomGroup(
-        [
-            _atom(index=4, mass=12.0, pos=[1, 0, 0]),
-            _atom(index=5, mass=12.0, pos=[0, 1, 0]),
-            _atom(index=6, mass=12.0, pos=[0, 0, 0]),
-        ]
-    )
-    edge_atom_set = _FakeAtomGroup(
-        [
-            _atom(index=4, mass=12.0, pos=[0, 0, 0]),
-        ]
-    )
 
     def _select_atoms(q):
         if q == "mass 2 to 999":
-            # return heavy atoms group
-            return heavy_atoms
-        if q.endswith("(bonded resindex 1 or resindex 3)"):
-            return edge_atom_set
-        if q == ("(mass 2 to 999) and bonded index 6"):
-            return [heavy_atoms[1]]
-        if q == ("(mass 2 to 999) and bonded index 5"):
-            return [heavy_atoms[0], heavy_atoms[2]]
+            return uas
+        if q.startswith("(mass 2 to 999) and bonded"):
+            return [uas[1]]
+        if q.startswith("resindex 0 and (bonded resindex"):
+            return [uas[2]]
 
-    backbone_atom = heavy_atoms[1]
     u.atoms.principal_axes.return_value = np.eye(3)
     u.atoms.select_atoms.side_effect = _select_atoms
     residue.select_atoms.side_effect = _select_atoms
-    residue.atoms.select_atoms.side_effect = _select_atoms
-    monkeypatch.setattr(ax, "get_chain", backbone_atom)
+
     monkeypatch.setattr(
         ax,
         "get_custom_residue_moment_of_inertia",
@@ -1430,10 +1373,9 @@ def _select_atoms(q):
     )
 
     trans_axes, rot_axes, rot_center, moi = ax.get_residue_axes(
-        u, index=2, relative_index=0
+        u, index=0, relative_index=0
     )
 
-    assert len(edge_atom_set) == 1
     assert np.allclose(trans_axes, np.eye(3))
     assert rot_axes.shape == (3, 3)
     assert rot_center.shape == (3,)
@@ -1473,8 +1415,10 @@ def _select_atoms(q):
             return [heavy_atoms[1]]
 
     residue_group.select_atoms.side_effect = _select_atoms
+    residue.select_atoms.side_effect = _select_atoms
     residue.atoms.select_atoms.side_effect = _select_atoms
-    monkeypatch.setattr(ax, "get_chain", lambda residue, first, last: heavy_atoms[1])
+
+    monkeypatch.setattr(ax, "get_chain", lambda residue, first, last: [heavy_atoms[1]])
     monkeypatch.setattr(
         ax,
         "get_bonded_axes",
@@ -1560,25 +1504,23 @@ def test_get_ua_axes_bonded_axes_first_resid(monkeypatch):
 
     residue.atoms[0] = heavy_atoms[0]
     residue.atoms[0].position = heavy_atoms[0].position
-    edge_atom_set = _FakeAtomGroup(
-        [
-            _atom(index=2, mass=12.0, pos=(0, 0, 1)),
-        ],
-    )
+    edge_atom_set = [heavy_atoms[2]]
+    bonded_atoms = [heavy_atoms[1]]
 
     def _select_atoms(q):
         if q == "mass 2 to 999":
             # return heavy atoms group
             return heavy_atoms
+        if q.startswith("index"):
+            return [heavy_atoms[0]]
         if q.startswith("resindex "):
             return edge_atom_set
-        if q.startswith("index "):
-            return [heavy_atoms[0]]
+        if q.startswith("(mass 2 to 999) and bonded index "):
+            return bonded_atoms
 
     residue_group.select_atoms.side_effect = _select_atoms
     residue.atoms.select_atoms.side_effect = _select_atoms
-    edge_atom_set.atoms = [edge_atom_set[0]]
-    monkeypatch.setattr(ax, "get_chain", lambda residue, first, last: heavy_atoms[1])
+
     monkeypatch.setattr(
         ax,
         "get_bonded_axes",
@@ -1598,7 +1540,7 @@ def test_get_ua_axes_bonded_axes_last_resid(monkeypatch):
     ax = AxesCalculator()
     residue_group = MagicMock()
     residue_group.__len__ = 2
-    residue = residue_group.residues[1]
+    residue = residue_group.residues[0]
     heavy_atoms = _FakeAtomGroup(
         [
             _atom(index=0, mass=12.0, pos=(1, 0, 0)),
@@ -1607,29 +1549,25 @@ def test_get_ua_axes_bonded_axes_last_resid(monkeypatch):
         ],
     )
 
-    edge_atom_set = _FakeAtomGroup(
-        [
-            _atom(index=0, mass=12.0, pos=(1, 0, 0)),
-        ],
-    )
+    residue.atoms[0] = heavy_atoms[0]
+    residue.atoms[0].position = heavy_atoms[0].position
+    edge_atom_set = [heavy_atoms[0]]
+    bonded_atoms = [heavy_atoms[1]]
 
     def _select_atoms(q):
         if q == "mass 2 to 999":
             # return heavy atoms group
             return heavy_atoms
+        if q.startswith("index"):
+            return [heavy_atoms[0]]
         if q.startswith("resindex "):
             return edge_atom_set
-        if q.startswith("index "):
-            return [heavy_atoms[0]]
-        if q == ("(mass 2 to 999) and bonded index 2"):
-            return [heavy_atoms[1]]
-        if q == ("(mass 2 to 999) and bonded index 1"):
-            return [heavy_atoms[0], heavy_atoms[2]]
+        if q.startswith("(mass 2 to 999) and bonded index "):
+            return bonded_atoms
 
     residue_group.select_atoms.side_effect = _select_atoms
     residue.atoms.select_atoms.side_effect = _select_atoms
-    edge_atom_set.atoms = [edge_atom_set[0]]
-    monkeypatch.setattr(ax, "get_chain", lambda residue, first, last: heavy_atoms[1])
+
     monkeypatch.setattr(
         ax,
         "get_bonded_axes",

From 648c5d034eaec9991e3ad6411b5687f7f8c6ddb3 Mon Sep 17 00:00:00 2001
From: Ioana Papa <iapapa1@sheffield.ac.uk>
Date: Wed, 26 Aug 2026 16:06:31 +0100
Subject: [PATCH 13/17] added edge cases for two points

---
 CodeEntropy/levels/axes.py | 69 +++++++++++++++++++++++++++-----------
 1 file changed, 50 insertions(+), 19 deletions(-)

diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py
index 1b040ba..aa3dc75 100644
--- a/CodeEntropy/levels/axes.py
+++ b/CodeEntropy/levels/axes.py
@@ -84,15 +84,23 @@ def get_residue_axes(
               all heavy atoms bonded to edge heavy atom and compute their average
               position. Find all other heavy atoms in residue and compute their average
               position.The three points are now used to obtain determine residue
-              rotational axes. (see get_residue_custom_axes) Compute a custom MOI,
-              using heavy atom positions and heavy atom + hydrogen masses.
+              rotational axes. (see get_residue_custom_axes) If there are
+              only two heavy atoms in the residue/all heavy atoms are bonded to edge
+              atom, x-axis is set along the vector between the edge atom and average
+              position of bonded atoms, y-axis is arbitrary and z-axis is paralel
+              to the two. This is the same as case 2 in get_bonded_axes. Compute a
+              custom MOI, using heavy atom positions and heavy atom + hydrogen masses.
 
         - If bonded to at least two other residues:
             * Translational axes are principal axes of data_container.
             * 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 residue rotational axes.
-              (see get_residue_custom_axes). Compute a custom MOI, using heavy
+              (see get_residue_custom_axes). If the two edge heavy atoms
+              are bonded to each other (i.e. there is no backbone), x-axis is set
+              along the vector between the edge atom and average position of bonded
+              atoms, y-axis is arbitrary and z-axis is paralel to the two. This is the
+              same as case 2 in get_bonded_axes. Compute a custom MOI, using heavy
               atom positions and heavy atom + hydrogen masses.
 
         Args:
@@ -175,10 +183,13 @@ def get_residue_axes(
                     for atom in other_atoms:
                         average_other_atoms += atom.position
                     average_other_atoms /= len(average_other_atoms)
-
-                rot_center, rot_axes = self.get_residue_custom_axes(
-                    [edge_atom.position, average_other_atoms], average_bonded_atom
-                )
+                    rot_center, rot_axes = self.get_residue_custom_axes(
+                        [edge_atom.position, average_other_atoms], average_bonded_atom
+                    )
+                else:
+                    rot_center, rot_axes = self.get_custom_axes(
+                        a=edge_atom.position, b=[average_bonded_atom], c=np.zeros(3)
+                    )
 
             else:
                 edges = [edge_atom_set[0].position, edge_atom_set[1].position]
@@ -188,9 +199,13 @@ def get_residue_axes(
                     for heavy_atom in backbone:
                         backbone_center += heavy_atom.position
                     backbone_center /= len(backbone)
-                rot_center, rot_axes = self.get_residue_custom_axes(
-                    edges, backbone_center
-                )
+                    rot_center, rot_axes = self.get_residue_custom_axes(
+                        edges, backbone_center
+                    )
+                else:
+                    rot_center, rot_axes = self.get_custom_axes(
+                        a=edges[0], b=[edges[1]], c=np.zeros(3)
+                    )
 
             moment_of_inertia = self.get_custom_residue_moment_of_inertia(
                 center_of_mass=rot_center,
@@ -272,6 +287,11 @@ def get_UA_axes(self, data_container, index: int, res_position):
             - If there are *no* bonds to other residues, use a custom principal axes
             from a moment-of-inertia (MOI) tensor that uses positions of heavy atoms
             only, but includes masses of heavy atom + bonded hydrogens.
+            - If bonded to only one other residue and there are only two heavy atoms
+            in the residue/all heavy atoms are bonded to edge atom,
+            x-axis is set along the vector between the edge atom and average position
+            of bonded atoms, y-axis is arbitrary and z-axis is paralel to the two.
+            This is the same as case 2 in get_bonded_axes.
             - If bonded to only one other residue, find edge heavy atom
             (i.e. heavy atom bonded to neighbour residue). Find all heavy atoms
             bonded to edge heavy atom and compute their average position.
@@ -282,6 +302,11 @@ def get_UA_axes(self, data_container, index: int, res_position):
             (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 residue rotational axes. (see get_residue_custom_axes).
+            - If bonded to at least two other residues and the two edge heavy atoms
+            are bonded to each other (i.e. there is no backbone), x-axis is set along
+            the vector between the edge atom and average position of bonded atoms,
+            y-axis is arbitrary and z-axis is paralel to the two. This is the same
+            as case 2 in get_bonded_axes.
 
         - Rotational axes:
             Identify heavy atoms in the residue/molecule of interest and choose
@@ -361,10 +386,14 @@ def get_UA_axes(self, data_container, index: int, res_position):
                         for atom in other_atoms:
                             average_other_atoms += atom.position
                         average_other_atoms /= len(other_atoms)
-                    trans_center, trans_axes = self.get_residue_custom_axes(
-                        [edge_atom.position, average_other_atoms],
-                        average_bonded_atom,
-                    )
+                        trans_center, trans_axes = self.get_residue_custom_axes(
+                            [edge_atom.position, average_other_atoms],
+                            average_bonded_atom,
+                        )
+                    else:
+                        trans_center, trans_axes = self.get_custom_axes(
+                            a=edge_atom.position, b=[average_bonded_atom], c=np.zeros(3)
+                        )
                 else:
                     # between 2 residues
                     residue = data_container.residues[1]
@@ -377,7 +406,6 @@ def get_UA_axes(self, data_container, index: int, res_position):
                         f"(bonded resindex {resindex_prev} or "
                         f"resindex {resindex_next})"
                     )
-
                     edges = edge_set.positions
                     backbone = self.get_chain(residue, edge_set[0], edge_set[1])
                     if len(backbone) > 0:
@@ -385,10 +413,13 @@ def get_UA_axes(self, data_container, index: int, res_position):
                         for heavy_atom in backbone:
                             backbone_center += heavy_atom.position
                         backbone_center /= len(backbone)
-
-                    trans_center, trans_axes = self.get_residue_custom_axes(
-                        edges, backbone_center
-                    )
+                        trans_center, trans_axes = self.get_residue_custom_axes(
+                            edges, backbone_center
+                        )
+                    else:
+                        trans_center, trans_axes = self.get_custom_axes(
+                            a=edges[0], b=[edges[1]], c=np.zeros(3)
+                        )
 
             # look for heavy atoms in residue of interest
             heavy_atom_indices = []

From e0272a853677108ad693c3fc8c4787b095f2b04e Mon Sep 17 00:00:00 2001
From: Ioana Papa <iapapa1@sheffield.ac.uk>
Date: Wed, 26 Aug 2026 17:03:04 +0100
Subject: [PATCH 14/17] fix centre for 2 heavy atoms situation

---
 CodeEntropy/levels/axes.py | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py
index aa3dc75..102cb14 100644
--- a/CodeEntropy/levels/axes.py
+++ b/CodeEntropy/levels/axes.py
@@ -187,7 +187,8 @@ def get_residue_axes(
                         [edge_atom.position, average_other_atoms], average_bonded_atom
                     )
                 else:
-                    rot_center, rot_axes = self.get_custom_axes(
+                    rot_center = edge_atom.position
+                    rot_axes = self.get_custom_axes(
                         a=edge_atom.position, b=[average_bonded_atom], c=np.zeros(3)
                     )
 
@@ -203,8 +204,9 @@ def get_residue_axes(
                         edges, backbone_center
                     )
                 else:
-                    rot_center, rot_axes = self.get_custom_axes(
-                        a=edges[0], b=[edges[1]], c=np.zeros(3)
+                    rot_center = (edges[0] + edges[1]) / 2
+                    rot_axes = self.get_custom_axes(
+                        a=rot_center, b=[edges[1]], c=np.zeros(3)
                     )
 
             moment_of_inertia = self.get_custom_residue_moment_of_inertia(
@@ -391,7 +393,8 @@ def get_UA_axes(self, data_container, index: int, res_position):
                             average_bonded_atom,
                         )
                     else:
-                        trans_center, trans_axes = self.get_custom_axes(
+                        trans_center = edge_atom.position
+                        trans_axes = self.get_custom_axes(
                             a=edge_atom.position, b=[average_bonded_atom], c=np.zeros(3)
                         )
                 else:
@@ -417,8 +420,9 @@ def get_UA_axes(self, data_container, index: int, res_position):
                             edges, backbone_center
                         )
                     else:
-                        trans_center, trans_axes = self.get_custom_axes(
-                            a=edges[0], b=[edges[1]], c=np.zeros(3)
+                        trans_center = (edges[0] + edges[1]) / 2
+                        trans_axes = self.get_custom_axes(
+                            a=trans_center, b=[edges[1]], c=np.zeros(3)
                         )
 
             # look for heavy atoms in residue of interest

From 127c097c477e41795611a5c1069a2c3b80fc5ee8 Mon Sep 17 00:00:00 2001
From: Ioana Papa <iapapa1@sheffield.ac.uk>
Date: Thu, 27 Aug 2026 09:32:12 +0100
Subject: [PATCH 15/17] added edge case unit test

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

diff --git a/tests/unit/CodeEntropy/levels/test_axes.py b/tests/unit/CodeEntropy/levels/test_axes.py
index f9f0d71..f64b48c 100644
--- a/tests/unit/CodeEntropy/levels/test_axes.py
+++ b/tests/unit/CodeEntropy/levels/test_axes.py
@@ -1269,7 +1269,7 @@ def _select_atoms(q):
 
     u.atoms.select_atoms.side_effect = _select_atoms
     u.atoms.principal_axes.return_value = np.eye(3)
-    monkeypatch.setattr(ax, "get_chain", backbone_atom)
+    monkeypatch.setattr(ax, "get_chain", lambda residue, first, last: [backbone_atom])
     monkeypatch.setattr(
         ax,
         "get_custom_residue_moment_of_inertia",
@@ -1639,3 +1639,49 @@ def _sel(q):
 
     with pytest.raises(ValueError):
         ax.get_UA_axes(u, index=0, res_position=None)
+
+
+def get_residue_bonded_axes_terminal_2_points(monkeypatch):
+    ax = AxesCalculator()
+    u = MagicMock()
+    u.dimensions = np.array([10.0, 10.0, 10.0, 90, 90, 90])
+    monkeypatch.setattr("CodeEntropy.levels.axes.make_whole", lambda _ag: None)
+    residue = u.select_atoms("resindex 0")
+    residue.__len__.return_value = 3
+    uas = _FakeAtomGroup(
+        [
+            _atom(index=0, mass=12.0, pos=[1, 0, 0]),
+            _atom(index=1, mass=12.0, pos=[0, 1, 0]),
+        ]
+    )
+
+    def _select_atoms(q):
+        if q == "mass 2 to 999":
+            return uas
+        if q.startswith("(mass 2 to 999) and bonded"):
+            # the bonded atom
+            return [uas[1]]
+        if q.startswith("resindex 0 and (bonded resindex"):
+            # edge atom
+            return [uas[2]]
+
+    u.atoms.principal_axes.return_value = np.eye(3)
+    u.atoms.select_atoms.side_effect = _select_atoms
+    residue.select_atoms.side_effect = _select_atoms
+
+    monkeypatch.setattr(
+        ax,
+        "get_custom_residue_moment_of_inertia",
+        lambda center_of_mass, positions, masses, custom_rot_axes, dimensions: np.array(
+            [1, 1, 1]
+        ),
+    )
+
+    trans_axes, rot_axes, rot_center, moi = ax.get_residue_axes(
+        u, index=0, relative_index=0
+    )
+
+    assert np.allclose(trans_axes, np.eye(3))
+    assert rot_axes.shape == (3, 3)
+    assert np.allclose(rot_center, [1, 1, 0])
+    assert np.allclose(moi, np.array([1, 1, 1]))

From 2ea2de5e27455aa6baa2bb0218f2fa647692ad6c Mon Sep 17 00:00:00 2001
From: ioanaapapa <ioanaalexandrapapa@gmail.com>
Date: Thu, 27 Aug 2026 11:50:30 +0100
Subject: [PATCH 16/17] unit tests for edge cases

---
 tests/unit/CodeEntropy/levels/test_axes.py | 164 +++++++++++++++++++--
 1 file changed, 148 insertions(+), 16 deletions(-)

diff --git a/tests/unit/CodeEntropy/levels/test_axes.py b/tests/unit/CodeEntropy/levels/test_axes.py
index f64b48c..2e4b52f 100644
--- a/tests/unit/CodeEntropy/levels/test_axes.py
+++ b/tests/unit/CodeEntropy/levels/test_axes.py
@@ -1641,7 +1641,139 @@ def _sel(q):
         ax.get_UA_axes(u, index=0, res_position=None)
 
 
-def get_residue_bonded_axes_terminal_2_points(monkeypatch):
+def test_get_ua_axes_bonded_terminal_2_points(monkeypatch):
+    ax = AxesCalculator()
+    residue_group = MagicMock()
+    residue_group.__len__ = 2
+    residue = residue_group.residues[1]
+    heavy_atoms = _FakeAtomGroup(
+        [
+            _atom(index=0, mass=12.0, pos=(1, 0, 0)),
+            _atom(index=1, mass=12.0, pos=(0, 1, 0)),
+        ],
+    )
+    edge_atom_set = [heavy_atoms[0]]
+    bonded_atoms = [heavy_atoms[1]]
+
+    def _select_atoms(q):
+        if q == "mass 2 to 999":
+            # return heavy atoms group
+            return heavy_atoms
+        if q.startswith("index"):
+            return [heavy_atoms[0]]
+        if q.startswith("resindex "):
+            return edge_atom_set
+        if q.startswith("(mass 2 to 999) and bonded index "):
+            return bonded_atoms
+
+    residue_group.select_atoms.side_effect = _select_atoms
+    residue.atoms.select_atoms.side_effect = _select_atoms
+
+    monkeypatch.setattr(
+        ax,
+        "get_bonded_axes",
+        lambda system, atom, dimensions: (np.eye(3), np.array([1.0, 1.0, 1.0])),
+    )
+
+    monkeypatch.setattr(ax, "get_custom_axes", lambda a, b, c: 2 * np.eye(3))
+
+    trans_axes, rot_axes, rot_center, moi = ax.get_UA_axes(
+        data_container=residue_group, index=0, res_position=1
+    )
+
+    assert np.allclose(trans_axes, 2 * np.eye(3))
+    assert np.allclose(rot_axes, np.eye(3))
+    assert np.allclose(rot_center, [1, 0, 0])
+    assert np.allclose(moi, np.array([1, 1, 1]))
+
+
+def test_get_ua_axes_non_terminal_2_atoms(monkeypatch):
+    ax = AxesCalculator()
+    residue_group = MagicMock()
+    residue_group.__len__ = 3
+    residue = residue_group.residues[1]
+    heavy_atoms = _FakeAtomGroup(
+        [
+            _atom(index=0, mass=12.0, pos=(1, 1, 1)),
+            _atom(index=1, mass=12.0, pos=(3, 3, 3)),
+        ],
+    )
+
+    def _select_atoms(q):
+        if q == "mass 2 to 999":
+            # return heavy atoms group
+            return heavy_atoms
+        if q.startswith("index"):
+            return [heavy_atoms[0]]
+        if q.startswith("resindex "):
+            return heavy_atoms
+
+    residue_group.select_atoms.side_effect = _select_atoms
+    residue.atoms.select_atoms.side_effect = _select_atoms
+    monkeypatch.setattr(
+        ax,
+        "get_bonded_axes",
+        lambda system, atom, dimensions: (np.eye(3), 3 * np.eye(3)),
+    )
+    monkeypatch.setattr(ax, "get_custom_axes", lambda a, b, c: 2 * np.eye(3))
+    monkeypatch.setattr(ax, "get_chain", lambda residue, first, last: [])
+    trans_axes, rot_axes, rot_center, moi = ax.get_UA_axes(
+        data_container=residue_group, index=0, res_position=0
+    )
+
+    assert np.allclose(trans_axes, 2 * np.eye(3))
+    assert np.allclose(rot_axes, np.eye(3))
+    assert np.allclose(rot_center, [1, 1, 1])
+    assert np.allclose(moi, 3 * np.eye(3))
+
+
+def test_get_residue_axes_non_terminal_2_atoms(monkeypatch):
+    ax = AxesCalculator()
+    u = MagicMock()
+    u.dimensions = np.array([10.0, 10.0, 10.0, 90, 90, 90])
+    monkeypatch.setattr("CodeEntropy.levels.axes.make_whole", lambda _ag: None)
+    residue = u.select_atoms("resindex 5")
+    residue.__len__.return_value = 2
+    print(f"The residue should be: {residue}")
+    u.atoms.principal_axes.return_value = np.eye(3)
+    uas = _FakeAtomGroup(
+        [
+            _atom(index=0, mass=12.0, pos=(1, 1, 1)),
+            _atom(index=1, mass=12.0, pos=(3, 3, 3)),
+        ],
+    )
+
+    def _select_atoms(q):
+        if q == "mass 2 to 999":
+            return uas
+        if q.startswith("resindex 5 and (bonded resindex"):
+            return uas
+
+    u.atoms.select_atoms.side_effect = _select_atoms
+    residue.select_atoms.side_effect = residue
+    monkeypatch.setattr(ax, "get_chain", lambda residue, first, last: [])
+    monkeypatch.setattr(ax, "get_custom_axes", lambda a, b, c: 2 * np.eye(3))
+    monkeypatch.setattr(
+        ax,
+        "get_custom_residue_moment_of_inertia",
+        lambda center_of_mass, positions, masses, custom_rot_axes, dimensions: np.array(
+            [1, 1, 1]
+        ),
+    )
+
+    trans_axes, rot_axes, rot_center, moi = ax.get_residue_axes(
+        data_container=u,
+        index=5,
+        relative_index=0,
+    )
+
+    assert np.allclose(trans_axes, np.eye(3))
+    assert np.allclose(rot_axes, 2 * np.eye(3))
+    assert np.allclose(rot_center, [2, 2, 2])
+    assert np.allclose(moi, [1, 1, 1])
+
+
+def test_get_residue_axes_terminal_2_atoms(monkeypatch):
     ax = AxesCalculator()
     u = MagicMock()
     u.dimensions = np.array([10.0, 10.0, 10.0, 90, 90, 90])
@@ -1650,25 +1782,25 @@ def get_residue_bonded_axes_terminal_2_points(monkeypatch):
     residue.__len__.return_value = 3
     uas = _FakeAtomGroup(
         [
-            _atom(index=0, mass=12.0, pos=[1, 0, 0]),
-            _atom(index=1, mass=12.0, pos=[0, 1, 0]),
-        ]
+            _atom(index=0, mass=12.0, pos=(1, 0, 0)),
+            _atom(index=1, mass=12.0, pos=(0, 1, 0)),
+            _atom(index=2, mass=12.0, pos=(0, 0, 1)),
+        ],
     )
+    u.atoms.principal_axes.return_value = np.eye(3)
 
     def _select_atoms(q):
         if q == "mass 2 to 999":
             return uas
-        if q.startswith("(mass 2 to 999) and bonded"):
-            # the bonded atom
-            return [uas[1]]
         if q.startswith("resindex 0 and (bonded resindex"):
-            # edge atom
+            # the edge atom
             return [uas[2]]
+        if q.startswith("(mass 2 to 999) and bonded index "):
+            return uas[0:2]
 
-    u.atoms.principal_axes.return_value = np.eye(3)
     u.atoms.select_atoms.side_effect = _select_atoms
     residue.select_atoms.side_effect = _select_atoms
-
+    monkeypatch.setattr(ax, "get_custom_axes", lambda a, b, c: 2 * np.eye(3))
     monkeypatch.setattr(
         ax,
         "get_custom_residue_moment_of_inertia",
@@ -1676,12 +1808,12 @@ def _select_atoms(q):
             [1, 1, 1]
         ),
     )
-
     trans_axes, rot_axes, rot_center, moi = ax.get_residue_axes(
-        u, index=0, relative_index=0
+        data_container=u,
+        index=0,
+        relative_index=0,
     )
-
     assert np.allclose(trans_axes, np.eye(3))
-    assert rot_axes.shape == (3, 3)
-    assert np.allclose(rot_center, [1, 1, 0])
-    assert np.allclose(moi, np.array([1, 1, 1]))
+    assert np.allclose(rot_axes, 2 * np.eye(3))
+    assert np.allclose(rot_center, [0, 0, 1])
+    assert np.allclose(moi, [1, 1, 1])

From 95dce270be301e7120f9b3f2c73d28b6dc8ff1ff Mon Sep 17 00:00:00 2001
From: ioanaapapa <ioanaalexandrapapa@gmail.com>
Date: Thu, 27 Aug 2026 13:06:01 +0100
Subject: [PATCH 17/17] changed sign of x axis for consistency with
 non-terminal residues

---
 CodeEntropy/levels/axes.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py
index 102cb14..6497337 100644
--- a/CodeEntropy/levels/axes.py
+++ b/CodeEntropy/levels/axes.py
@@ -206,7 +206,7 @@ def get_residue_axes(
                 else:
                     rot_center = (edges[0] + edges[1]) / 2
                     rot_axes = self.get_custom_axes(
-                        a=rot_center, b=[edges[1]], c=np.zeros(3)
+                        a=rot_center, b=[edges[0]], c=np.zeros(3)
                     )
 
             moment_of_inertia = self.get_custom_residue_moment_of_inertia(
@@ -422,7 +422,7 @@ def get_UA_axes(self, data_container, index: int, res_position):
                     else:
                         trans_center = (edges[0] + edges[1]) / 2
                         trans_axes = self.get_custom_axes(
-                            a=trans_center, b=[edges[1]], c=np.zeros(3)
+                            a=trans_center, b=[edges[0]], c=np.zeros(3)
                         )
 
             # look for heavy atoms in residue of interest