updated tests and removed redundant 1 heavy atom code

Author: ioanaapapa

CodeEntropy/levels/axes.py

@@ -123,6 +123,7 @@ def get_residue_axes(self, data_container, index: int, residue=None):
             # No UAS are bonded to other residues
             # Use a custom principal axes, from a MOI tensor that uses positions of
             # heavy atoms only, but including masses of heavy atom + bonded H.
+
             moi_tensor = self.get_moment_of_inertia_tensor(
                 center_of_mass=np.array(residue.center_of_mass()),
                 positions=uas.positions,
@@ -225,7 +226,6 @@ def get_UA_axes(self, data_container, index: int, res_position):
         """
         index = int(index)  # bead index
         heavy_atoms = data_container.atoms.select_atoms("mass 2 to 999")
-
         # use the same customPI trans axes as the residue level
         if len(heavy_atoms) > 1:
             if len(data_container.residues) == 1:
@@ -289,15 +289,6 @@ def get_UA_axes(self, data_container, index: int, res_position):
                 trans_center = np.array(backbone.center_of_mass())
                 trans_axes = self.get_residue_custom_axes(edges, trans_center)
 
-        else:
-            # only one heavy atom or hydrogen molecule
-            make_whole(data_container.atoms)
-            residue = data_container
-            # trans_center is center of mass
-            trans_center = np.array(data_container.center_of_mass())
-            trans_axes = data_container.atoms.principal_axes()
-
-        if len(heavy_atoms) > 1:
             residue_heavy_atoms = residue.atoms.select_atoms("mass 2 to 999")
             # look for heavy atoms in residue of interest
             heavy_atom_indices = []
@@ -307,19 +298,30 @@ def get_UA_axes(self, data_container, index: int, res_position):
             # where n is the bead index
             heavy_atom_index = heavy_atom_indices[index]
             heavy_atom = residue.atoms.select_atoms(f"index {heavy_atom_index}")
+            rot_center = heavy_atom.positions[0]
+            rot_axes, moment_of_inertia = self.get_bonded_axes(
+                system=data_container,
+                atom=heavy_atom[0],
+                dimensions=data_container.dimensions[:3],
+            )
+
         else:
-            # only the one heavy atom
+            # 1 heavy atom in the data_container
             heavy_atom = heavy_atoms[0]
+            # trans and rot centres are centre of mass
+            rot_center = data_container.center_of_mass()
+            rot_axes, moment_of_inertia = self.get_bonded_axes(
+                system=data_container,
+                atom=heavy_atom[0],
+                dimensions=data_container.dimensions[:3],
+            )
+            trans_center = rot_center
+            # principal axes
+            trans_axes = rot_axes
+
         if trans_axes is None:
             raise ValueError("Unable to compute translation axes for UA bead.")
 
-        rot_center = heavy_atom.positions[0]
-        rot_axes, moment_of_inertia = self.get_bonded_axes(
-            system=data_container,
-            atom=heavy_atom[0],
-            dimensions=data_container.dimensions[:3],
-        )
-
         if rot_axes is None or moment_of_inertia is None:
             raise ValueError("Unable to compute bonded axes for UA bead.")
 
@@ -765,7 +767,6 @@ def get_moment_of_inertia_tensor(
         """
         r = self.get_vector(center_of_mass, positions, dimensions)
         r2 = np.sum(r**2, axis=1)
-
         masses_arr = np.asarray(list(masses), dtype=float)
         moment_of_inertia_tensor = np.eye(3) * np.sum(masses_arr * r2)
         moment_of_inertia_tensor -= np.einsum("i,ij,ik->jk", masses_arr, r, r)
@@ -797,6 +798,7 @@ def get_custom_principal_axes(
                 - principal_axes: (3, 3) principal axes (rows).
                 - moment_of_inertia: (3,) principal moments.
         """
+
         eigenvalues, eigenvectors = np.linalg.eig(moment_of_inertia_tensor)
         order = np.abs(eigenvalues).argsort()[::-1]  # descending order
         transposed = np.transpose(eigenvectors)  # columns -> rows

tests/unit/CodeEntropy/levels/test_axes.py

@@ -166,13 +166,15 @@ def test_get_UA_axes_uses_principal_axes_when_single_heavy(monkeypatch):
     u = MagicMock()
     u.dimensions = np.array([10.0, 10.0, 10.0, 90, 90, 90])
     u.atoms.principal_axes.return_value = np.eye(3)
+    u.center_of_mass.return_value = np.array([[4.0, 0.0, 0.0]])
 
     # heavy_atoms length <= 1 => principal_axes path
-    heavy_atom = MagicMock(index=5)
+    heavy_atom = MagicMock(index=5, position=np.array([4.0, 0.0, 0.0]))
     heavy_atoms = [heavy_atom]
 
     def _sel(q):
-        if q == "prop mass > 1.1":
+        if q == "mass 2 to 999":
+            # return heavy atoms group
             return heavy_atoms
         if q.startswith("index "):
             # return atom group with positions
@@ -181,6 +183,7 @@ def _sel(q):
             ag.__getitem__.return_value = MagicMock(
                 mass=12.0, position=np.array([4.0, 0.0, 0.0]), index=5
             )
+
             return ag
         return []
 
@@ -513,18 +516,16 @@ def _select_atoms(q):
     assert np.allclose(moi, np.array([3.0, 2.0, 1.0]))
 
 
-def test_get_residue_axes_with_bonds_vanilla_path(monkeypatch):
+def test_get_residue_axes_vanilla_path(monkeypatch):
     ax = AxesCalculator()
-
     residue = MagicMock()
     residue.__len__.return_value = 1
-    residue.atoms.principal_axes.return_value = np.eye(3) * 2
     residue.atoms.center_of_mass.return_value = np.array([1.0, 2.0, 3.0])
     residue.center_of_mass.return_value = np.array([1.0, 2.0, 3.0])
+    residue.select_atoms.return_value = MagicMock(positions=np.zeros((1, 3)))
 
     u = MagicMock()
     u.dimensions = np.array([10.0, 10.0, 10.0, 90, 90, 90])
-    u.atoms.principal_axes.return_value = np.eye(3) * 2
 
     def _select_atoms(q):
         if q.startswith("(resindex"):
@@ -537,7 +538,9 @@ def _select_atoms(q):
 
     monkeypatch.setattr("CodeEntropy.levels.axes.make_whole", lambda _ag: None)
     monkeypatch.setattr(
-        ax, "get_vanilla_axes", lambda mol: (np.eye(3) * 2, np.array([9.0, 8.0, 7.0]))
+        ax,
+        "get_custom_principal_axes",
+        lambda mol: (np.eye(3) * 2, np.array([9.0, 8.0, 7.0])),
     )
 
     trans, rot, center, moi = ax.get_residue_axes(u, index=10, residue=residue)
@@ -647,19 +650,13 @@ def center_of_mass(self, *args, **kwargs):
         def __getitem__(self, idx):
             return system_atom
 
-        def _select_atoms(q):
-            if q == "prop mass > 1.1":
-                return heavy_atoms
-            if q.startswith("index "):
-                return heavy_atom_selection
-            return _FakeAtomGroup([])
-
     data_container = MagicMock()
     data_container.atoms = _Atoms()
     data_container.dimensions = np.array([10.0, 10.0, 10.0, 90, 90, 90], dtype=float)
+    _FakeAtomGroup.atoms = heavy_atom_selection
 
     def _select_atoms(q):
-        if q == "prop mass > 1.1":
+        if q == "mass 2 to 999":
             return heavy_atoms
         if q.startswith("index "):
             return heavy_atom_selection