Implementation of WaterSelection for water resiude selection (MDAnalysis#4854)

* Update selection.py

Addition of WaterSelection

* Update test_atomselections.py

added tests

* Update CHANGELOG

changelog addition mentioned

* Update selection.py

Author: talagayev

package/CHANGELOG

@@ -27,6 +27,7 @@ Fixes
    the function to prevent shared state. (Issue #4655)
 
 Enhancements
+ * Addition of 'water' token for water selection (Issue #4839)
  * Enables parallelization for analysis.density.DensityAnalysis (Issue #4677, PR #4729)
  * Enables parallelization for analysis.contacts.Contacts (Issue #4660)
  * Enable parallelization for analysis.nucleicacids.NucPairDist (Issue #4670)

package/MDAnalysis/core/selection.py

@@ -1094,6 +1094,37 @@ def _apply(self, group):
         return group[mask]
 
 
+class WaterSelection(Selection):
+    """All atoms in water residues with recognized water residue names.
+
+    Recognized residue names:
+
+    * recognized 3 Letter resnames: 'H2O', 'HOH', 'OH2', 'HHO', 'OHH'
+        'TIP', 'T3P', 'T4P', 'T5P', 'SOL', 'WAT'
+
+    * recognized 4 Letter resnames: 'TIP2', 'TIP3', 'TIP4'
+
+    .. versionadded:: 2.9.0
+    """
+    token = 'water'
+
+    # Recognized water resnames
+    water_res = {
+        'H2O', 'HOH', 'OH2', 'HHO', 'OHH',
+        'T3P', 'T4P', 'T5P', 'SOL', 'WAT',
+        'TIP', 'TIP2', 'TIP3', 'TIP4'
+    }
+
+    def _apply(self, group):
+        resnames = group.universe._topology.resnames
+        nmidx = resnames.nmidx[group.resindices]
+
+        matches = [ix for (nm, ix) in resnames.namedict.items()
+                   if nm in self.water_res]
+        mask = np.isin(nmidx, matches)
+
+        return group[mask]
+
 class BackboneSelection(ProteinSelection):
     """A BackboneSelection contains all atoms with name 'N', 'CA', 'C', 'O'.
 

testsuite/MDAnalysisTests/core/test_atomselections.py

@@ -48,6 +48,8 @@
     PDB_helix,
     PDB_elements,
     PDB_charges,
+    waterPSF,
+    PDB_full,
 )
 from MDAnalysisTests import make_Universe
 
@@ -650,6 +652,38 @@ def test_nucleicsugar(self, universe):
         assert_equal(rna.n_atoms, rna.n_residues * 5)
 
 
+class TestSelectionsWater(object):
+    @pytest.fixture(scope='class')
+    def universe(self):
+        return MDAnalysis.Universe(GRO)
+
+    @pytest.fixture(scope='class')
+    def universe2(self):
+        return MDAnalysis.Universe(waterPSF)
+
+    @pytest.fixture(scope='class')
+    def universe3(self):
+        return MDAnalysis.Universe(PDB_full)
+
+    def test_water_gro(self, universe):
+        # Test SOL water with 4 atoms
+        water_gro = universe.select_atoms("water")
+        assert_equal(water_gro.n_atoms, 44336)
+        assert_equal(water_gro.n_residues, 11084)
+
+    def test_water_tip3(self, universe2):
+        # Test TIP3 water with 3 atoms
+        water_tip3 = universe2.select_atoms('water')
+        assert_equal(water_tip3.n_atoms, 15)
+        assert_equal(water_tip3.n_residues, 5)
+
+    def test_water_pdb(self, universe3):
+        # Test HOH water with 1 atom
+        water_pdb = universe3.select_atoms("water")
+        assert_equal(water_pdb.n_residues, 188)
+        assert_equal(water_pdb.n_atoms, 188)
+
+
 class BaseDistanceSelection(object):
     """Both KDTree and distmat selections on orthogonal system