diff --git a/atomium/structures/chains.py b/atomium/structures/chains.py index 9dc4cc5e..97f98e28 100644 --- a/atomium/structures/chains.py +++ b/atomium/structures/chains.py @@ -249,3 +249,19 @@ def ligand(self, ligand=None): if not isinstance(ligand, Molecule): raise TypeError("ligand {} is not a Molecule".format(ligand)) self._ligand = ligand + + + def residues(self, *args, **kwargs): + """Returns the :py:class:`.Residue` objects in the structure, including + water molecules. It can be given search criteria if you wish. + + :param str residue_id: Filter by residue ID. + :param str name: Filter by name. + :rtype: ``set``""" + + residues = ResidueStructure.residues(self, *args, **kwargs) + for atom in self.atoms(): + molecule = atom.molecule() + if molecule and not isinstance(molecule, Chain): + residues.add(molecule) + return residues diff --git a/atomium/structures/molecules.py b/atomium/structures/molecules.py index 1a345df8..a7bbc99d 100644 --- a/atomium/structures/molecules.py +++ b/atomium/structures/molecules.py @@ -340,7 +340,7 @@ def model(self): return atom.model() - def site(self): + def site(self, include_water=False): """Returns the :py:class:`.Site` that encompasses this molecule. This is all the residues with a non-hydrogen atom within 4 Angstroms of a non-hydrogen atom in the molecule. @@ -352,7 +352,12 @@ def site(self): for atom in atoms: nearby.update(atom.nearby(4, exclude="H")) residues = [atom.residue() for atom in nearby if atom not in atoms] - residues = [residue for residue in residues if residue] + if include_water: + residues += [ + atom.molecule() for atom in nearby if atom not in atoms + and atom.molecule() != None and atom.molecule().name() == "HOH" + ] + residues = set([residue for residue in residues if residue]) return Site(*residues, ligand=self) diff --git a/tests/integration/test_pdb.py b/tests/integration/test_pdb.py index cb887d16..b31220b2 100644 --- a/tests/integration/test_pdb.py +++ b/tests/integration/test_pdb.py @@ -86,6 +86,12 @@ def test_can_read_pdb(self): model.residue("A42"), model.residue("A70"), model.residue("A72"), model.residue("A96"), model.residue("A123"), model.residue("A155") ])) + full_site = model.molecule("A5001").site(include_water=True) + self.assertEqual(full_site.residues(), set([ + model.residue("A42"), model.residue("A70"), model.residue("A72"), + model.residue("A96"), model.residue("A123"), model.residue("A155"), + model.molecule("A3015") + ])) # Bonding is correct residue = chaina[0] diff --git a/tests/unit/structure_tests/test_molecules.py b/tests/unit/structure_tests/test_molecules.py index fda32ddb..55c98caf 100644 --- a/tests/unit/structure_tests/test_molecules.py +++ b/tests/unit/structure_tests/test_molecules.py @@ -151,34 +151,53 @@ def test_can_get_no_model(self, mock_atoms): class MoleculeSiteTests(MoleculeTest): + def setUp(self): + MoleculeTest.setUp(self) + self.molecule = Molecule(self.atom1, self.atom2, self.atom3) + + other_atoms = [Mock(), Mock(), Mock(), Mock(), Mock(), Mock(), Mock()] + self.atom1.nearby.return_value = set(other_atoms[:3] + self.atoms[1:]) + self.atom2.nearby.return_value = set(other_atoms[2:5] + self.atoms[::1]) + self.atom3.nearby.return_value = set(other_atoms[4:] + self.atoms[:2]) + + self.residues = [Mock(), Mock(), Mock(), Mock(), Mock(), Mock(), Mock()] + self.waters = [Mock(), Mock(), Mock(), Mock()] + self.waters[0].name.return_value = "HOH" + other_atoms[0].residue.return_value = self.residues[0] + other_atoms[1].residue.return_value = self.residues[0] + other_atoms[2].residue.return_value = self.residues[1] + other_atoms[3].residue.return_value = self.residues[1] + other_atoms[4].residue.return_value = self.residues[2] + other_atoms[5].residue.return_value = self.residues[2] + other_atoms[6].residue.return_value = None + other_atoms[6].molecule.return_value = self.waters[0] + + @patch("atomium.structures.Molecule.atoms") @patch("atomium.structures.chains.Site") def test_can_get_site(self, mock_site, mock_atoms): - residues = [Mock(), Mock(), Mock(), Mock(), Mock(), Mock(), Mock()] - atoms = [] - for residue in residues: - atoms += [Mock(), Mock(), Mock()] - atoms[-3].residue.return_value = residue - atoms[-2].residue.return_value = residue - atoms[-1].residue.return_value = residue mock_atoms.return_value = set(self.atoms) - self.atom1.nearby.return_value = set([self.atom2, atoms[1], atoms[3]]) - self.atom2.nearby.return_value = set([self.atom1, atoms[3], atoms[6]]) - self.atom3.nearby.return_value = set([self.atom3, atoms[6], atoms[-1]]) - self.atom1.residue.return_value = None - self.atom2.residue.return_value = None - self.atom3.residue.return_value = None - site = Mock() - mock_site.return_value = site - mol = Molecule(self.atom1, self.atom2, self.atom3) - returned_site = mol.site() - self.assertIs(site, returned_site) + returned_site = self.molecule.site() + mock_atoms.assert_called_with(exclude="H") + self.atom1.nearby.assert_called_with(4, exclude="H") + self.atom2.nearby.assert_called_with(4, exclude="H") + self.atom3.nearby.assert_called_with(4, exclude="H") + residues_passed = mock_site.call_args_list[0][0] + self.assertEqual(set(residues_passed), set(self.residues[:3])) + kwargs = mock_site.call_args_list[0][1] + self.assertEqual(kwargs, {"ligand": self.molecule}) + + + @patch("atomium.structures.Molecule.atoms") + @patch("atomium.structures.chains.Site") + def test_can_get_site_with_water(self, mock_site, mock_atoms): + mock_atoms.return_value = set(self.atoms) + returned_site = self.molecule.site(include_water=True) mock_atoms.assert_called_with(exclude="H") self.atom1.nearby.assert_called_with(4, exclude="H") self.atom2.nearby.assert_called_with(4, exclude="H") self.atom3.nearby.assert_called_with(4, exclude="H") - site_args, site_kwargs = mock_site.call_args_list[0] - self.assertEqual(set(site_args), set([ - residues[0], residues[1], residues[2], residues[6] - ])) - self.assertEqual(site_kwargs, {"ligand": mol}) + residues_passed = mock_site.call_args_list[0][0] + self.assertEqual(set(residues_passed), set(self.residues[:3] + self.waters[:1])) + kwargs = mock_site.call_args_list[0][1] + self.assertEqual(kwargs, {"ligand": self.molecule}) diff --git a/tests/unit/structure_tests/test_sites.py b/tests/unit/structure_tests/test_sites.py index 984a6db3..89fd0cac 100644 --- a/tests/unit/structure_tests/test_sites.py +++ b/tests/unit/structure_tests/test_sites.py @@ -79,3 +79,28 @@ def test_ligand_must_be_molecule(self): site = Site(*self.residues, ligand=self.ligand) with self.assertRaises(TypeError): site.ligand("ligand") + + + +class SiteResiduesTests(SiteTest): + + @patch("atomium.structures.chains.ResidueStructure.residues") + def test_can_get_normal_residues(self, mock_res): + mock_res.return_value = [1, 2, 3] + site = Site(*self.residues) + self.assertEqual(site.residues(), [1, 2, 3]) + + + @patch("atomium.structures.chains.ResidueStructure.residues") + @patch("atomium.structures.chains.Site.atoms") + def test_can_get_water_residues(self, mock_atoms, mock_res): + mock_res.return_value = set([1, 2, 3]) + site = Site(*self.residues) + atoms = [Mock(), Mock(), Mock()] + site.atoms.return_value = atoms + water = Mock() + water.residue_name.return_value = "HOH" + atoms[0].molecule.return_value = water + atoms[1].molecule.return_value = 1 + atoms[2].molecule.return_value = 2 + self.assertEqual(site.residues(), set([1, 2, 3, water]))