Add sort method to AtomGroup

package/AUTHORS

@@ -163,7 +163,8 @@ Chronological list of authors
   - Estefania Barreto-Ojeda
   - Paarth Thadani
   - Henry Kobin
-
+  - Kosuke Kudo
+
 External code
 -------------
 

package/CHANGELOG

@@ -17,7 +17,8 @@ The rules for this file:
          lilyminium, daveminh, jbarnoud, yuxuanzhuang, VOD555, ianmkenney,
          calcraven，xiki-tempula, mieczyslaw, manuel.nuno.melo, PicoCentauri,
          hanatok, rmeli, aditya-kamath, tirkarthi, LeonardoBarneschi, hejamu,
-         biogen98, orioncohen, z3y50n, hp115, ojeda-e, thadanipaarth, HenryKobin
+         biogen98, orioncohen, z3y50n, hp115, ojeda-e, thadanipaarth, HenryKobin,
+         1ut
 
   * 2.0.0
 
@@ -98,6 +99,7 @@ Fixes
   * Fix syntax warning over comparison of literals using is (Issue #3066)
 
 Enhancements
+  * Added sort method to the atomgroup (Issue #2976, PR #3188) 
   * ITPParser now reads [ atomtypes ] sections in ITP files, used for charges
     and masses not defined in the [ atoms ] sections
   * Add `set_dimensions` transformation class for setting constant 

package/MDAnalysis/core/groups.py

@@ -3324,6 +3324,83 @@ def write(self, filename=None, file_format=None,
 
         raise ValueError("No writer found for format: {}".format(filename))
 
+    def sort(self, key='ix', keyfunc=None):
+        """
+        Returns a sorted ``AtomGroup`` using a specified attribute as the key.
+
+        Parameters
+        ----------
+        key: str, optional
+            The name of the ``AtomGroup`` attribute to sort by (e.g. ``ids``,
+            ``ix``. default= ``ix`` ).
+        keyfunc: callable, optional
+            A function to convert multidimensional arrays to a single
+            dimension. This 1D array will be used as the sort key and
+            is required when sorting with an ``AtomGroup`` attribute
+            key which has multiple dimensions. Note: this argument
+            is ignored when the attribute is one dimensional.
+
+        Returns
+        ----------
+        :class:`AtomGroup`
+            Sorted ``AtomGroup``.
+
+        Example
+        ----------
+
+        .. code-block:: python
+
+            >>> import MDAnalysis as mda
+            >>> from MDAnalysisTests.datafiles import PDB_small
+            >>> u = mda.Universe(PDB_small)
+            >>> ag = u.atoms[[3, 2, 1, 0]]
+            >>> ag.ix
+            array([3 2 1 0])
+            >>> ag = ag.sort()
+            >>> ag.ix
+            array([0 1 2 3])
+            >>> ag.positions
+            array([[-11.921,  26.307,  10.41 ],
+                   [-11.447,  26.741,   9.595],
+                   [-12.44 ,  27.042,  10.926],
+                   [-12.632,  25.619,  10.046]], dtype=float32)
+            >>> ag = ag.sort("positions", lambda x: x[:, 1])
+            >>> ag.positions
+            array([[-12.632,  25.619,  10.046],
+                   [-11.921,  26.307,  10.41 ],
+                   [-11.447,  26.741,   9.595],
+                   [-12.44 ,  27.042,  10.926]], dtype=float32)
+
+        Note
+        ----------
+        This uses a stable sort as implemented by
+        `numpy.argsort(kind='stable')`.
+
+
+        .. versionadded:: 2.0.0
+        """
+        idx = getattr(self.atoms, key)
+        if len(idx) != len(self.atoms):
+            raise ValueError("The array returned by the attribute '{}' "
+                             "must have the same length as the number of "
+                             "atoms in the input AtomGroup".format(key))
+        if idx.ndim == 1:
+            order = np.argsort(idx, kind='stable')
+        elif idx.ndim > 1:
+            if keyfunc is None:
+                raise NameError("The {} attribute returns a multidimensional "
+                                "array. In order to sort it, a function "
+                                "returning a 1D array (to be used as the sort "
+                                "key) must be passed to the keyfunc argument"
+                                .format(key))
+            sortkeys = keyfunc(idx)
+            if sortkeys.ndim != 1:
+                raise ValueError("The function assigned to the argument "
+                                 "'keyfunc':{} doesn't return a 1D array."
+                                 .format(keyfunc))
+            order = np.argsort(sortkeys, kind='stable')
+        return self.atoms[order]
+
 
 class ResidueGroup(GroupBase):
     """ResidueGroup base class.

testsuite/MDAnalysisTests/core/test_atomgroup.py

@@ -1674,3 +1674,80 @@ def test_partial_timestep(self, universe):
                                 ag.ts.velocities,
                                 self.prec,
                                 err_msg="Partial timestep coordinates wrong")
+
+
+class TestAtomGroupSort(object):
+    """Tests the AtomGroup.sort attribute"""
+
+    @pytest.fixture()
+    def universe(self):
+        u = mda.Universe.empty(
+            n_atoms=7,
+            n_residues=3,
+            n_segments=2,
+            atom_resindex=np.array([0, 0, 0, 1, 1, 1, 2]),
+            residue_segindex=np.array([0, 0, 1]),
+            trajectory=True,
+            velocities=True,
+            forces=True
+        )
+        attributes = ["id", "charge", "mass", "tempfactor"]
+
+        for i in (attributes):
+            u.add_TopologyAttr(i, [6, 5, 4, 3, 2, 1, 0])
+
+        u.add_TopologyAttr('resid', [2, 1, 0])
+        u.add_TopologyAttr('segid', [1, 0])
+        u.add_TopologyAttr('bonds', [(0, 1)])
+
+        return u
+
+    @pytest.fixture()
+    def ag(self, universe):
+        ag = universe.atoms
+        ag.positions = (-np.arange(21)).reshape(7, 3)
+        return ag
+
+    test_ids = [
+       "ix",
+       "ids",
+       "resids",
+       "segids",
+       "charges",
+       "masses",
+       "tempfactors"
+    ]
+
+    test_data = [
+        ("ix", np.array([0, 1, 2, 3, 4, 5, 6])),
+        ("ids", np.array([6, 5, 4, 3, 2, 1, 0])),
+        ("resids", np.array([6, 3, 4, 5, 0, 1, 2])),
+        ("segids", np.array([6, 0, 1, 2, 3, 4, 5])),
+        ("charges", np.array([6, 5, 4, 3, 2, 1, 0])),
+        ("masses", np.array([6, 5, 4, 3, 2, 1, 0])),
+        ("tempfactors", np.array([6, 5, 4, 3, 2, 1, 0])),
+    ]
+
+    @pytest.mark.parametrize("inputs, expected", test_data, ids=test_ids)
+    def test_sort(self, ag, inputs, expected):
+        agsort = ag.sort(inputs)
+        assert np.array_equal(expected, agsort.ix)
+
+    def test_sort_bonds(self, ag):
+        with pytest.raises(ValueError, match=r"The array returned by the "
+                           "attribute"):
+            ag.sort("bonds")
+
+    def test_sort_positions_2D(self, ag):
+        with pytest.raises(ValueError, match=r"The function assigned to"):
+            ag.sort("positions", keyfunc=lambda x: x)
+
+    def test_sort_position_no_keyfunc(self, ag):
+        with pytest.raises(NameError, match=r"The .* attribute returns a "
+                           "multidimensional array. In order to sort it, "):
+            ag.sort("positions")
+
+    def test_sort_position(self, ag):
+        ref = [6, 5, 4, 3, 2, 1, 0]
+        agsort = ag.sort("positions", keyfunc=lambda x: x[:, 1])
+        assert np.array_equal(ref, agsort.ix)