diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index 7d93fd9..6497337 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -73,17 +73,35 @@ 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). 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. Compute a + 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. + 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). 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. Compute a custom MOI, using heavy + atom positions and heavy atom + hydrogen masses. Args: data_container (MDAnalysis.Universe or AtomGroup): @@ -145,30 +163,51 @@ def get_residue_axes( 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] + edge_atom = edge_atom_set[0] + 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 + 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) + 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 ) 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) + rot_center = edge_atom.position + rot_axes = self.get_custom_axes( + a=edge_atom.position, b=[average_bonded_atom], c=np.zeros(3) + ) + else: - edges = [edge_atom_set[0], edge_atom_set[1]] + 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) - rot_center, rot_axes = self.get_residue_custom_axes(edges, backbone_center) + 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, backbone_center + ) + else: + rot_center = (edges[0] + edges[1]) / 2 + rot_axes = self.get_custom_axes( + a=rot_center, b=[edges[0]], c=np.zeros(3) + ) moment_of_inertia = self.get_custom_residue_moment_of_inertia( center_of_mass=rot_center, @@ -247,18 +286,35 @@ 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 and 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 bonded to only one other residue, find edge heavy atom + (i.e. heavy atom bonded to neighbour residue). 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 bonded to at least two other residues, 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). + - If bonded to at least two other residues and 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. - 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): @@ -293,66 +349,81 @@ def get_UA_axes(self, data_container, index: int, res_position): 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}" + 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] + 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_atom.index}" ) - - edges = [residue.atoms[0], second_edge[0]] - backbone = self.get_chain( - residue, residue.atoms[0], second_edge.atoms[0] - ) - - elif res_position == 0: + # find the average position of heavy atoms bonded to edge atom + 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) + 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.position, average_other_atoms], + average_bonded_atom, + ) + else: + trans_center = edge_atom.position + trans_axes = self.get_custom_axes( + a=edge_atom.position, b=[average_bonded_atom], c=np.zeros(3) + ) + else: # between 2 residues residue = data_container.residues[1] resindex = residue.resindex resindex_next = resindex + 1 resindex_prev = resindex - 1 - + residue_heavy_atoms = residue.atoms.select_atoms("mass 2 to 999") edge_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]] + edges = edge_set.positions 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}" - ) - - 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") + 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 + ) + else: + trans_center = (edges[0] + edges[1]) / 2 + trans_axes = self.get_custom_axes( + a=trans_center, b=[edges[0]], c=np.zeros(3) + ) # look for heavy atoms in residue of interest heavy_atom_indices = [] @@ -580,9 +651,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 +669,7 @@ 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_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..2e4b52f 100644 --- a/tests/unit/CodeEntropy/levels/test_axes.py +++ b/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,) @@ -1269,7 +1269,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 +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", @@ -1701,3 +1639,181 @@ def _sel(q): 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 + 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]) + monkeypatch.setattr("CodeEntropy.levels.axes.make_whole", lambda _ag: None) + residue = u.select_atoms("resindex 0") + 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=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 + 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])