update unit tests + fix atom selection bug

Author: ioanaapapa

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,37 +329,40 @@ 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:
                         # last residue
                         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

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",