diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index 7d93fd9..a6676bd 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -73,17 +73,23 @@ 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). + 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) - 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. + * 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. Args: data_container (MDAnalysis.Universe or AtomGroup): @@ -143,33 +149,13 @@ def get_residue_axes( else: make_whole(data_container.atoms) 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] - ) - 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) + edge_atom = edge_atom_set[0] + rot_center, rot_axes = self.get_terminal_axes(residue, edge_atom) else: - edges = [edge_atom_set[0], edge_atom_set[1]] - 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) - + rot_center, rot_axes = self.get_non_terminal_axes( + residue, edge_atom_set + ) moment_of_inertia = self.get_custom_residue_moment_of_inertia( center_of_mass=rot_center, positions=uas.positions, @@ -247,18 +233,17 @@ 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, 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 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): @@ -290,71 +275,47 @@ def get_UA_axes(self, data_container, index: int, res_position): residue = data_container trans_center = data_container.atoms.center_of_mass(unwrap=True) trans_axes = data_container.atoms.principal_axes() - 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}" - ) - - edges = [residue.atoms[0], second_edge[0]] - backbone = self.get_chain( - residue, residue.atoms[0], second_edge.atoms[0] + 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] + trans_center, trans_axes = self.get_terminal_axes( + residue, edge_atom ) - - elif res_position == 0: + else: # between 2 residues residue = data_container.residues[1] resindex = residue.resindex resindex_next = resindex + 1 resindex_prev = resindex - 1 - - edge_set = data_container.select_atoms( + residue_heavy_atoms = residue.atoms.select_atoms("mass 2 to 999") + edge_atom_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]] - 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}" + trans_center, trans_axes = self.get_non_terminal_axes( + residue, edge_atom_set ) - - 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") - # look for heavy atoms in residue of interest + 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) @@ -580,9 +541,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 +559,99 @@ 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_terminal_axes(self, residue, edge): + """ + Compute rotation axes at the residue level/translation axes at the UA level + for the terminal residues in a polymer, given the edge atom + (i.e. atom bonded to neighbour residue) and residue of interest. + 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. + If there no heavy atoms bonded to the edge atom (i.e. the edge atom is the only + heavy atom in the residue), centre is set on edge atom and axes are principal + axes. + + Args: + residue: MDAnalysis AtomGroup + edge: MDAnalysis atom + + Returns: + rot_center: (3,) rotation centre, + rot_axes: (3,3) rotation axes of 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.index}" + ) + if len(bonded_atoms) == 0: + # there is only one heavy atom in the residue + rot_center = edge.position + rot_axes = residue.atoms.principal_axes() + else: + average_bonded = np.zeros(3) + for bonded_atom in bonded_atoms: + average_bonded += bonded_atom.position + average_bonded /= len(bonded_atoms) + # find the average position of all other heavy atoms in residue + other_atoms = [] + for atom in heavy_atoms: + if atom != edge 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(other_atoms) + rot_center, rot_axes = self.get_residue_custom_axes( + [edge.position, average_other_atoms], average_bonded + ) + else: + rot_center = edge.position + rot_axes = self.get_custom_axes( + a=edge.position, b=[average_bonded], c=np.zeros(3) + ) + return rot_center, rot_axes + + def get_non_terminal_axes(self, residue, edges): + """ + Compute rotation axes at the residue level/ translation axes at + the UA level for the non-terminal residues in a linear polymer, given the + edge atoms (i.e. heavy atoms bonded to neighbour residues) and + residue of interest. Find the shortest chain between edge atoms: the backbone. + Edges + backbone average position determine the 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. + Args: + residue: MDAnalysis AtomGroup + edges: MDAnalysis AtomGroup + + Returns: + rot_center: (3,) rotation centre, + rot_axes: (3,3) rotation axes of residue + """ + backbone = self.get_chain(residue, edges[0], edges[1]) + 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.positions, backbone_center + ) + else: + rot_center = (edges[0].position + edges[1].position) / 2 + rot_axes = self.get_custom_axes(a=rot_center, b=[edges[0]], c=np.zeros(3)) return rot_center, rot_axes def get_bonded_axes(self, system, atom, dimensions: np.ndarray): diff --git a/docs/science.rst b/docs/science.rst index 7aa219d..b42af0a 100644 --- a/docs/science.rst +++ b/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. diff --git a/tests/unit/CodeEntropy/levels/test_axes.py b/tests/unit/CodeEntropy/levels/test_axes.py index c926f95..e54a850 100644 --- a/tests/unit/CodeEntropy/levels/test_axes.py +++ b/tests/unit/CodeEntropy/levels/test_axes.py @@ -218,6 +218,7 @@ def _sel(q): return [] u.select_atoms.side_effect = _sel + monkeypatch.setattr("CodeEntropy.levels.axes.make_whole", lambda _ag: None) monkeypatch.setattr(ax, "get_bonded_axes", lambda **kwargs: (None, None)) @@ -1220,7 +1221,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,) @@ -1269,7 +1270,7 @@ def _select_atoms(q): u.atoms.select_atoms.side_effect = _select_atoms u.atoms.principal_axes.return_value = np.eye(3) - monkeypatch.setattr(ax, "get_chain", backbone_atom) + monkeypatch.setattr(ax, "get_chain", lambda residue, first, last: [backbone_atom]) monkeypatch.setattr( ax, "get_custom_residue_moment_of_inertia", @@ -1337,90 +1338,34 @@ 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 +1375,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 +1417,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 +1506,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 +1542,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 +1551,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", @@ -1684,20 +1624,419 @@ def test_get_UA_axes_raises_when_only_rot_axes_fail(monkeypatch): u = MagicMock() u.atoms.principal_axes.return_value = np.eye(3) u.dimensions = np.array([10.0, 10.0, 10.0, 90, 90, 90]) - residue = MagicMock() - u.residues = [residue] heavy_atoms = [ _atom(index=0, mass=12.0, pos=(1, 0, 0)), _atom(index=1, mass=12.0, pos=(0, 1, 0)), ] + u.residues = [heavy_atoms] def _sel(q): if q == "mass 2 to 999": return heavy_atoms + if q.startswith("index"): + return [heavy_atoms[0]] u.select_atoms.side_effect = _sel + u.atoms.select_atoms.side_effect = _sel + monkeypatch.setattr("CodeEntropy.levels.axes.make_whole", lambda _ag: None) monkeypatch.setattr(ax, "get_bonded_axes", lambda **kwargs: (None, None)) with pytest.raises(ValueError): ax.get_UA_axes(u, index=0, res_position=None) + + +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 + 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]) + monkeypatch.setattr("CodeEntropy.levels.axes.make_whole", lambda _ag: None) + residue = u.select_atoms("resindex 0") + residue.__len__.return_value = 2 + 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, 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("resindex 0 and (bonded resindex"): + # the edge atom + return [uas[2]] + if q.startswith("(mass 2 to 999) and bonded index "): + return uas[0:2] + + 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_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=0, + 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, [0, 0, 1]) + assert np.allclose(moi, [1, 1, 1]) + + +def test_get_res_axes_terminal_1_atom(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 = 1 + uas = _FakeAtomGroup( + [ + _atom(index=0, mass=12.0, pos=(1, 0, 0)), + ], + ) + u.atoms.principal_axes.return_value = np.eye(3) + residue.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 index "): + return [] + if q.startswith("resindex "): + return [uas[0]] + + 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_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=0, + relative_index=0, + ) + assert np.allclose(trans_axes, np.eye(3)) + assert np.allclose(rot_axes, np.eye(3)) + assert np.allclose(rot_center, [1, 0, 0]) + assert np.allclose(moi, [1, 1, 1]) + + +def test_get_UA_axes_terminal_1_atom(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, 1, 1)), + _atom(index=1, mass=12.0, pos=(3, 3, 3)), + ], + ) + + def _select_atoms_data_container(q): + if q == "mass 2 to 999": + # return heavy atoms group + return heavy_atoms + if q.startswith("resindex "): + return [heavy_atoms[0]] + + def _select_atoms_residue(q): + if q == "mass 2 to 999": + # return heavy atoms group + return [heavy_atoms[0]] + if q.startswith("index"): + return [heavy_atoms[0]] + if q.startswith("(mass 2 to 999) and bonded "): + return [] + + residue_group.select_atoms.side_effect = _select_atoms_data_container + residue_group.atoms.select_atoms.side_effect = _select_atoms_data_container + residue.select_atoms.side_effect = _select_atoms_residue + residue.atoms.select_atoms.side_effect = _select_atoms_residue + monkeypatch.setattr( + ax, + "get_bonded_axes", + lambda system, atom, dimensions: (np.eye(3), 3 * np.eye(3)), + ) + residue.atoms.principal_axes.return_value = 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, 1, 1]) + assert np.allclose(moi, 3 * np.eye(3)) + + +def test_get_terminal_axes_1point(monkeypatch): + ax = AxesCalculator() + residue = MagicMock() + residue.__len__ = 1 + heavy_atoms = _FakeAtomGroup( + [ + _atom(index=0, mass=12.0, pos=(1, 0, 0)), + ], + ) + + def _select_atoms(q): + if q == "mass 2 to 999": + return heavy_atoms + if q.startswith("(mass 2 to 999) and bonded"): + return [] + + residue.atoms.select_atoms.side_effect = _select_atoms + residue.atoms.principal_axes.return_value = np.eye(3) + centre, axes = ax.get_terminal_axes(residue, heavy_atoms[0]) + assert np.allclose(centre, [1, 0, 0]) + assert np.allclose(axes, np.eye(3)) + + +def test_get_terminal_axes_2points(monkeypatch): + ax = AxesCalculator() + residue = MagicMock() + residue.__len__ = 2 + heavy_atoms = _FakeAtomGroup( + [ + _atom(index=0, mass=12.0, pos=(1, 0, 0)), + _atom(index=1, mass=12.0, pos=(0, 1, 0)), + ], + ) + + def _select_atoms(q): + if q == "mass 2 to 999": + return heavy_atoms + if q.startswith("(mass 2 to 999) and bonded"): + return [heavy_atoms[1]] + + residue.atoms.select_atoms.side_effect = _select_atoms + monkeypatch.setattr(ax, "get_custom_axes", lambda a, b, c: 2 * np.eye(3)) + centre, axes = ax.get_terminal_axes(residue, heavy_atoms[0]) + assert np.allclose(centre, [1, 0, 0]) + assert np.allclose(axes, 2 * np.eye(3)) + + +def test_get_terminal_axes_3points(monkeypatch): + ax = AxesCalculator() + residue = MagicMock() + residue.__len__ = 4 + heavy_atoms = _FakeAtomGroup( + [ + _atom(index=0, mass=12.0, pos=(0, 0, 0)), + _atom(index=1, mass=12.0, pos=(1, 1, 0)), + _atom(index=2, mass=12.0, pos=(1, 0, 0)), + _atom(index=3, mass=14.0, pos=(3, 0, 0)), + ], + ) + + def _select_atoms(q): + if q == "mass 2 to 999": + return heavy_atoms + if q.startswith("(mass 2 to 999) and bonded"): + return [heavy_atoms[1]] + + residue.atoms.select_atoms.side_effect = _select_atoms + monkeypatch.setattr( + ax, "get_residue_custom_axes", lambda edges, center: ([1, 0, 0], 3 * np.eye(3)) + ) + + centre, axes = ax.get_terminal_axes(residue, heavy_atoms[0]) + + assert np.allclose(centre, [1, 0, 0]) + assert np.allclose(axes, 3 * np.eye(3)) + + +def test_get_non_terminal_axes_2points(monkeypatch): + ax = AxesCalculator() + residue = MagicMock() + residue.__len__ = 2 + heavy_atoms = _FakeAtomGroup( + [ + _atom(index=0, mass=12.0, pos=(2, 0, 0)), + _atom(index=1, mass=12.0, pos=(0, 2, 0)), + ], + ) + monkeypatch.setattr(ax, "get_chain", lambda residue, first, last: []) + monkeypatch.setattr(ax, "get_custom_axes", lambda a, b, c: 2 * np.eye(3)) + centre, axes = ax.get_non_terminal_axes(residue, heavy_atoms) + assert np.allclose(centre, [1, 1, 0]) + assert np.allclose(axes, 2 * np.eye(3)) + + +def test_get_non_terminal_axes_3points(monkeypatch): + ax = AxesCalculator() + residue = MagicMock() + residue.__len__ = 4 + heavy_atoms = _FakeAtomGroup( + [ + _atom(index=0, mass=12.0, pos=(0, 0, 0)), + _atom(index=1, mass=12.0, pos=(1, 0, 0)), + _atom(index=2, mass=14.0, pos=(3, 0, 0)), + _atom(index=3, mass=12.0, pos=(1, 1, 0)), + ], + ) + edges = _FakeAtomGroup( + [ + _atom(index=0, mass=12.0, pos=(0, 0, 0)), + _atom(index=3, mass=12.0, pos=(1, 1, 0)), + ], + ) + monkeypatch.setattr( + ax, "get_chain", lambda residue, first, last: [heavy_atoms[1], heavy_atoms[2]] + ) + monkeypatch.setattr( + ax, "get_residue_custom_axes", lambda edges, center: ([1, 0, 0], 3 * np.eye(3)) + ) + centre, axes = ax.get_non_terminal_axes(residue, edges) + + assert np.allclose(centre, [1, 0, 0]) + assert np.allclose(axes, 3 * np.eye(3))