Add option to include water residues in binding sites

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

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)
 
 

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]

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})

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]))