Residue rotation and UA translation axes for terminal residues

CodeEntropy/levels/axes.py

@@ -73,17 +73,35 @@ 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 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). 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 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 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). 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:
             data_container (MDAnalysis.Universe or AtomGroup):
@@ -145,30 +163,51 @@ 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]
+                edge_atom = edge_atom_set[0]
+                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
+                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)
+                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
                     )
                 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)
+                    rot_center = edge_atom.position
+                    rot_axes = self.get_custom_axes(
+                        a=edge_atom.position, b=[average_bonded_atom], c=np.zeros(3)
+                    )
+
             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)
+                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
+                    )
+                else:
+                    rot_center = (edges[0] + edges[1]) / 2
+                    rot_axes = self.get_custom_axes(
+                        a=rot_center, b=[edges[0]], c=np.zeros(3)
+                    )
 
             moment_of_inertia = self.get_custom_residue_moment_of_inertia(
                 center_of_mass=rot_center,
@@ -247,18 +286,35 @@ 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, 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 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.
+            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).
+            - 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
             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.
+            axes (see :meth:`get_bonded_axes`). Compute a custom MOI tensor.
 
         Args:
             data_container (MDAnalysis.Universe or AtomGroup):
@@ -293,66 +349,81 @@ 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_set = 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_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.atoms.select_atoms(
+                        f"(mass 2 to 999) and bonded index {edge_atom.index}"
                     )
-
-                    edges = [residue.atoms[0], second_edge[0]]
-                    backbone = self.get_chain(
-                        residue, residue.atoms[0], second_edge.atoms[0]
-                    )
-
-                elif res_position == 0:
+                    # find the average position of heavy atoms bonded to edge atom
+                    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)
+                    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.position, average_other_atoms],
+                            average_bonded_atom,
+                        )
+                    else:
+                        trans_center = edge_atom.position
+                        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]
                     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])
-
-                else:
-                    # last resid
-                    # always resindex 1 in data_container
-                    residue = data_container.residues[1]
-                    resindex = residue.resindex
-                    resindex_prev = resindex - 1
-                    first_edge = data_container.select_atoms(
-                        f"resindex {resindex} and bonded resindex {resindex_prev}"
-                    )
-
-                    last_index = len(heavy_atoms) - 1
-                    last = None
-                    # look for last heavy atom
-                    # with only one bond to another
-                    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")
+                    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
+                        )
+                    else:
+                        trans_center = (edges[0] + edges[1]) / 2
+                        trans_axes = self.get_custom_axes(
+                            a=trans_center, b=[edges[0]], c=np.zeros(3)
+                        )
 
             # look for heavy atoms in residue of interest
             heavy_atom_indices = []
@@ -580,9 +651,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 +669,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):

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.