Added the 'same ... as' selection keyword (Fixes #217)

package/AUTHORS

@@ -54,6 +54,7 @@ Chronological list of authors
   - Joe Jordan
 2015
   - Alex Nesterenko
+  - Caio S. Souza
 
 External code
 -------------

package/CHANGELOG

@@ -10,11 +10,12 @@ The rules for this file:
     use tabs but use spaces for formatting
 
 ------------------------------------------------------------------------------
-??/??/??  richardjgowers
+??/??/??  richardjgowers, caio s. souza, manuel.nuno.melo
   * 0.10.0
 
   Enhancements
 
+  * Added the 'same ... as' selection keyword (Issue #217)
   * Added guess_bonds keyword argument to Universe creation.  This will attempt to
     guess all topology information on Universe creation. (Issue #245)
   * Added guess_bonds method to AtomGroup. (Issue #245)

package/MDAnalysis/core/Selection.py

@@ -815,6 +815,32 @@ def __repr__(self):
         return "<'PropertySelection' " + abs_str + repr(self.prop) + " " + repr(self.operator.__name__) + " " + repr(
             self.value) + ">"
 
+class SameSelection(Selection):
+    def __init__(self, sel, prop):
+        Selection.__init__(self)
+        self.sel = sel
+        self.prop = prop
+    def _apply(self, group):
+        res = self.sel._apply(group)
+        if not res:
+            return set([])
+        if self.prop in ("residue", "fragment"):
+            atoms = set([])
+            for a in res:
+                atoms |= set(getattr(a, self.prop).atoms)
+            return Selection._group_atoms & atoms
+        elif self.prop in ("name", "type", "resname", "resid", "segid", "mass", "charge", "radius", "bfactor", "resnum"):
+            props = [getattr(a, self.prop) for a in res]
+            result_set = (a for a in Selection._group_atoms if getattr(a, self.prop) in props)
+        elif self.prop in ("x", "y", "z"):
+            p = getattr(Selection.coord, "_"+self.prop)
+            res_indices = numpy.array([a.number for a in res])
+            sel_indices = numpy.array([a.number for a in Selection._group_atoms])
+            result_set = group.atoms[numpy.where(numpy.in1d(p[sel_indices], p[res_indices]))[0]]._atoms
+        return set(result_set)
+    def __repr__(self):
+        return "<'SameSelection' of "+ repr(self.prop)+" >"
+
 
 class ParseError(Exception):
     pass
@@ -872,6 +898,7 @@ class SelectionParser:
     NBB = 'nucleicbackbone'
     BASE = 'nucleicbase'
     SUGAR = 'nucleicsugar'
+    SAME = 'same'
     EOF = 'EOF'
     GT = '>'
     LT = '<'
@@ -894,11 +921,12 @@ class SelectionParser:
         (BB, BackboneSelection), (NBB, NucleicBackboneSelection),
         (BASE, BaseSelection), (SUGAR, NucleicSugarSelection),
         #(BONDED, BondedSelection), not supported yet, need a better way to walk the bond lists
-        (ATOM, AtomSelection), (SELGROUP, SelgroupSelection), (FULLSELGROUP, FullSelgroupSelection)])
+        (ATOM, AtomSelection), (SELGROUP, SelgroupSelection), (FULLSELGROUP, FullSelgroupSelection),
+        (SAME, SameSelection)])
     associativity = dict([(AND, "left"), (OR, "left")])
     precedence = dict(
         [(AROUND, 1), (SPHLAYER, 1), (SPHZONE, 1), (CYLAYER, 1), (CYZONE, 1), (POINT, 1), (BYRES, 1), (BONDED, 1),
-            (AND, 3), (OR, 3), (NOT, 5)])
+            (SAME, 1), (AND, 3), (OR, 3), (NOT, 5)])
 
     # Borg pattern: http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/66531
     _shared_state = {}
@@ -1049,6 +1077,15 @@ def __parse_subexp(self):
             resid = int(self.__consume_token())
             name = self.__consume_token()
             return self.classdict[op](name, resid, segid)
+        elif op == self.SAME:
+            prop = self.__consume_token()
+            self.__expect("as")
+            if prop in ("name", "type", "resname", "resid", "segid", "mass", "charge", "radius", "bfactor",
+                        "resnum", "residue", "segment", "fragment", "x", "y", "z"):
+                exp = self.__parse_expression(self.precedence[op])
+                return self.classdict[op](exp, prop)
+            else:
+                self.__error(prop, expected=False)
         else:
             self.__error(op)
 

