Adds Unwrap keyword to center

package/CHANGELOG

@@ -56,6 +56,7 @@ Enhancements
     PR #2221)
   * survival probability additions: residues, intermittency, step with performance,
     (PR #2226)   
+  * added unwrap keyword to center (PR #2275)
 
 Changes
   * added official support for Python 3.7 (PR #1963)

package/MDAnalysis/core/groups.py

@@ -349,7 +349,16 @@ class _ImmutableBase(object):
     #  cache lookup if the class is reused (as in ag._derived_class(...)).
     __new__ = object.__new__
 
-
+def check_pbc_and_unwrap(function):
+    """Decorator to raise ValueError when both 'pbc' and 'unwrap' are set to True.
+    """
+    @functools.wraps(function)
+    def wrapped(group, *args, **kwargs):
+        if kwargs.get('compound') == 'group':
+            if kwargs.get('pbc') and kwargs.get('unwrap'):
+                raise ValueError("both 'pbc' and 'unwrap' can not be set to true")
+        return function(group, *args, **kwargs)
+    return wrapped
 
 def _only_same_level(function):
     @functools.wraps(function)
@@ -648,8 +657,10 @@ def isunique(self):
         #    return ``not np.any(mask)`` here but using the following is faster:
         return not np.count_nonzero(mask)
 
+
     @warn_if_not_unique
-    def center(self, weights, pbc=None, compound='group'):
+    @check_pbc_and_unwrap
+    def center(self, weights, pbc=None, compound='group', unwrap=False):
         """Weighted center of (compounds of) the group
 
         Computes the weighted center of :class:`Atoms<Atom>` in the group.
@@ -679,6 +690,8 @@ def center(self, weights, pbc=None, compound='group'):
             will be returned as an array of position vectors, i.e. a 2d array.
             Note that, in any case, *only* the positions of :class:`Atoms<Atom>`
             *belonging to the group* will be taken into account.
+        unwrap : bool, optional
+            If ``True``, compounds will be unwrapped before computing their centers.
 
         Returns
         -------
@@ -695,6 +708,8 @@ def center(self, weights, pbc=None, compound='group'):
         ValueError
             If `compound` is not one of ``'group'``, ``'segments'``,
             ``'residues'``, ``'molecules'``, or ``'fragments'``.
+        ValueError
+            If both 'pbc' and 'unwrap' set to true.
         ~MDAnalysis.exceptions.NoDataError
             If `compound` is ``'molecule'`` but the topology doesn't
             contain molecule information (molnums) or if `compound` is
@@ -722,6 +737,7 @@ def center(self, weights, pbc=None, compound='group'):
         .. versionchanged:: 0.19.0 Added `compound` parameter
         .. versionchanged:: 0.20.0 Added ``'molecules'`` and ``'fragments'``
             compounds
+        .. versionchanged:: 0.20.0 Added `unwrap` parameter
         """
 
         if pbc is None:
@@ -736,6 +752,8 @@ def center(self, weights, pbc=None, compound='group'):
         if comp == 'group':
             if pbc:
                 coords = atoms.pack_into_box(inplace=False)
+            elif unwrap:
+                coords = atoms.unwrap(compound=comp, reference=None, inplace=False)
             else:
                 coords = atoms.positions
             # If there's no atom, return its (empty) coordinates unchanged.
@@ -773,7 +791,12 @@ def center(self, weights, pbc=None, compound='group'):
         # required:
         sort_indices = np.argsort(compound_indices)
         compound_indices = compound_indices[sort_indices]
-        coords = atoms.positions[sort_indices]
+
+        # Unwrap Atoms
+        if unwrap:
+            coords = atoms.unwrap(compound=comp, reference=None, inplace=False)
+        else:
+            coords = atoms.positions[sort_indices]
         if weights is None:
             coords = coords.astype(dtype, copy=False)
         else:

testsuite/MDAnalysisTests/core/test_atomgroup.py

@@ -52,7 +52,7 @@
     GRO
 )
 from MDAnalysisTests import make_Universe, no_deprecated_call
-
+from MDAnalysisTests.core.util import UnWrapUniverse
 import pytest
 
 
@@ -518,6 +518,37 @@ def test_center_wrong_shape(self, ag):
         with pytest.raises(ValueError):
             ag.center(weights)
 
+    @pytest.mark.parametrize('level', ('atoms', 'residues', 'segments'))
+    @pytest.mark.parametrize('compound', ('fragments', 'molecules', 'residues',
+                                          'group', 'segments'))
+    @pytest.mark.parametrize('is_triclinic', (False, True))
+    def test_center_unwrap(self, level, compound, is_triclinic):
+        u = UnWrapUniverse(is_triclinic=is_triclinic)
+        # select group appropriate for compound:
+        if compound == 'group':
+            group = u.atoms[39:47] # molecule 12
+        elif compound == 'segments':
+            group = u.atoms[23:47] # molecules 10, 11, 12
+        else:
+            group = u.atoms
+        # select topology level:
+        if level == 'residues':
+            group = group.residues
+        elif level == 'segments':
+            group = group.segments
+
+        # get the expected results
+        center = group.center(weights=None, pbc=False, compound=compound, unwrap=True)
+
+        ref_center = u.center(compound=compound)
+        assert_almost_equal(ref_center, center, decimal=4)
+
+    def test_center_unwrap_pbc_true_group(self):
+        u = UnWrapUniverse(is_triclinic=False)
+        # select group appropriate for compound:
+        group = u.atoms[39:47]  # molecule 12
+        with pytest.raises(ValueError):
+            group.center(weights=None, compound="group", unwrap=True, pbc=True)
 
 class TestSplit(object):
 

testsuite/MDAnalysisTests/core/test_groups.py

@@ -1407,3 +1407,21 @@ def test_VE_no_uni_2(self, components):
 
         with pytest.raises(TypeError):
             cls(0, 2, 4)  # missing Universe
+
+
+class TestDecorator(object):
+    @groups.check_pbc_and_unwrap
+    def dummy_funtion(cls, compound="group", pbc=True, unwrap=True):
+        return 0
+
+    @pytest.mark.parametrize('compound', ('fragments', 'molecules', 'residues',
+                                          'group', 'segments'))
+    @pytest.mark.parametrize('pbc', (True, False))
+    @pytest.mark.parametrize('unwrap', (True, False))
+    def test_decorator(self, compound, pbc, unwrap):
+
+        if compound == 'group' and pbc and unwrap:
+            with pytest.raises(ValueError):
+                self.dummy_funtion(compound=compound, pbc=pbc, unwrap=unwrap)
+        else:
+            assert_equal(self.dummy_funtion(compound=compound, pbc=pbc, unwrap=unwrap), 0)

testsuite/MDAnalysisTests/core/util.py

@@ -333,6 +333,7 @@ def __new__(cls, have_bonds=True, have_masses=True, have_molnums=True,
         # bind custom methods to universe:
         u.unwrapped_coords = cls.unwrapped_coords.__get__(u)
         u.wrapped_coords = cls.wrapped_coords.__get__(u)
+        u.center = cls.center.__get__(u)
 
         return u
 
@@ -545,3 +546,67 @@ def wrapped_coords(self, compound, center):
         positions = relpos * np.array([a, a, a])
 
         return positions.astype(np.float32)
+
+    def center(self, compound):
+        """Returns centers which correspond to the unwrapped system.
+
+        Parameters
+        ----------
+        compound : {'atoms', 'group', 'segments', 'residues', 'molecules', \
+                    'fragments'}
+            Which type of component is unwrapped. Note that for ``'group'``,
+            the result will only be correct *if the group is the entire system*.
+
+        Note
+        ----
+        This function assumes that all atom masses are equal. Therefore, the
+        returned coordinates for ``center='com'`` and ``center='cog'`` are
+        identical.
+        """
+
+        relpos = self.unwrapped_coords(compound, reference=None)
+
+        comp = compound.lower()
+        if comp not in ['group', 'segments', 'residues', 'molecules',
+                        'fragments']:
+            raise ValueError("Unknown unwrap compound: {}".format(compound))
+
+        pos = 0
+
+        if compound=="residues":
+            center_pos = np.zeros((15, 3), dtype=np.float32)
+        else:
+            center_pos = np.zeros((12, 3), dtype=np.float32)
+
+        for base in range(3):
+            loc_center = relpos[base, :]
+            center_pos[pos,:] = loc_center
+            pos+=1
+
+        for base in range(3, 15, 3):
+            loc_center = np.mean(relpos[base:base + 3, :], axis=0)
+            center_pos[pos,:] = loc_center
+            pos+=1
+
+        if compound=="residues":
+            for base in range(15, 47, 4):
+                loc_center = np.mean(relpos[base:base + 4, :], axis=0)
+                center_pos[pos,:] = loc_center
+                pos+=1
+        else:
+            for base in range(15, 23, 4):
+                loc_center = np.mean(relpos[base:base + 4, :], axis=0)
+                center_pos[pos,:] = loc_center
+                pos+=1
+            for base in range(23, 47, 8):
+                loc_center = np.mean(relpos[base:base + 8, :], axis=0)
+                center_pos[pos,:] = loc_center
+                pos+=1
+
+        if compound == "group":
+            center_pos = center_pos[11]
+        elif compound == "segments":
+            center_pos = center_pos[9:]
+
+        return center_pos
+