unit tests for edge cases

Author: ioanaapapa

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,38 +1782,38 @@ 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",
         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
+        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])