fixed calling get_custom_axes in get_non_terminal_axes and get_terminal_axes

Author: ioanaapapa

CodeEntropy/levels/axes.py

@@ -73,13 +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 includes 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:
+
+        - 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).
+              Compute rotation centre and axes as in get_terminal_axes.
+              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).
+              Compute rotation centre and axes as in get_non_terminal_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).
               Compute rotation centre and axes as in get_terminal_axes.
               custom MOI, using heavy atom positions and heavy atom + hydrogen masses.
@@ -244,12 +257,18 @@ def get_UA_axes(self, data_container, index: int, res_position):
             only, but includes masses of heavy atom + bonded hydrogens.
             - If bonded to only one other residue, see get_terminal_axes.
             - If bonded to at least two other residues, see get_non_terminal_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, see get_terminal_axes.
+            - If bonded to at least two other residues, see get_non_terminal_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.
+            axes (see :meth:`get_bonded_axes`). Compute a custom MOI tensor.
 
         Args:
             data_container (MDAnalysis.Universe or AtomGroup):
@@ -326,6 +345,7 @@ def get_UA_axes(self, data_container, index: int, res_position):
                     )
             # look for heavy atoms in residue of interest
             residue_heavy_atoms = residue.atoms.select_atoms("mass 2 to 999")
+            residue_heavy_atoms = residue.atoms.select_atoms("mass 2 to 999")
             heavy_atom_indices = []
             for atom in residue_heavy_atoms:
                 heavy_atom_indices.append(atom.index)