package/doc/sphinx/source/conf.py

@@ -90,8 +90,8 @@ def __getattr__(cls, name):
     David Caplan, Matthieu Chavent, Xavier Deupi, Jan Domański, David L. Dotson
     Lennard van der Feltz, Philip Fowler, Joseph Goose, Richard J. Gowers, Lukas Grossar,
     Benjamin Hall, Joe Jordan, Jinju Lu, Robert McGibbon, Alex Nesterenko,
-    Manuel Nuno Melo, Danny Parton, Joshua L. Phillips, Tyler Reddy, Paul Rigor, Andy Somogyi,
-    Lukas Stelzl, Zhuyi Xue, and Oliver Beckstein"""
+    Manuel Nuno Melo, Caio S. Souza, Danny Parton, Joshua L. Phillips, Tyler Reddy,
+    Paul Rigor, Andy Somogyi, Lukas Stelzl, Zhuyi Xue, and Oliver Beckstein"""
 project = u'MDAnalysis'
 copyright = u'2005-2015, ' + authors
 

testsuite/MDAnalysisTests/test_selections.py

@@ -15,7 +15,7 @@
 #
 import MDAnalysis
 import MDAnalysis.core.Selection
-from MDAnalysis.tests.datafiles import PSF, DCD, PRMpbc, TRJpbc_bz2, PSF_NAMD, PDB_NAMD, GRO, NUCL
+from MDAnalysis.tests.datafiles import PSF, DCD, PRMpbc, TRJpbc_bz2, PSF_NAMD, PDB_NAMD, GRO, NUCL, TPR, XTC
 
 from numpy.testing import *
 from numpy import array, float32
@@ -148,6 +148,16 @@ def test_prop(self):
     # add more test cases for byres, bynum, point
     # and also for selection keywords such as 'nucleic'
 
+    def test_same_resname(self):
+        """Test the 'same ... as' construct (Issue 217)"""
+        sel = self.universe.selectAtoms("same resname as resid 10 or resid 11")
+        assert_equal(len(sel), 331, "Found a wrong number of atoms with same resname as resids 10 or 11")
+        target_resids = array([ 7, 8, 10, 11, 12, 14, 17, 25, 32, 37, 38, 42, 46,
+                               49, 55, 56, 66, 73, 80, 85, 93, 95, 99, 100, 122, 127,
+                              130, 144, 150, 176, 180, 186, 188, 189, 194, 198, 203, 207, 214])
+        assert_array_equal(sel.resids(), target_resids, "Found wrong residues with same resname as resids 10 or 11")
+
+
     def test_empty_selection(self):
         """Test that empty selection can be processed (see Issue 12)"""
         assert_equal(len(self.universe.selectAtoms('resname TRP')), 0)  # no Trp in AdK
@@ -260,6 +270,29 @@ def test_atom(self):
         assert_equal(len(sel), 1)
         assert_equal(sel.resnames(), ['GLY'])
 
+    @dec.slow
+    def test_same_coordinate(self):
+        """Test the 'same ... as' construct (Issue 217)"""
+        # This test comes here because it's hard to get same _x with full precision formats.
+        #  The 'same' construct uses numpy.in1d to compare floats. It might be sensitive to
+        #  precision issues, but I am expecting .gro coordinates with the same values to
+        #  be converted to the exact same floats, at least in the same machine.
+        sel = self.universe.selectAtoms("same x as bynum 1 or bynum 10")
+        assert_equal(len(sel), 12, "Found a wrong number of atoms with same x as ids 1 or 10")
+        target_ids = array([ 0, 8, 9, 224, 643, 3515, 11210, 14121, 18430, 25418, 35811, 43618])
+        assert_array_equal(sel.indices(), target_ids, "Found wrong atoms with same x as ids 1 or 10")
+
+
+class TestSelectionsXTC(TestCase):
+    def setUp(self):
+        self.universe = MDAnalysis.Universe(TPR,XTC)
+
+    def test_same_fragment(self):
+        """Test the 'same ... as' construct (Issue 217)"""
+        # This test comes here because it's a system with solvent, and thus multiple fragments.
+        sel = self.universe.selectAtoms("same fragment as bynum 1")
+        assert_equal(len(sel), 3341, "Found a wrong number of atoms on the same fragment as id 1")
+        assert_equal(sel._atoms, self.universe.atoms[0].fragment._atoms, "Found a different set of atoms when using the 'same fragment as' construct vs. the .fragment prperty")
 
 class TestSelectionsNucleicAcids(TestCase):
     def setUp(self):